def create_and_train(hidden_layer_sizes: List[int], batch_normalized: bool,
regularization: float, initial_learning_rate: float,
decay_factor: float, momentum: float, train_id: str,
target: str):
# Build network
net = build_network(hidden_layer_sizes, batch_normalized, regularization)
network_desc = {
'hidden_layers': hidden_layer_sizes,
'batch_norm': batch_normalized,
'regularization': regularization
}
# Train
opt = MomentumSGD(net, initial_learning_rate, decay_factor, True, momentum)
if target == 'coarse':
tr = training.subset(5000)
val = validation.subset(1000)
epochs = 10
elif target == 'fine':
tr = training.subset(10000)
val = validation.subset(1000)
epochs = 15
elif target == 'overfit':
tr = val = training.subset(100)
epochs = 400
else:
tr = training
val = validation
epochs = 20
opt.train(tr, val, epochs, 500)
training_desc = {
'epochs': epochs,
'initial_learning_rate': initial_learning_rate,
'decay_factor': decay_factor,
'momentum': momentum,
'final_cost_train': opt.cost_train[-1],
'final_acc_train': opt.acc_train[-1],
'final_cost_val': opt.cost_val[-1],
'final_acc_val': opt.acc_val[-1],
'plot': costs_accuracies_plot(opt, '{}.png'.format(train_id))
}
# Use the test dataset
test_res = {
'final_cost_test': 0,
'final_acc_test': 0,
'confusion_matrix': ''
}
if target == 'final':
test_res['final_cost_test'], test_res[
'final_acc_test'] = net.cost_accuracy(test)
test_res['confusion_matrix'] = confusion_matrix_plot(
net, test,
CIFAR10().labels, '{}_conf.png'.format(train_id))
return {**network_desc, **training_desc, **test_res}
def test_cifar():
# type: () -> None
"""
Performs One-class classification tests on CIFAR
"""
# Build dataset and model
dataset = CIFAR10(path='data/CIFAR10')
model = LSACIFAR10(input_shape=dataset.shape,
code_length=64,
cpd_channels=100).cuda().eval()
# Set up result helper and perform test
helper = OneClassResultHelper(dataset,
model,
checkpoints_dir='checkpoints/cifar10/',
output_file='cifar10.txt')
helper.test_one_class_classification()
def main():
print(f"\nStart training ...\n")
start_epoch = 0 # start from epoch 0 or last checkpoint epoch
print('==> Building model..')
net = BuildNet(backbone=args.arch,
num_classes=args.train_class_num,
embed_dim=args.embed_dim)
net = net.to(device)
if device == 'cuda':
net = torch.nn.DataParallel(net)
cudnn.benchmark = True
optimizer = torch.optim.SGD(net.parameters(),
lr=args.lr,
momentum=0.9,
weight_decay=5e-4)
if args.resume:
# Load checkpoint.
if os.path.isfile(args.resume):
print('==> Resuming from checkpoint..')
checkpoint = torch.load(args.resume)
net.load_state_dict(checkpoint['net'])
optimizer.load_state_dict(checkpoint['optimizer'])
start_epoch = checkpoint['epoch']
loggerList = []
for i in range(args.train_class_num, args.test_class_num + 1):
loggerList.append(
Logger(os.path.join(args.checkpoint, f'log{i}.txt'),
resume=True))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
else:
loggerList = []
for i in range(args.train_class_num, args.test_class_num + 1):
logger = Logger(os.path.join(args.checkpoint, f'log{i}.txt'))
logger.set_names([
'Epoch', 'Train Loss', 'Train Acc.', "Pos-F1", 'Norm-F1',
'Energy-F1'
])
loggerList.append(logger)
if not args.evaluate:
for epoch in range(start_epoch, args.es):
adjust_learning_rate(optimizer,
epoch,
args.lr,
factor=args.lr_factor,
step=args.lr_step)
print('\nEpoch: %d | Learning rate: %f ' %
(epoch + 1, optimizer.param_groups[0]['lr']))
train_out = train(net, trainloader, optimizer, criterion, device)
save_model(net, optimizer, epoch,
os.path.join(args.checkpoint, 'last_model.pth'))
for test_class_num in range(args.train_class_num,
args.test_class_num + 1):
testset = CIFAR10(
root='../../data',
train=False,
download=True,
transform=transform_test,
train_class_num=args.train_class_num,
test_class_num=test_class_num,
includes_all_train_class=args.includes_all_train_class)
testloader = torch.utils.data.DataLoader(testset,
batch_size=args.bs,
shuffle=False,
num_workers=4)
test_out = test(net, testloader, criterion, device)
logger = loggerList[test_class_num - args.train_class_num]
logger.append([
epoch + 1, train_out["train_loss"], train_out["accuracy"],
test_out["best_F1_possibility"], test_out["best_F1_norm"],
test_out["best_F1_energy"]
])
logger.close()
print(f"\nFinish training...\n")
def main():
global args, best_prec1, train_rec, test_rec
args = parser.parse_args()
args.root = "work"
args.folder = osp.join(args.root, args.arch)
setproctitle.setproctitle(args.arch)
# if osp.exists(args.folder):
# shutil.rmtree(args.folder)
os.makedirs(args.folder, exist_ok=True)
if args.dataset == "cifar10":
CIFAR = CIFAR10(args.data)
else:
CIFAR = CIFAR100(args.data)
# create model
if args.pretrained:
print("=> using pre-trained model '{}'".format(args.arch))
raise NotImplementedError("pre-trained is not supported on CIFAR")
# model = models.__dict__[args.arch](pretrained=True)
else:
print("=> creating model '{}'".format(args.arch))
model = models.__dict__[args.arch](CIFAR.num_classes)
args.distributed = args.world_size > 1
if not args.distributed:
if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
model.features = torch.nn.DataParallel(model.features)
model.cuda()
else:
model = torch.nn.DataParallel(model).cuda()
else:
model.cuda()
model = torch.nn.parallel.DistributedDataParallel(model)
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# optionally resume from a checkpoint
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_prec1 = checkpoint['best_prec1']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
cudnn.benchmark = True
from trainers.classification import ClassificationTrainer
train_loader, valid_loader = CIFAR.get_loader(args)
trainer = ClassificationTrainer(model, criterion, args, optimizer)
if args.evaluate:
trainer.evaluate(valid_loader, model, criterion)
return
from torch.optim.lr_scheduler import MultiStepLR, StepLR
step1 = int(args.epochs * 0.5)
step2 = int(args.epochs * 0.75)
lr_scheduler = MultiStepLR(optimizer, milestones=[step1, step2], gamma=0.1)
trainer.fit(train_loader,
valid_loader,
start_epoch=0,
max_epochs=200,
lr_scheduler=lr_scheduler)
gradients = np.array([
[0, 1],
[3, 4],
[6, 7],
[9, 2],
])
grr = batch_norm_grad_slow(bn, gradients)
backprop = bn.backward(gradients)
assert np.allclose(backprop, grr), \
'Different\n' + str(grr) + '\n' + str(backprop)
print('Passed')
if __name__ == '__main__':
cifar = CIFAR10()
training = cifar.get_named_batches('data_batch_1').subset(20)
np.random.seed(123)
one_layer_no_reg()
one_layer_with_bn()
one_layer_with_reg()
two_layer_with_reg()
two_layer_with_bn()
three_layer_with_bn()
four_layer_with_bn()
# -*- coding: utf-8 -*-
# @Author: Andre Goncalves
# @Date: 2019-10-30 10:09:00
# @Last Modified by: Andre Goncalves
# @Last Modified time: 2019-10-30 10:11:34
import sys
sys.path.append('../')
from datasets import CIFAR10
from classifiers.deterministic.LogisticClassifier import LogisticClassifier
from classifiers.deterministic.DenseNet import DenseNet
from design import ActiveLearningLoop
if __name__ == '__main__':
dataset = CIFAR10()
dataset.prepare()
n_epochs = 2
methods = [{'model': DenseNet(name='DN-Random', epochs=n_epochs), 'acq_func': 'random'},
{'model': DenseNet(name='DN-Entropy', epochs=n_epochs), 'acq_func': 'entropy'},
{'model': DenseNet(name='DN-Abstention', epochs=n_epochs), 'acq_func': 'abstention'},
{'model': DenseNet(name='DN-Abstention-Entropy-A', epochs=n_epochs), 'acq_func': 'abstention_entropy_amean'},
{'model': DenseNet(name='DN-Abstention-Entropy-H', epochs=n_epochs), 'acq_func': 'abstention_entropy_hmean'}, ]
exp_folder = __file__.strip('.py')
exp_loop = ActiveLearningLoop(exp_folder)
exp_loop.execute(dataset, methods,
train_perc=0.2,
test_perc=0.3,
def main(results_path, network_config: dict, learningrate: int = 1e-3, weight_decay: float = 1e-5,
n_updates: int = int(1e5), device: torch.device = torch.device("cuda:0")):
"""Main function that takes hyperparameters and performs training and evaluation of model"""
# Prepare a path to plot to
plotpath = os.path.join(results_path, 'plots')
os.makedirs(plotpath, exist_ok=True)
# Load or download CIFAR10 dataset
cifar10_dataset = CIFAR10(data_folder='cifar10')
# Split dataset into training, validation, and test set randomly
trainingset = torch.utils.data.Subset(cifar10_dataset,
indices=np.arange(int(len(cifar10_dataset) * (3 / 5))))
validationset = torch.utils.data.Subset(cifar10_dataset,
indices=np.arange(int(len(cifar10_dataset) * (3 / 5)),
int(len(cifar10_dataset) * (4 / 5))))
testset = torch.utils.data.Subset(cifar10_dataset,
indices=np.arange(int(len(cifar10_dataset) * (4 / 5)),
len(cifar10_dataset)))
# Create datasets and dataloaders with rotated targets without augmentation (for evaluation)
trainingset_eval = RotatedImages(dataset=trainingset, rotation_angle=45.)
validationset = RotatedImages(dataset=validationset, rotation_angle=45.)
testset = RotatedImages(dataset=testset, rotation_angle=45.)
trainloader = torch.utils.data.DataLoader(trainingset_eval, batch_size=1, shuffle=False,
num_workers=0)
valloader = torch.utils.data.DataLoader(validationset, batch_size=1, shuffle=False,
num_workers=0)
testloader = torch.utils.data.DataLoader(testset, batch_size=1, shuffle=False, num_workers=0)
# Create datasets and dataloaders with rotated targets with augmentation (for training)
trainingset_augmented = RotatedImages(dataset=trainingset, rotation_angle=45.,
transform_chain=transforms.Compose(
[transforms.RandomHorizontalFlip(),
transforms.RandomVerticalFlip()]))
trainloader_augmented = torch.utils.data.DataLoader(trainingset_augmented, batch_size=16,
shuffle=True,
num_workers=0)
# Define a tensorboard summary writer that writes to directory "results_path/tensorboard"
writer = SummaryWriter(log_dir=os.path.join(results_path, 'tensorboard'))
# Create Network
net = SimpleCNN(**network_config)
net.to(device)
# Get mse loss function
mse = torch.nn.MSELoss()
# Get adam optimizer
optimizer = torch.optim.Adam(net.parameters(), lr=learningrate, weight_decay=weight_decay)
print_stats_at = 1e2 # print status to tensorboard every x updates
plot_at = 1e4 # plot every x updates
validate_at = 5e3 # evaluate model on validation set and check for new best model every x updates
update = 0 # current update counter
best_validation_loss = np.inf # best validation loss so far
update_progess_bar = tqdm.tqdm(total=n_updates, desc=f"loss: {np.nan:7.5f}",
position=0) # progressbar
# Save initial model as "best" model (will be overwritten later)
torch.save(net, os.path.join(results_path, 'best_model.pt'))
# Train until n_updates update have been reached
while update < n_updates:
for data in trainloader_augmented:
# Get next samples in `trainloader_augmented`
inputs, targets, ids = data
inputs = inputs.to(device)
targets = targets.to(device)
# Reset gradients
optimizer.zero_grad()
# Get outputs for network
outputs = net(inputs)
# Calculate loss, do backward pass, and update weights
loss = mse(outputs, targets)
loss.backward()
optimizer.step()
# Print current status and score
if update % print_stats_at == 0 and update > 0:
writer.add_scalar(tag="training/loss",
scalar_value=loss.cpu(),
global_step=update)
# Plot output
if update % plot_at == 0:
plot(inputs.detach().cpu().numpy(), targets.detach().cpu().numpy(),
outputs.detach().cpu().numpy(),
plotpath, update)
# Evaluate model on validation set
if update % validate_at == 0 and update > 0:
val_loss = evaluate_model(net, dataloader=valloader, device=device)
writer.add_scalar(tag="validation/loss", scalar_value=val_loss.cpu(),
global_step=update)
# Add weights as arrays to tensorboard
for i, param in enumerate(net.parameters()):
writer.add_histogram(tag=f'validation/param_{i}', values=param.cpu(),
global_step=update)
# Add gradients as arrays to tensorboard
for i, param in enumerate(net.parameters()):
writer.add_histogram(tag=f'validation/gradients_{i}',
values=param.grad.cpu(),
global_step=update)
# Save best model for early stopping
if best_validation_loss > val_loss:
best_validation_loss = val_loss
torch.save(net, os.path.join(results_path, 'best_model.pt'))
update_progess_bar.set_description(f"loss: {loss:7.5f}", refresh=True)
update_progess_bar.update()
# Increment update counter, exit if maximum number of updates is reached
update += 1
if update >= n_updates:
break
update_progess_bar.close()
print('Finished Training!')
# Load best model and compute score on test set
print(f"Computing scores for best model")
net = torch.load(os.path.join(results_path, 'best_model.pt'))
test_loss = evaluate_model(net, dataloader=testloader, device=device)
val_loss = evaluate_model(net, dataloader=valloader, device=device)
train_loss = evaluate_model(net, dataloader=trainloader, device=device)
print(f"Scores:")
print(f"test loss: {test_loss}")
print(f"validation loss: {val_loss}")
print(f"training loss: {train_loss}")
# Write result to file
with open(os.path.join(results_path, 'results.txt'), 'w') as fh:
print(f"Scores:", file=fh)
print(f"test loss: {test_loss}", file=fh)
print(f"validation loss: {val_loss}", file=fh)
print(f"training loss: {train_loss}", file=fh)
def main(args):
# Load CIFAR-10 dataset
data = CIFAR10()
# Instantiate networks
f_online = encoders[args.encoder]()
g_online = ProjectionHead()
q_online = ProjectionHead()
f_target = encoders[args.encoder]()
g_target = ProjectionHead()
# Initialize the weights of the networks
x = tf.random.normal((256, 32, 32, 3))
h = f_online(x, training=False)
print('Initializing online networks...')
print('Shape of h:', h.shape)
z = g_online(h, training=False)
print('Shape of z:', z.shape)
p = q_online(z, training=False)
print('Shape of p:', p.shape)
h = f_target(x, training=False)
print('Initializing target networks...')
print('Shape of h:', h.shape)
z = g_target(h, training=False)
print('Shape of z:', z.shape)
num_params_f = tf.reduce_sum(
[tf.reduce_prod(var.shape) for var in f_online.trainable_variables])
print(
'The encoders have {} trainable parameters each.'.format(num_params_f))
# Define optimizer
lr = 1e-3 * args.batch_size / 512
opt = tf.keras.optimizers.Adam(learning_rate=lr)
print('Using Adam optimizer with learning rate {}.'.format(lr))
@tf.function
def train_step_pretraining(x1, x2): # (bs, 32, 32, 3), (bs, 32, 32, 3)
# Forward pass
h_target_1 = f_target(x1, training=True)
z_target_1 = g_target(h_target_1, training=True)
h_target_2 = f_target(x2, training=True)
z_target_2 = g_target(h_target_2, training=True)
with tf.GradientTape(persistent=True) as tape:
h_online_1 = f_online(x1, training=True)
z_online_1 = g_online(h_online_1, training=True)
p_online_1 = q_online(z_online_1, training=True)
h_online_2 = f_online(x2, training=True)
z_online_2 = g_online(h_online_2, training=True)
p_online_2 = q_online(z_online_2, training=True)
p_online = tf.concat([p_online_1, p_online_2], axis=0)
z_target = tf.concat([z_target_2, z_target_1], axis=0)
loss = byol_loss(p_online, z_target)
# Backward pass (update online networks)
grads = tape.gradient(loss, f_online.trainable_variables)
opt.apply_gradients(zip(grads, f_online.trainable_variables))
grads = tape.gradient(loss, g_online.trainable_variables)
opt.apply_gradients(zip(grads, g_online.trainable_variables))
grads = tape.gradient(loss, q_online.trainable_variables)
opt.apply_gradients(zip(grads, q_online.trainable_variables))
del tape
return loss
batches_per_epoch = data.num_train_images // args.batch_size
log_every = 10 # batches
save_every = 100 # epochs
losses = []
for epoch_id in range(args.num_epochs):
data.shuffle_training_data()
for batch_id in range(batches_per_epoch):
x1, x2 = data.get_batch_pretraining(batch_id, args.batch_size)
loss = train_step_pretraining(x1, x2)
losses.append(float(loss))
# Update target networks (exponential moving average of online networks)
beta = 0.99
f_target_weights = f_target.get_weights()
f_online_weights = f_online.get_weights()
for i in range(len(f_online_weights)):
f_target_weights[i] = beta * f_target_weights[i] + (
1 - beta) * f_online_weights[i]
f_target.set_weights(f_target_weights)
g_target_weights = g_target.get_weights()
g_online_weights = g_online.get_weights()
for i in range(len(g_online_weights)):
g_target_weights[i] = beta * g_target_weights[i] + (
1 - beta) * g_online_weights[i]
g_target.set_weights(g_target_weights)
if (batch_id + 1) % log_every == 0:
print('[Epoch {}/{} Batch {}/{}] Loss={:.5f}.'.format(
epoch_id + 1, args.num_epochs, batch_id + 1,
batches_per_epoch, loss))
if (epoch_id + 1) % save_every == 0:
f_online.save_weights('f_online_{}.h5'.format(epoch_id + 1))
print('Weights of f saved.')
np.savetxt('losses.txt', tf.stack(losses).numpy())
def main(args):
# Load CIFAR-10 dataset
data = CIFAR10()
# Define hyperparameters
num_epochs = 50
batch_size = 512
# Instantiate networks f and c
f_net = encoders[args.encoder]()
c_net = ClassificationHead()
# Initialize the weights of f and c
x, y = data.get_batch_finetuning(batch_id=0, batch_size=batch_size)
h = f_net(x, training=False)
print('Shape of h:', h.shape)
s = c_net(h)
print('Shape of s:', s.shape)
# Load the weights of f from pretraining
f_net.load_weights(args.encoder_weights)
print('Weights of f loaded.')
# Define optimizer
batches_per_epoch = data.num_train_images // batch_size
total_update_steps = num_epochs * batches_per_epoch
lr_schedule = tf.keras.optimizers.schedules.PolynomialDecay(
5e-2, total_update_steps, 5e-4, power=2)
opt = tf.keras.optimizers.Adam(learning_rate=lr_schedule)
@tf.function
def train_step_evaluation(x, y): # (bs, 32, 32, 3), (bs)
# Forward pass
with tf.GradientTape() as tape:
h = f_net(x, training=False) # (bs, 512)
y_pred_logits = c_net(h) # (bs, 10)
loss = tf.reduce_mean(
tf.nn.sparse_softmax_cross_entropy_with_logits(
labels=y, logits=y_pred_logits))
# Backward pass
grads = tape.gradient(loss, c_net.trainable_variables)
opt.apply_gradients(zip(grads, c_net.trainable_variables))
return loss
log_every = 10 # batches
for epoch_id in range(num_epochs):
data.shuffle_training_data()
for batch_id in range(batches_per_epoch):
x, y = data.get_batch_finetuning(batch_id, batch_size)
loss = train_step_evaluation(x, y)
if (batch_id + 1) % log_every == 0:
print('[Epoch {}/{} Batch {}/{}] Loss: {:.4f}'.format(
epoch_id + 1, num_epochs, batch_id + 1, batches_per_epoch,
loss))
# Compute classification accuracy on test set
test_accuracy = compute_test_accuracy(data, f_net, c_net)
print('Test Accuracy: {:.4f}'.format(test_accuracy))
model_parser.add_argument('--net', action='store_true', default=False)
model_parser.add_argument('--grouprgb', action='store_true', default=False)
model_parser.add_argument('--otherrgb', action='store_true', default=False)
dataset_parser = parser.add_mutually_exclusive_group(required=True)
dataset_parser.add_argument('--mnist', action='store_true', default=False)
dataset_parser.add_argument('--cifar10', action='store_true', default=False)
args = parser.parse_args()
num_epochs = args.epochs
save_step = args.save_step
kernel_size = args.kernel_size
maxpool = 1 if args.disable_pool else 2
dropout = not args.disable_dropout
if args.cifar10:
dataset = CIFAR10(args.batch_size)
else:
dataset = MNIST(args.batch_size)
comment = args.comment
if args.other:
model_class = OtherNet
if args.group:
model_class = GroupNet
if args.net:
model_class = Net
if args.grouprgb:
model_class = GroupNetRGB
if args.otherrgb:
model_class = OtherNetRGB
print(args)
model = model_class(dataset.batch_size,
def create_and_train(training: Batch,
validation: Batch,
epochs: int,
hidden_size: int,
regularization: float,
initial_learning_rate: float,
decay_factor: float,
momentum: float,
train_id: str,
test: Batch = None):
"""
Create and train a 2 layer network:
- subtract mean of the training set
- linear layer
- relu
- linear layer
- softmax
The only parameters that are fixed are the layer initializers
and the batch size.
:param train_id:
:param training:
:param validation:
:param epochs:
:param hidden_size:
:param regularization:
:param initial_learning_rate:
:param decay_factor:
:param momentum:
:return:
"""
# Mean of the training set
mu = training.mean()
# Definition of the network
net = Network()
net.add_layer(layers.BatchNormalization(CIFAR10.input_size, mu))
net.add_layer(layers.Linear(CIFAR10.input_size, hidden_size, regularization, initializers.Xavier()))
net.add_layer(layers.ReLU(hidden_size))
net.add_layer(layers.Linear(hidden_size, CIFAR10.output_size, regularization, initializers.Xavier()))
net.add_layer(layers.Softmax(CIFAR10.output_size))
# Training
opt = optimizers.MomentumSGD(net, initial_learning_rate, decay_factor, True, momentum)
opt.train(training, validation, epochs, 10000)
# Plotting
plot = costs_accuracies_plot(opt.epoch_nums, opt.acc_train, opt.acc_val,
opt.cost_train, opt.cost_val,
'images/{}.png'.format(train_id))
result = {
'epochs': epochs,
'hidden_size': hidden_size,
'regularization': regularization,
'initial_learning_rate': initial_learning_rate,
'decay_factor': decay_factor,
'momentum': momentum,
# 'net': net,
# 'opt': opt,
'epoch_nums': opt.epoch_nums,
'cost_train': opt.cost_train,
'acc_train': opt.acc_train,
'cost_val': opt.cost_val,
'acc_val': opt.acc_val,
'final_cost_train': opt.cost_train[-1],
'final_acc_train': opt.acc_train[-1],
'final_cost_val': opt.cost_val[-1],
'final_acc_val': opt.acc_val[-1],
'plot': plot
}
# Test set
if test is not None:
result['final_cost_test'], result['final_acc_test'] = net.cost_accuracy(test)
result['confusion_matrix'] = confusion_matrix_plot(net, test,
CIFAR10().labels,
'images/{}_conf.png'.format(train_id))
return result
def main():
args.checkpoint = './checkpoints/cifar/' + args.arch
if not os.path.isdir(args.checkpoint):
mkdir_p(args.checkpoint)
# folder to save figures
args.plotfolder = './checkpoints/cifar/' + args.arch + '/plotter'
if not os.path.isdir(args.plotfolder):
mkdir_p(args.plotfolder)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(device)
start_epoch = 0 # start from epoch 0 or last checkpoint epoch
print('==> Preparing data..')
print('==> Preparing data..')
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
trainset = CIFAR10(root='../../data',
train=True,
download=True,
transform=transform_train,
train_class_num=args.train_class_num,
test_class_num=args.test_class_num,
includes_all_train_class=args.includes_all_train_class)
trainloader = torch.utils.data.DataLoader(trainset,
batch_size=args.bs,
shuffle=True,
num_workers=4)
testset = CIFAR10(root='../../data',
train=False,
download=True,
transform=transform_test,
train_class_num=args.train_class_num,
test_class_num=args.test_class_num,
includes_all_train_class=args.includes_all_train_class)
testloader = torch.utils.data.DataLoader(testset,
batch_size=args.bs,
shuffle=False,
num_workers=4)
print('==> Building model..')
net = Network(backbone=args.arch,
num_classes=args.train_class_num,
embed_dim=args.embed_dim)
fea_dim = net.classifier.in_features
net = net.to(device)
if device == 'cuda':
net = torch.nn.DataParallel(net)
cudnn.benchmark = True
criterion_softamx = nn.CrossEntropyLoss()
print(fea_dim)
criterion_rpl = RPLoss(number_class=args.train_class_num,
feat_dim=fea_dim).to(device)
optimizer_rpl = torch.optim.SGD(criterion_rpl.parameters(),
lr=args.center_lr,
momentum=0.9,
weight_decay=5e-4)
criterion_centerloss = CenterLoss(num_classes=args.train_class_num,
feat_dim=fea_dim).to(device)
optimizer_softmax = optim.SGD(net.parameters(),
lr=args.lr,
momentum=0.9,
weight_decay=5e-4)
optimizer_centerloss = torch.optim.SGD(criterion_centerloss.parameters(),
lr=args.center_lr,
momentum=0.9,
weight_decay=5e-4)
if args.resume:
# Load checkpoint.
if os.path.isfile(args.resume):
print('==> Resuming from checkpoint..')
checkpoint = torch.load(args.resume)
net.load_state_dict(checkpoint['net'])
criterion_rpl.load_state_dict(checkpoint['centerloss'])
# best_acc = checkpoint['acc']
# print("BEST_ACCURACY: "+str(best_acc))
start_epoch = checkpoint['epoch']
logger = Logger(os.path.join(args.checkpoint, 'log.txt'),
resume=True)
else:
print("=> no checkpoint found at '{}'".format(args.resume))
else:
logger = Logger(os.path.join(args.checkpoint, 'log.txt'))
logger.set_names(['Epoch', 'Total Loss', 'train Acc.'])
if not args.evaluate:
scheduler = lr_scheduler.StepLR(optimizer_softmax,
step_size=30,
gamma=0.1)
for epoch in range(start_epoch, start_epoch + args.es):
print('\nEpoch: %d Learning rate: %f' %
(epoch + 1, optimizer_softmax.param_groups[0]['lr']))
train_loss, train_acc = train(net, trainloader, optimizer_softmax,
optimizer_rpl, criterion_rpl, device)
save_model(net, criterion_centerloss, epoch,
os.path.join(args.checkpoint, 'last_model.pth'))
# plot the training data
# if args.plot:
# plot_feature(net,criterion_centerloss, trainloader, device, args.plotfolder, epoch=epoch,
# plot_class_num=args.train_class_num,maximum=args.plot_max, plot_quality=args.plot_quality)
logger.append([epoch + 1, train_loss, train_acc])
scheduler.step()
test(net, testloader, criterion_rpl, device)
# if args.plot:
# plot_feature(net, criterion_centerloss, testloader, device, args.plotfolder, epoch="test",
# plot_class_num=args.train_class_num+1, maximum=args.plot_max, plot_quality=args.plot_quality)
logger.close()
def main_stage2(net, mid_known, mid_unknown):
print("Starting stage-2 fine-tuning ...")
start_epoch = 0
criterion = FinetuneLoss(mid_known=mid_known,
mid_unknown=mid_unknown,
gamma=args.gamma,
temperature=args.temperature,
feature='energy')
criterion = criterion.to(device)
optimizer = torch.optim.SGD(net.parameters(),
lr=args.stage2_lr,
momentum=0.9,
weight_decay=5e-4)
if args.stage2_resume:
# Load checkpoint.
if os.path.isfile(args.stage2_resume):
print('==> Resuming from checkpoint..')
checkpoint = torch.load(args.stage2_resume)
net.load_state_dict(checkpoint['net'])
optimizer.load_state_dict(checkpoint['optimizer'])
start_epoch = checkpoint['epoch']
loggerList = []
for i in range(args.train_class_num, args.test_class_num + 1):
loggerList.append(
Logger(os.path.join(args.checkpoint, f'log{i}_stage2.txt'),
resume=True))
else:
print("=> no checkpoint found at '{}'".format(args.stage2_resume))
else:
loggerList = []
for i in range(args.train_class_num, args.test_class_num + 1):
logger = Logger(os.path.join(args.checkpoint,
f'log{i}_stage2.txt'))
logger.set_names(
['Epoch', 'Train Loss', 'Train Acc.', 'Energy-F1'])
loggerList.append(logger)
if not args.evaluate:
for epoch in range(start_epoch, args.stage2_es):
adjust_learning_rate(optimizer,
epoch,
args.stage2_lr,
factor=args.stage2_lr_factor,
step=args.stage2_lr_step)
print('\nStage_2 Epoch: %d | Learning rate: %f ' %
(epoch + 1, optimizer.param_groups[0]['lr']))
train_out = stage2_train(net, trainloader, optimizer, criterion,
device)
save_model(net, optimizer, epoch,
os.path.join(args.checkpoint, 'stage_2_last_model.pth'))
for test_class_num in range(args.train_class_num,
args.test_class_num + 1):
testset = CIFAR10(
root='../../data',
train=False,
download=True,
transform=transform_test,
train_class_num=args.train_class_num,
test_class_num=test_class_num,
includes_all_train_class=args.includes_all_train_class)
testloader = torch.utils.data.DataLoader(
testset,
batch_size=args.stage2_bs,
shuffle=False,
num_workers=4)
test_out = test(net, testloader, device)
logger = loggerList[test_class_num - args.train_class_num]
logger.append([
epoch + 1, train_out["train_loss"], train_out["accuracy"],
test_out["best_F1"]
])
logger.close()
print(f"\nFinish Stage-2 training...\n")
print('==> Preparing data..')
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
trainset = CIFAR10(root='../../data',
train=True,
download=True,
transform=transform_train,
train_class_num=args.train_class_num,
test_class_num=args.test_class_num,
includes_all_train_class=args.includes_all_train_class)
testset = CIFAR10(root='../../data',
train=False,
download=True,
transform=transform_test,
train_class_num=args.train_class_num,
test_class_num=args.test_class_num,
includes_all_train_class=args.includes_all_train_class)
# data loader
trainloader = torch.utils.data.DataLoader(trainset,
batch_size=args.stage1_bs,
shuffle=True,
num_workers=4)
def main():
print(args)
# Check the save_dir exists or not
if not os.path.exists(args.save_dir):
os.makedirs(args.save_dir)
torch.cuda.manual_seed(args.seed)
cudnn.benchmark = True
global best_prec1
# prepare dataset
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
])
transform_test = transforms.Compose([
transforms.ToTensor(),
])
trainset = CIFAR10(root='~/datasets/CIFAR10',
train=True,
download=True,
transform=transform_train)
train_loader = DataLoader(trainset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
num_classes = trainset.num_classes
targets = np.asarray(trainset.targets)
testset = CIFAR10(root='~/datasets/CIFAR10',
train=False,
download=True,
transform=transform_test)
test_loader = DataLoader(testset,
batch_size=args.batch_size * 4,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
model = get_model(args, num_classes)
if torch.cuda.device_count() > 1:
model = nn.DataParallel(model)
model.cuda()
criterion = get_loss(args, labels=targets, num_classes=num_classes)
optimizer = get_optimizer(model, args)
# optionally resume from a checkpoint
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch'] + 1
best_prec1 = checkpoint['best_prec1']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optim_dict'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
scheduler = get_scheduler(optimizer, args)
if args.evaluate:
validate(test_loader, model)
return
print("*" * 40)
for epoch in range(args.start_epoch, args.epochs + 1):
scheduler.step(epoch)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch)
print("*" * 40)
# evaluate on validation sets
print("train:", end="\t")
prec1 = validate(train_loader, model)
print("test:", end="\t")
prec1 = validate(test_loader, model)
print("*" * 40)
# remember best prec@1 and save checkpoint
is_best = prec1 > best_prec1
best_prec1 = max(prec1, best_prec1)
if (epoch < 70
and epoch % 10 == 0) or (epoch >= 70
and epoch % args.save_freq == 0):
filename = 'checkpoint_{}.tar'.format(epoch)
else:
filename = None
save_checkpoint(args.save_dir, {
'epoch': epoch,
'state_dict': model.state_dict(),
'optim_dict': optimizer.state_dict(),
'best_prec1': best_prec1,
},
is_best,
filename=filename)
def main():
logging.basicConfig(level=logging.INFO,
format="%(asctime)s | %(message)s",
handlers=[
logging.FileHandler(
os.path.join(args.save_dir, 'training.log')),
logging.StreamHandler()
])
logger = logging.getLogger()
logging.info('Robust self-training')
logging.info('Args: %s', args)
# Check the save_dir exists or not
if not os.path.exists(args.save_dir):
os.makedirs(args.save_dir)
cudnn.benchmark = True
global best_prec1
# prepare dataset
transform_test = transforms.Compose([
transforms.ToTensor(),
])
num_classes = 10
testset = CIFAR10(root='../data',
train=False,
download=True,
transform=transform_test)
test_loader = DataLoader(testset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
model = get_model(args, num_classes)
if torch.cuda.device_count() > 1:
model = nn.DataParallel(model)
model.cuda()
path = os.path.join(args.save_dir, 'cifar10.tar')
model.load_state_dict(torch.load(path)["state_dict"])
loss, accuracy, robust_accuracy, robust_clean_accuracy = eval(
args, model, device, 'test', test_loader)
print('clean accuracy : {}, robust accuracy : {}'.format(
robust_clean_accuracy, robust_accuracy))
log_path = os.path.join(args.save_dir, 'log_best.txt')
adversary = AutoAttack(args,
model,
norm='Linf',
eps=args.epsilon,
log_path=log_path,
version='standard',
device=device)
l = [data[0] for data in test_loader]
x_test = torch.cat(l, 0)
l = [data[1] for data in test_loader]
y_test = torch.cat(l, 0)
#print(x_test.shape)
with torch.no_grad():
if not args.individual:
adv_complete = adversary.run_standard_evaluation(
x_test, y_test, bs=int(args.test_batch_size))
torch.save({'adv_complete': adv_complete},
'{}/{}_{}_1_{}_eps_{:.5f}.pth'.format(
args.save_dir, 'aa', 'standard',
adv_complete.shape[0], args.epsilon))
else:
# individual version, each attack is run on all test points
adv_complete = adversary.run_standard_evaluation_individual(
x_test, y_test, bs=int(args.val_batch_size))
torch.save(
adv_complete,
'{}/{}_{}_individual_1_{}_eps_{:.5f}_plus_{}_cheap_{}_{}.pth'.
format(args.save_dir, 'aa', 'standard',
len(test_loader.dataset), args.epsilon))
def main():
global best_prec1, train_rec, test_rec
conf = get_configs()
conf.root = "work"
conf.folder = osp.join(conf.root, conf.arch)
conf.device = torch.device(
"cuda:0" if torch.cuda.is_available() else "cpu")
setproctitle.setproctitle(conf.arch)
os.makedirs(conf.folder, exist_ok=True)
if conf.dataset == "cifar10":
CIFAR = CIFAR10(conf.data) # Datasets object
else:
CIFAR = CIFAR100(conf.data)
# create model
if conf.pretrained:
print("=> using pre-trained model '{}'".format(conf.arch))
raise NotImplementedError("pre-trained is not supported on CIFAR")
# model = models.__dict__[args.arch](pretrained=True)
else:
print("=> creating model '{}'".format(conf.arch))
model = models.__dict__[conf.arch](CIFAR.num_classes)
# print(model.features)
conf.distributed = conf.distributed_processes > 1
if not conf.distributed:
if conf.gpus > 0:
model = nn.DataParallel(model)
model.to(conf.device)
else:
model.cuda()
model = torch.nn.parallel.DistributedDataParallel(model)
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(),
conf.lr,
momentum=conf.momentum,
weight_decay=conf.weight_decay)
# optionally resume from a checkpoint
if conf.resume:
if os.path.isfile(conf.resume):
print("=> loading checkpoint from '{}'".format(conf.resume))
checkpoint = torch.load(conf.resume)
conf.start_epoch = checkpoint['epoch']
best_prec1 = checkpoint['best_prec1']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint done. (epoch {})".format(
checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(conf.resume))
cudnn.benchmark = True # improve the efficiency of the program
train_loader, valid_loader = CIFAR.get_loader(conf)
trainer = ClassificationTrainer(model, criterion, conf, optimizer)
if conf.evaluate:
trainer.evaluate(valid_loader, model, criterion)
return
step1 = int(conf.epochs * 0.5)
step2 = int(conf.epochs * 0.75)
lr_scheduler = MultiStepLR(optimizer, milestones=[step1, step2], gamma=0.1)
trainer.fit(train_loader,
valid_loader,
start_epoch=0,
max_epochs=200,
lr_scheduler=lr_scheduler)