def checkpoint(self, name=None, timestamp=False, tags=None, verbose=False):
if name is None:
name = self.config["name"]
return save_checkpoint(self.get_state(),
name=name, tag=tags, timestamp=timestamp,
verbose=verbose)
def main():
args = parse_arguments()
random.seed(args.seed)
torch.manual_seed(args.seed)
if args.use_cuda:
torch.cuda.manual_seed_all(args.seed)
cudnn.benchmark = True
model_path = get_model_path(args.dataset, args.arch, args.seed)
# Init logger
log_file_name = os.path.join(model_path, 'log.txt')
print("Log file: {}".format(log_file_name))
log = open(log_file_name, 'w')
print_log('model path : {}'.format(model_path), log)
state = {k: v for k, v in args._get_kwargs()}
for key, value in state.items():
print_log("{} : {}".format(key, value), log)
print_log("Random Seed: {}".format(args.seed), log)
print_log("Python version : {}".format(sys.version.replace('\n', ' ')),
log)
print_log("Torch version : {}".format(torch.__version__), log)
print_log("Cudnn version : {}".format(torch.backends.cudnn.version()),
log)
# Data specifications for the webistes dataset
mean = [0., 0., 0.]
std = [1., 1., 1.]
input_size = 224
num_classes = 4
# Dataset
traindir = os.path.join(WEBSITES_DATASET_PATH, 'train')
valdir = os.path.join(WEBSITES_DATASET_PATH, 'val')
train_transform = transforms.Compose([
transforms.Resize(input_size),
transforms.ToTensor(),
transforms.Normalize(mean, std)
])
test_transform = transforms.Compose([
transforms.Resize(input_size),
transforms.ToTensor(),
transforms.Normalize(mean, std)
])
data_train = dset.ImageFolder(root=traindir, transform=train_transform)
data_test = dset.ImageFolder(root=valdir, transform=test_transform)
# Dataloader
data_train_loader = torch.utils.data.DataLoader(data_train,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
data_test_loader = torch.utils.data.DataLoader(data_test,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
# Network
if args.arch == "vgg16":
net = models.vgg16(pretrained=True)
elif args.arch == "vgg19":
net = models.vgg19(pretrained=True)
elif args.arch == "resnet18":
net = models.resnet18(pretrained=True)
elif args.arch == "resnet50":
net = models.resnet50(pretrained=True)
elif args.arch == "resnet101":
net = models.resnet101(pretrained=True)
elif args.arch == "resnet152":
net = models.resnet152(pretrained=True)
else:
raise ValueError("Network {} not supported".format(args.arch))
if num_classes != 1000:
net = manipulate_net_architecture(model_arch=args.arch,
net=net,
num_classes=num_classes)
# Loss function
if args.loss_function == "ce":
criterion = torch.nn.CrossEntropyLoss()
else:
raise ValueError
# Cuda
if args.use_cuda:
net.cuda()
criterion.cuda()
# Optimizer
momentum = 0.9
decay = 5e-4
optimizer = torch.optim.SGD(net.parameters(),
lr=args.learning_rate,
momentum=momentum,
weight_decay=decay,
nesterov=True)
recorder = RecorderMeter(args.epochs)
start_time = time.time()
epoch_time = AverageMeter()
# Main loop
for epoch in range(args.epochs):
current_learning_rate = adjust_learning_rate(args.learning_rate,
momentum, optimizer,
epoch, args.gammas,
args.schedule)
need_hour, need_mins, need_secs = convert_secs2time(
epoch_time.avg * (args.epochs - epoch))
need_time = '[Need: {:02d}:{:02d}:{:02d}]'.format(
need_hour, need_mins, need_secs)
print_log('\n==>>{:s} [Epoch={:03d}/{:03d}] {:s} [learning_rate={:6.4f}]'.format(time_string(), epoch, args.epochs, need_time, current_learning_rate) \
+ ' [Best : Accuracy={:.2f}, Error={:.2f}]'.format(recorder.max_accuracy(False), 100-recorder.max_accuracy(False)), log)
# train for one epoch
train_acc, train_los = train_model(data_loader=data_train_loader,
model=net,
criterion=criterion,
optimizer=optimizer,
epoch=epoch,
log=log,
print_freq=200,
use_cuda=True)
# evaluate on test set
print_log("Validation on test dataset:", log)
val_acc, val_loss = validate(data_test_loader,
net,
criterion,
log=log,
use_cuda=args.use_cuda)
recorder.update(epoch, train_los, train_acc, val_loss, val_acc)
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': net.state_dict(),
'optimizer': optimizer.state_dict(),
'args': copy.deepcopy(args),
}, model_path, 'checkpoint.pth.tar')
# measure elapsed time
epoch_time.update(time.time() - start_time)
start_time = time.time()
recorder.plot_curve(os.path.join(model_path, 'curve.png'))
log.close()
def main():
args = parse_arguments()
random.seed(args.pretrained_seed)
torch.manual_seed(args.pretrained_seed)
if args.use_cuda:
torch.cuda.manual_seed_all(args.pretrained_seed)
cudnn.benchmark = True
# get the result path to store the results
result_path = get_result_path(dataset_name=args.dataset,
network_arch=args.pretrained_arch,
random_seed=args.pretrained_seed,
result_subfolder=args.result_subfolder,
postfix=args.postfix)
# Init logger
log_file_name = os.path.join(result_path, 'log.txt')
print("Log file: {}".format(log_file_name))
log = open(log_file_name, 'w')
print_log('save path : {}'.format(result_path), log)
state = {k: v for k, v in args._get_kwargs()}
for key, value in state.items():
print_log("{} : {}".format(key, value), log)
print_log("Random Seed: {}".format(args.pretrained_seed), log)
print_log("Python version : {}".format(sys.version.replace('\n', ' ')), log)
print_log("Torch version : {}".format(torch.__version__), log)
print_log("Cudnn version : {}".format(torch.backends.cudnn.version()), log)
_, pretrained_data_test = get_data(args.pretrained_dataset, args.pretrained_dataset)
pretrained_data_test_loader = torch.utils.data.DataLoader(pretrained_data_test,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
##### Dataloader for training ####
num_classes, (mean, std), input_size, num_channels = get_data_specs(args.pretrained_dataset)
data_train, _ = get_data(args.dataset, args.pretrained_dataset)
data_train_loader = torch.utils.data.DataLoader(data_train,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
####################################
# Init model, criterion, and optimizer
print_log("=> Creating model '{}'".format(args.pretrained_arch), log)
# get a path for loading the model to be attacked
model_path = get_model_path(dataset_name=args.pretrained_dataset,
network_arch=args.pretrained_arch,
random_seed=args.pretrained_seed)
model_weights_path = os.path.join(model_path, "checkpoint.pth.tar")
target_network = get_network(args.pretrained_arch,
input_size=input_size,
num_classes=num_classes,
finetune=False)
print_log("=> Network :\n {}".format(target_network), log)
target_network = torch.nn.DataParallel(target_network, device_ids=list(range(args.ngpu)))
# Set the target model into evaluation mode
target_network.eval()
# Imagenet models use the pretrained pytorch weights
if args.pretrained_dataset != "imagenet":
network_data = torch.load(model_weights_path)
target_network.load_state_dict(network_data['state_dict'])
# Set all weights to not trainable
set_parameter_requires_grad(target_network, requires_grad=False)
non_trainale_params = get_num_non_trainable_parameters(target_network)
trainale_params = get_num_trainable_parameters(target_network)
total_params = get_num_parameters(target_network)
print_log("Target Network Trainable parameters: {}".format(trainale_params), log)
print_log("Target Network Non Trainable parameters: {}".format(non_trainale_params), log)
print_log("Target Network Total # parameters: {}".format(total_params), log)
print_log("=> Inserting Generator", log)
generator = UAP(shape=(input_size, input_size),
num_channels=num_channels,
mean=mean,
std=std,
use_cuda=args.use_cuda)
print_log("=> Generator :\n {}".format(generator), log)
non_trainale_params = get_num_non_trainable_parameters(generator)
trainale_params = get_num_trainable_parameters(generator)
total_params = get_num_parameters(generator)
print_log("Generator Trainable parameters: {}".format(trainale_params), log)
print_log("Generator Non Trainable parameters: {}".format(non_trainale_params), log)
print_log("Generator Total # parameters: {}".format(total_params), log)
perturbed_net = nn.Sequential(OrderedDict([('generator', generator), ('target_model', target_network)]))
perturbed_net = torch.nn.DataParallel(perturbed_net, device_ids=list(range(args.ngpu)))
non_trainale_params = get_num_non_trainable_parameters(perturbed_net)
trainale_params = get_num_trainable_parameters(perturbed_net)
total_params = get_num_parameters(perturbed_net)
print_log("Perturbed Net Trainable parameters: {}".format(trainale_params), log)
print_log("Perturbed Net Non Trainable parameters: {}".format(non_trainale_params), log)
print_log("Perturbed Net Total # parameters: {}".format(total_params), log)
# Set the target model into evaluation mode
perturbed_net.module.target_model.eval()
perturbed_net.module.generator.train()
if args.loss_function == "ce":
criterion = torch.nn.CrossEntropyLoss()
elif args.loss_function == "neg_ce":
criterion = NegativeCrossEntropy()
elif args.loss_function == "logit":
criterion = LogitLoss(num_classes=num_classes, use_cuda=args.use_cuda)
elif args.loss_function == "bounded_logit":
criterion = BoundedLogitLoss(num_classes=num_classes, confidence=args.confidence, use_cuda=args.use_cuda)
elif args.loss_function == "bounded_logit_fixed_ref":
criterion = BoundedLogitLossFixedRef(num_classes=num_classes, confidence=args.confidence, use_cuda=args.use_cuda)
elif args.loss_function == "bounded_logit_neg":
criterion = BoundedLogitLoss_neg(num_classes=num_classes, confidence=args.confidence, use_cuda=args.use_cuda)
else:
raise ValueError
if args.use_cuda:
target_network.cuda()
generator.cuda()
perturbed_net.cuda()
criterion.cuda()
optimizer = torch.optim.Adam(perturbed_net.parameters(), lr=state['learning_rate'])
# Measure the time needed for the UAP generation
start = time.time()
train(data_loader=data_train_loader,
model=perturbed_net,
criterion=criterion,
optimizer=optimizer,
epsilon=args.epsilon,
num_iterations=args.num_iterations,
targeted=args.targeted,
target_class=args.target_class,
log=log,
print_freq=args.print_freq,
use_cuda=args.use_cuda)
end = time.time()
print_log("Time needed for UAP generation: {}".format(end - start), log)
# evaluate
print_log("Final evaluation:", log)
metrics_evaluate(data_loader=pretrained_data_test_loader,
target_model=target_network,
perturbed_model=perturbed_net,
targeted=args.targeted,
target_class=args.target_class,
log=log,
use_cuda=args.use_cuda)
save_checkpoint({
'arch' : args.pretrained_arch,
# 'state_dict' : perturbed_net.state_dict(),
'state_dict' : perturbed_net.module.generator.state_dict(),
'optimizer' : optimizer.state_dict(),
'args' : copy.deepcopy(args),
}, result_path, 'checkpoint.pth.tar')
log.close()
test_loss, test_err, test_iou = train_utils.test(
model, loaders['test'], criterion, alphas, betas, biOptimizer,
ngpus, dbsn)
logging.info('Test - Loss: {:.4f} | Acc: {:.4f} | IOU: {:.4f}'.format(
test_loss, 1 - test_err, test_iou))
# time_elapsed = time.time() - since
# logging.info('Total Time {:.0f}m {:.0f}s'.format(
# time_elapsed // 60, time_elapsed % 60))
### Checkpoint ###
if epoch % args.save_freq is 0 or epoch == args.epochs:
logging.info('Saving model at Epoch: {}'.format(epoch))
train_utils.save_checkpoint(dir=args.dir,
epoch=epoch,
state_dict=model.state_dict(),
optimizer=optimizer.state_dict(),
alphas=alphas,
betas=betas)
if args.optimizer == 'RMSProp':
### Adjust Lr ###
if epoch < args.ft_start:
scheduler.step(epoch=epoch)
else:
#scheduler.step(epoch=-1) #reset to args.lr_init for fine-tuning
train_utils.adjust_learning_rate(optimizer, args.ft_lr)
elif args.optimizer == 'SGD':
lr = train_utils.schedule(epoch, args.lr_init, args.epochs)
train_utils.adjust_learning_rate(optimizer, lr)
"Val - Loss: {:.4f} | Acc: {:.4f} | IOU: {:.4f}".format(
val_loss, 1 - val_err, val_iou
)
)
writer.add_scalar("val/loss", val_loss, epoch)
writer.add_scalar("val/error", val_err, epoch)
time_elapsed = time.time() - since
print("Total Time {:.0f}m {:.0f}s\n".format(time_elapsed // 60, time_elapsed % 60))
### Checkpoint ###
if epoch % args.save_freq == 0:
print("Saving model at Epoch: ", epoch)
save_checkpoint(
dir=args.dir,
epoch=epoch,
state_dict=model.state_dict(),
optimizer=optimizer.state_dict(),
)
lr = schedule(
epoch, args.lr_init, args.epochs
)
adjust_learning_rate(optimizer, lr)
writer.add_scalar("hypers/lr", lr, epoch)
### Test set ###
test_loss, test_err, test_iou = train_utils.test(model, loaders["test"], criterion)
print(
"SGD Test - Loss: {:.4f} | Acc: {:.4f} | IOU: {:.4f}".format(
test_loss, 1 - test_err, test_iou
def main():
args = parse_arguments()
random.seed(args.pretrained_seed)
torch.manual_seed(args.pretrained_seed)
if args.use_cuda:
torch.cuda.manual_seed_all(args.pretrained_seed)
cudnn.benchmark = True
# Get data specs
num_classes, (mean, std), input_size, num_channels = get_data_specs(args.pretrained_dataset, args.pretrained_arch)
# Construct the array other classes:
if args.pretrained_dataset in ["imagenet", "ycb"]:
other_classes = args.source_classes
else:
all_classes = np.arange(num_classes)
other_classes = [int(cl) for cl in all_classes if cl not in args.source_classes]
half_batch_size = args.batch_size//2
# get the result path to store the results
result_path = get_result_path(dataset_name=args.pretrained_dataset,
network_arch=args.pretrained_arch,
random_seed=args.pretrained_seed,
result_subfolder=args.result_subfolder,
source_class=args.source_classes,
sink_class=args.sink_classes,
postfix=args.postfix)
# Init logger
log_file_name = os.path.join(result_path, 'log.txt')
print("Log file: {}".format(log_file_name))
log = open(log_file_name, 'w')
print_log('save path : {}'.format(result_path), log)
state = {k: v for k, v in args._get_kwargs()}
for key, value in state.items():
print_log("{} : {}".format(key, value), log)
print_log("Random Seed: {}".format(args.pretrained_seed), log)
print_log("Python version : {}".format(sys.version.replace('\n', ' ')), log)
print_log("Torch version : {}".format(torch.__version__), log)
print_log("Cudnn version : {}".format(torch.backends.cudnn.version()), log)
data_train_sources, data_test_sources = get_data(args.pretrained_dataset,
mean=mean,
std=std,
input_size=input_size,
classes=args.source_classes,
train_samples_per_class=args.num_train_samples_per_class)
data_train_sources_loader = torch.utils.data.DataLoader(data_train_sources,
batch_size=half_batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
data_test_sources_loader = torch.utils.data.DataLoader(data_test_sources,
batch_size=half_batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
data_train_others, data_test_others = get_data(args.pretrained_dataset,
mean=mean,
std=std,
input_size=input_size,
classes=other_classes,
others=True,
train_samples_per_class=args.num_train_samples_per_class)
data_train_others_loader = torch.utils.data.DataLoader(data_train_others,
batch_size=half_batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
data_test_others_loader = torch.utils.data.DataLoader(data_test_others,
batch_size=half_batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
# Init model, criterion, and optimizer
print_log("=> Creating model '{}'".format(args.pretrained_arch), log)
# get a path for loading the model to be attacked
model_path = get_model_path(dataset_name=args.pretrained_dataset,
network_arch=args.pretrained_arch,
random_seed=args.pretrained_seed)
model_weights_path = os.path.join(model_path, "checkpoint.pth.tar")
target_network = get_network(args.pretrained_arch, input_size=input_size, num_classes=num_classes, finetune=args.finetune)
# print_log("=> Network :\n {}".format(target_network), log)
target_network = torch.nn.DataParallel(target_network, device_ids=list(range(args.ngpu)))
# Set the target model into evaluation mode
target_network.eval()
# Imagenet models use the pretrained pytorch weights
if args.pretrained_dataset != "imagenet":
network_data = torch.load(model_weights_path)
target_network.load_state_dict(network_data['state_dict'])
# Set all weights to not trainable
set_parameter_requires_grad(target_network, requires_grad=False)
non_trainale_params = get_num_non_trainable_parameters(target_network)
trainale_params = get_num_trainable_parameters(target_network)
total_params = get_num_parameters(target_network)
print_log("Target Network Trainable parameters: {}".format(trainale_params), log)
print_log("Target Network Non Trainable parameters: {}".format(non_trainale_params), log)
print_log("Target Network Total # parameters: {}".format(total_params), log)
print_log("=> Inserting Generator", log)
generator = UAP(shape=(input_size, input_size),
num_channels=num_channels,
mean=mean,
std=std,
use_cuda=args.use_cuda)
print_log("=> Generator :\n {}".format(generator), log)
non_trainale_params = get_num_non_trainable_parameters(generator)
trainale_params = get_num_trainable_parameters(generator)
total_params = get_num_parameters(generator)
print_log("Generator Trainable parameters: {}".format(trainale_params), log)
print_log("Generator Non Trainable parameters: {}".format(non_trainale_params), log)
print_log("Generator Total # parameters: {}".format(total_params), log)
perturbed_net = nn.Sequential(OrderedDict([('generator', generator), ('target_model', target_network)]))
perturbed_net = torch.nn.DataParallel(perturbed_net, device_ids=list(range(args.ngpu)))
non_trainale_params = get_num_non_trainable_parameters(perturbed_net)
trainale_params = get_num_trainable_parameters(perturbed_net)
total_params = get_num_parameters(perturbed_net)
print_log("Perturbed Net Trainable parameters: {}".format(trainale_params), log)
print_log("Perturbed Net Non Trainable parameters: {}".format(non_trainale_params), log)
print_log("Perturbed Net Total # parameters: {}".format(total_params), log)
# Set the target model into evaluation mode
perturbed_net.module.target_model.eval()
perturbed_net.module.generator.train()
criterion = LossConstructor(source_classes=args.source_classes,
sink_classes=args.sink_classes,
num_classes=num_classes,
source_loss=args.source_loss,
others_loss=args.others_loss,
confidence=args.confidence,
alpha=args.alpha,
use_cuda=args.use_cuda)
if args.use_cuda:
target_network.cuda()
generator.cuda()
perturbed_net.cuda()
criterion.cuda()
optimizer = torch.optim.Adam(perturbed_net.parameters(),
lr=state['learning_rate']) # betas=(0.5, 0.999)
if args.pretrained_dataset not in ["imagenet", "ycb"]:
metrics_evaluate(source_loader=data_test_sources_loader,
others_loader=data_test_others_loader,
target_model=target_network,
perturbed_model=perturbed_net,
source_classes=args.source_classes,
sink_classes=args.sink_classes,
log=log,
use_cuda=args.use_cuda)
start_time = time.time()
train_half_half(sources_data_loader=data_train_sources_loader,
others_data_loader=data_train_others_loader,
model=perturbed_net,
target_model=target_network,
criterion=criterion,
optimizer=optimizer,
epsilon=args.epsilon,
num_iterations=args.num_iterations,
log=log,
print_freq=args.print_freq,
use_cuda=args.use_cuda)
end_time = time.time()
print_log("Elapsed generation time: {}".format(end_time-start_time), log)
# evaluate
print_log("Final evaluation:", log)
metrics_evaluate(source_loader=data_test_sources_loader,
others_loader=data_test_others_loader,
target_model=target_network,
perturbed_model=perturbed_net,
source_classes=args.source_classes,
sink_classes=args.sink_classes,
log=log,
use_cuda=args.use_cuda)
save_checkpoint({
'arch' : args.pretrained_arch,
'state_dict' : perturbed_net.state_dict(),
'optimizer' : optimizer.state_dict(),
'args' : copy.deepcopy(args),
}, result_path, 'checkpoint.pth.tar')
# Plot the adversarial perturbation
uap_numpy = perturbed_net.module.generator.uap.detach().cpu().numpy()
# Calculate the norm
uap_norm = np.linalg.norm(uap_numpy.reshape(-1), np.inf)
print_log("Norm of UAP: {}".format(uap_norm), log)
log.close()
experiment.metrics["loss_lm_" + name].append(tag="train", value=avg_loss)
experiment.metrics["ppl_lm_" + name].append(tag="train",
value=math.exp(avg_loss))
experiment.metrics["loss_lm_" + name].append(tag="val", value=avg_val_loss)
experiment.metrics["ppl_lm_" + name].append(tag="val",
value=math.exp(avg_val_loss))
############################################################
# epoch summary
############################################################
epoch_summary("train", avg_loss)
epoch_summary("val", avg_val_loss)
# after updating all the values, refresh the plots
experiment.update_plots()
# Save the model if the validation loss is the best we've seen so far.
if not best_loss or avg_val_loss < best_loss:
print("saving checkpoint...")
save_checkpoint("{}".format(name),
model,
optimizer,
train_set.vocab,
loss=avg_val_loss,
timestamp=True)
best_loss = avg_val_loss
print()
def main():
args = parse_arguments()
random.seed(args.pretrained_seed)
torch.manual_seed(args.pretrained_seed)
if args.use_cuda:
torch.cuda.manual_seed_all(args.pretrained_seed)
cudnn.benchmark = True
# get a path for saving the model to be trained
model_path = get_model_path(dataset_name=args.pretrained_dataset,
network_arch=args.pretrained_arch,
random_seed=args.pretrained_seed)
# Init logger
log_file_name = os.path.join(model_path, 'log_seed_{}.txt'.format(args.pretrained_seed))
print("Log file: {}".format(log_file_name))
log = open(log_file_name, 'w')
print_log('save path : {}'.format(model_path), log)
state = {k: v for k, v in args._get_kwargs()}
for key, value in state.items():
print_log("{} : {}".format(key, value), log)
print_log("Random Seed: {}".format(args.pretrained_seed), log)
print_log("Python version : {}".format(sys.version.replace('\n', ' ')), log)
print_log("Torch version : {}".format(torch.__version__), log)
print_log("Cudnn version : {}".format(torch.backends.cudnn.version()), log)
# Get data specs
num_classes, (mean, std), input_size, num_channels = get_data_specs(args.pretrained_dataset, args.pretrained_arch)
pretrained_data_train, pretrained_data_test = get_data(args.pretrained_dataset,
mean=mean,
std=std,
input_size=input_size,
train_target_model=True)
pretrained_data_train_loader = torch.utils.data.DataLoader(pretrained_data_train,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
pretrained_data_test_loader = torch.utils.data.DataLoader(pretrained_data_test,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
print_log("=> Creating model '{}'".format(args.pretrained_arch), log)
# Init model, criterion, and optimizer
net = get_network(args.pretrained_arch, input_size=input_size, num_classes=num_classes, finetune=args.finetune)
print_log("=> Network :\n {}".format(net), log)
net = torch.nn.DataParallel(net, device_ids=list(range(args.ngpu)))
non_trainale_params = get_num_non_trainable_parameters(net)
trainale_params = get_num_trainable_parameters(net)
total_params = get_num_parameters(net)
print_log("Trainable parameters: {}".format(trainale_params), log)
print_log("Non Trainable parameters: {}".format(non_trainale_params), log)
print_log("Total # parameters: {}".format(total_params), log)
# define loss function (criterion) and optimizer
criterion_xent = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(net.parameters(), state['learning_rate'], momentum=state['momentum'],
weight_decay=state['decay'], nesterov=True)
if args.use_cuda:
net.cuda()
criterion_xent.cuda()
recorder = RecorderMeter(args.epochs)
# Main loop
start_time = time.time()
epoch_time = AverageMeter()
for epoch in range(args.epochs):
current_learning_rate = adjust_learning_rate(args.learning_rate, args.momentum, optimizer, epoch, args.gammas, args.schedule)
need_hour, need_mins, need_secs = convert_secs2time(epoch_time.avg * (args.epochs-epoch))
need_time = '[Need: {:02d}:{:02d}:{:02d}]'.format(need_hour, need_mins, need_secs)
print_log('\n==>>{:s} [Epoch={:03d}/{:03d}] {:s} [learning_rate={:6.4f}]'.format(time_string(), epoch, args.epochs, need_time, current_learning_rate) \
+ ' [Best : Accuracy={:.2f}, Error={:.2f}]'.format(recorder.max_accuracy(False), 100-recorder.max_accuracy(False)), log)
# train for one epoch
train_acc, train_los = train_target_model(pretrained_data_train_loader, net, criterion_xent, optimizer, epoch, log,
print_freq=args.print_freq,
use_cuda=args.use_cuda)
# evaluate on validation set
print_log("Validation on pretrained test dataset:", log)
val_acc = validate(pretrained_data_test_loader, net, criterion_xent, log, use_cuda=args.use_cuda)
is_best = recorder.update(epoch, train_los, train_acc, 0., val_acc)
save_checkpoint({
'epoch' : epoch + 1,
'arch' : args.pretrained_arch,
'state_dict' : net.state_dict(),
'recorder' : recorder,
'optimizer' : optimizer.state_dict(),
'args' : copy.deepcopy(args),
}, model_path, 'checkpoint.pth.tar')
# measure elapsed time
epoch_time.update(time.time() - start_time)
start_time = time.time()
recorder.plot_curve(os.path.join(model_path, 'curve.png') )
log.close()
def train(self):
# 训练的每个epoch，start_epoch是根据断点来的，在没有断点的时候就是0
for epoch in range(self.start_epoch, self.args.epochs):
# optimizer.param_groups[0]：长度为6的字典，包括[‘amsgrad’, ‘params’, ‘lr’, ‘betas’, ‘weight_decay’, ‘eps’]这6个参数
# 这里不过是获取一下学习率打印出来罢了
lr = self.g_optimizer.param_groups[0]['lr']
print('learning rate = %.7f' % lr)
for i, (a_real,
b_real) in enumerate(zip(self.a_loader, self.b_loader)):
# Generator Computations
training.set_grad([self.Da, self.Db],
False) # 固定D训练G，猜测D的loss和G无关所以不需要锁定G
self.g_optimizer.zero_grad()
a_real = Variable(a_real[0])
b_real = Variable(b_real[0])
a_real, b_real = network.cuda([a_real, b_real])
# Forward pass through generators
a_fake = self.Gab(b_real) # 从b的真图片生成假的a
b_fake = self.Gba(a_real) # 从a的真图片生成假的b
a_recon = self.Gab(b_fake) # 从b的假图片重构a
b_recon = self.Gba(a_fake) # 从a的假图片重构b
a_idt = self.Gab(a_real) # 注意a_idt是a_real喂进去得到的，本来Gab应该是从b转换到a
b_idt = self.Gba(b_real)
# Identity losses
# 这个loss似乎在论文中没有提及，不过也是常用的，意义在于希望转换后基本上保证画面内要素不变
a_idt_loss = self.L1(
a_idt, a_real) * self.args.lamda * self.args.idt_coef
b_idt_loss = self.L1(
b_idt, b_real) * self.args.lamda * self.args.idt_coef
# Adversarial losses
# 对抗loss，从这里可以推断出D输出的是图片为真的概率（注意是G的视角）
a_fake_dis = self.Da(a_fake)
b_fake_dis = self.Db(b_fake)
# 这些假图片判为真的概率应为1
real_label = network.cuda(
Variable(torch.ones(a_fake_dis.size())))
# 按照论文的说法计算MSE作为G的对抗loss
a_gen_loss = self.MSE(a_fake_dis, real_label)
b_gen_loss = self.MSE(b_fake_dis, real_label)
# 循环不变loss
a_cycle_loss = self.L1(a_recon, a_real) * self.args.lamda
b_cycle_loss = self.L1(b_recon, b_real) * self.args.lamda
# 总的生成loss
gen_loss = a_gen_loss + b_gen_loss + a_cycle_loss + b_cycle_loss + a_idt_loss + b_idt_loss
# loss进行反向传播，优化器更新G的参数
gen_loss.backward()
self.g_optimizer.step()
# 开始训练D
training.set_grad([self.Da, self.Db], True) # 解锁D训练D
self.d_optimizer.zero_grad()
# buffer对象是可以call的，就是将新生成的一波假图到buffer中涮一下拿出一些给D更新
a_fake = Variable(
torch.Tensor(
self.a_fake_sample([a_fake.cpu().data.numpy()])[0]))
b_fake = Variable(
torch.Tensor(
self.b_fake_sample([b_fake.cpu().data.numpy()])[0]))
a_fake, b_fake = network.cuda([a_fake, b_fake])
# 让两个D分别对真实样本和生成样本进行判断
a_real_dis = self.Da(a_real)
a_fake_dis = self.Da(a_fake)
b_real_dis = self.Db(b_real)
b_fake_dis = self.Db(b_fake)
real_label = network.cuda(
Variable(torch.ones(a_real_dis.size())))
fake_label = network.cuda(
Variable(torch.zeros(a_fake_dis.size())))
# 这些loss都和标签都用MSE
a_dis_real_loss = self.MSE(a_real_dis, real_label)
a_dis_fake_loss = self.MSE(a_fake_dis, fake_label)
b_dis_real_loss = self.MSE(b_real_dis, real_label)
b_dis_fake_loss = self.MSE(b_fake_dis, fake_label)
# Total discriminators losses
a_dis_loss = (a_dis_real_loss + a_dis_fake_loss) * 0.5
b_dis_loss = (b_dis_real_loss + b_dis_fake_loss) * 0.5
# loss进行反向传播，优化器更新D的参数
a_dis_loss.backward()
b_dis_loss.backward()
self.d_optimizer.step()
print(
"Epoch: (%3d) (%5d/%5d) | Gen Loss:%.2e | Dis Loss:%.2e" %
(epoch, i + 1, min(len(self.a_loader), len(
self.b_loader)), gen_loss, a_dis_loss + b_dis_loss))
# 保存最新的断点
training.save_checkpoint(
{
'epoch': epoch + 1,
'Da': self.Da.state_dict(),
'Db': self.Db.state_dict(),
'Gab': self.Gab.state_dict(),
'Gba': self.Gba.state_dict(),
'd_optimizer': self.d_optimizer.state_dict(),
'g_optimizer': self.g_optimizer.state_dict()
}, '%s/latest.ckpt' % self.args.checkpoint_dir)
# 更新学习率
self.g_lr_scheduler.step()
self.d_lr_scheduler.step()