def validate(val_loader, model, criterion, log=None, use_cuda=True): top1 = AverageMeter() top5 = AverageMeter() # switch to evaluate mode model.eval() for i, (input, target) in enumerate(val_loader): if use_cuda: target = target.cuda() input = input.cuda() with torch.no_grad(): # compute output output = model(input) # measure accuracy and record loss prec1, prec5 = accuracy(output.data, target, topk=(1, 5)) top1.update(prec1.item(), input.size(0)) top5.update(prec5.item(), input.size(0)) if log: print_log(' **Test** Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f} Error@1 {error1:.3f}'.format(top1=top1, top5=top5, error1=100-top1.avg), log) else: print(' **Test** Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f} Error@1 {error1:.3f}'.format(top1=top1, top5=top5, error1=100-top1.avg)) return top1.avg
def save_checkpoint(self, val_loss, model, optimizer, log): '''Saves model when validation loss decrease.''' if self.verbose: print_log(( f'Validation loss decreased ({self.val_loss_min:.6f} --> {val_loss:.6f}). Saving model ...' ), log) state = { 'model': model.state_dict(), 'optimizer': optimizer.state_dict() } torch.save(state, self.save_name) self.val_loss_min = val_loss
def train_model(data_loader, model, criterion, optimizer, epoch, log, print_freq=200, use_cuda=True): # train function (forward, backward, update) batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter() top1 = AverageMeter() # switch to train mode model.train() end = time.time() for iteration, (input, target) in enumerate(data_loader): # measure data loading time data_time.update(time.time() - end) if use_cuda: target = target.cuda() input = input.cuda() # compute output output = model(input) loss = criterion(output, target) # measure accuracy and record loss if len(target.shape) > 1: target = torch.argmax(target, dim=-1) prec1, = accuracy(output.data, target, topk=(1,)) losses.update(loss.item(), input.size(0)) top1.update(prec1.item(), input.size(0)) # compute gradient and do SGD step optimizer.zero_grad() loss.backward() optimizer.step() # measure elapsed time batch_time.update(time.time() - end) end = time.time() if iteration % print_freq == 0: print_log(' Epoch: [{:03d}][{:03d}/{:03d}] ' 'Time {batch_time.val:.3f} ({batch_time.avg:.3f}) ' 'Data {data_time.val:.3f} ({data_time.avg:.3f}) ' 'Loss {loss.val:.4f} ({loss.avg:.4f}) ' 'Prec@1 {top1.val:.3f} ({top1.avg:.3f}) '.format( epoch, iteration, len(data_loader), batch_time=batch_time, data_time=data_time, loss=losses, top1=top1) + time_string(), log) print_log(' **Train** Prec@1 {top1.avg:.3f} Error@1 {error1:.3f}'.format(top1=top1, error1=100-top1.avg), log) return top1.avg, losses.avg
def test(args, model, device, test_loader): model.eval() test_loss = 0 correct = 0 with torch.no_grad(): for data, target in test_loader: data, target = data.to(device), target.to(device) output = model(data) test_loss += F.nll_loss(output, target, reduction='sum').item() # sum up batch loss pred = output.argmax(dim=1, keepdim=True) # get the index of the max log-probability correct += pred.eq(target.view_as(pred)).sum().item() test_loss /= len(test_loader.dataset) utils.print_log('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format( test_loss, correct, len(test_loader.dataset), 100. * correct / len(test_loader.dataset)), args)
def train(args, model, device, train_loader, optimizer, epoch): """ Training the classifier. """ model.train() for batch_idx, (data, target) in enumerate(train_loader): data, target = data.to(device), target.to(device) optimizer.zero_grad() output = model(data) loss = F.nll_loss(output, target) loss.backward() optimizer.step() if batch_idx % args.log_interval == 0: utils.print_log('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( epoch, batch_idx * len(data), len(train_loader.dataset), 100. * batch_idx / len(train_loader), loss.item()), args) if args.dry_run: break
def __call__(self, val_loss, model, optimizer, log): score = -val_loss if self.best_score is None: self.best_score = score self.save_checkpoint(val_loss, model, optimizer, log) elif score < self.best_score - self.delta: self.counter += 1 print_log(( f'EarlyStopping counter: {self.counter} out of {self.patience}' ), log) if self.counter >= self.patience: self.early_stop = True else: self.best_score = score self.save_checkpoint(val_loss, model, optimizer, log) self.counter = 0 return self.early_stop
# utils.load_state_dict(model, weight_file) # model.fc.reset_parameters() criterion = nn.CrossEntropyLoss() if cuda: model = model.cuda() criterion = criterion.cuda() optim = torch.optim.SGD(model.parameters(), lr=train_cfg['lr'], momentum=train_cfg['momentum'], weight_decay=train_cfg['weight_decay']) # lr_scheduler = torch.optim.lr_scheduler.StepLR(optim, train_cfg['step_size'], gamma=train_cfg['gamma'], last_epoch=-1) lr_scheduler = None log = utils.print_log(configure['log_dir'], [net_cfg['type'], timestamp]) log.write(str(net_cfg)) log.write(str(train_cfg)) epoch_time = utils.AverageMeter() batch_time = utils.AverageMeter() data_time = utils.AverageMeter() losses = utils.AverageMeter() top1 = utils.AverageMeter() # --------------------train & validation & save checkpoint---------------- # epoch = 0 last_iteration = 0 print_freq = 1 best_top1 = 0 max_epoch = train_cfg['max_epoch'] print("train max epoch {0}".format(max_epoch))
def train(): train_data = ACNet_data.SUNRGBD(transform=transforms.Compose([ ACNet_data.scaleNorm(), ACNet_data.RandomScale((1.0, 1.4)), ACNet_data.RandomHSV((0.9, 1.1), (0.9, 1.1), (25, 25)), ACNet_data.RandomCrop(image_h, image_w), ACNet_data.RandomFlip(), ACNet_data.ToTensor(), ACNet_data.Normalize() ]), phase_train=True, data_dir=args.data_dir) train_loader = DataLoader(train_data, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=False) num_train = len(train_data) if args.last_ckpt: model = ACNet_models_V1.ACNet(num_class=40, pretrained=False) else: model = ACNet_models_V1.ACNet(num_class=40, pretrained=True) if torch.cuda.device_count() > 1: print("Let's use", torch.cuda.device_count(), "GPUs!") model = nn.DataParallel(model) # CEL_weighted = utils.CrossEntropyLoss2d() CEL_weighted = utils.FocalLoss2d(weight=nyuv2_frq, gamma=2) model.train() model.to(device) CEL_weighted.to(device) optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) global_step = 0 if args.last_ckpt: global_step, args.start_epoch = load_ckpt(model, optimizer, args.last_ckpt, device) #hxx for finetuing # lr_decay_lambda = lambda epoch: 0.2 * args.lr_decay_rate ** ((epoch - args.start_epoch) // args.lr_epoch_per_decay) lr_decay_lambda = lambda epoch: (1 - (epoch - args.start_epoch) / (args.epochs - args.start_epoch))**0.9 scheduler = LambdaLR(optimizer, lr_lambda=lr_decay_lambda) writer = SummaryWriter(args.summary_dir) for epoch in range(int(args.start_epoch), args.epochs): # if (epoch - args.start_epoch) % args.lr_epoch_per_decay == 0: scheduler.step(epoch) local_count = 0 last_count = 0 end_time = time.time() if epoch % args.save_epoch_freq == 0 and epoch != args.start_epoch: save_ckpt(args.ckpt_dir, model, optimizer, global_step, epoch, local_count, num_train) for batch_idx, sample in enumerate(train_loader): image = sample['image'].to(device) depth = sample['depth'].to(device) target_scales = [ sample[s].to(device) for s in ['label', 'label2', 'label3', 'label4', 'label5'] ] optimizer.zero_grad() pred_scales = model(image, depth, args.checkpoint) loss = CEL_weighted(pred_scales, target_scales) loss.backward() optimizer.step() local_count += image.data.shape[0] global_step += 1 if global_step % args.print_freq == 0 or global_step == 1: time_inter = time.time() - end_time count_inter = local_count - last_count print_log(global_step, epoch, local_count, count_inter, num_train, loss, time_inter) end_time = time.time() last_count = local_count save_ckpt(args.ckpt_dir, model, optimizer, global_step, args.epochs, 0, num_train) print("Training completed ")
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()
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()
def metrics_evaluate(data_loader, target_model, perturbed_model, targeted, target_class, log=None, use_cuda=True): # switch to evaluate mode target_model.eval() perturbed_model.eval() perturbed_model.module.generator.eval() perturbed_model.module.target_model.eval() clean_acc = AverageMeter() perturbed_acc = AverageMeter() attack_success_rate = AverageMeter() # Among the correctly classified samples, the ratio of being different from clean prediction (same as gt) if targeted: all_to_target_success_rate = AverageMeter() # The ratio of samples going to the sink classes all_to_target_success_rate_filtered = AverageMeter() total_num_samples = 0 num_same_classified = 0 num_diff_classified = 0 for input, gt in data_loader: if use_cuda: gt = gt.cuda() input = input.cuda() # compute output with torch.no_grad(): clean_output = target_model(input) pert_output = perturbed_model(input) correctly_classified_mask = torch.argmax(clean_output, dim=-1).cpu() == gt.cpu() cl_acc = accuracy(clean_output.data, gt, topk=(1,)) clean_acc.update(cl_acc[0].item(), input.size(0)) pert_acc = accuracy(pert_output.data, gt, topk=(1,)) perturbed_acc.update(pert_acc[0].item(), input.size(0)) # Calculating Fooling Ratio params clean_out_class = torch.argmax(clean_output, dim=-1) pert_out_class = torch.argmax(pert_output, dim=-1) total_num_samples += len(clean_out_class) num_same_classified += torch.sum(clean_out_class == pert_out_class).cpu().numpy() num_diff_classified += torch.sum(~(clean_out_class == pert_out_class)).cpu().numpy() if torch.sum(correctly_classified_mask)>0: with torch.no_grad(): pert_output_corr_cl = perturbed_model(input[correctly_classified_mask]) attack_succ_rate = accuracy(pert_output_corr_cl, gt[correctly_classified_mask], topk=(1,)) attack_success_rate.update(attack_succ_rate[0].item(), pert_output_corr_cl.size(0)) # Calculate Absolute Accuracy Drop aad_source = clean_acc.avg - perturbed_acc.avg # Calculate Relative Accuracy Drop if clean_acc.avg !=0: rad_source = (clean_acc.avg - perturbed_acc.avg)/clean_acc.avg * 100. else: rad_source = 0. # Calculate fooling ratio fooling_ratio = num_diff_classified/total_num_samples * 100. if targeted: # 2. How many of all samples go the sink class (Only relevant for others loader) target_cl = torch.ones_like(gt) * target_class all_to_target_succ_rate = accuracy(pert_output, target_cl, topk=(1,)) all_to_target_success_rate.update(all_to_target_succ_rate[0].item(), pert_output.size(0)) # 3. How many of all samples go the sink class, except gt sink class (Only relevant for others loader) # Filter all idxs which are not belonging to sink class non_target_class_idxs = [i != target_class for i in gt] non_target_class_mask = torch.Tensor(non_target_class_idxs)==True if torch.sum(non_target_class_mask)>0: gt_non_target_class = gt[non_target_class_mask] pert_output_non_target_class = pert_output[non_target_class_mask] target_cl = torch.ones_like(gt_non_target_class) * target_class all_to_target_succ_rate_filtered = accuracy(pert_output_non_target_class, target_cl, topk=(1,)) all_to_target_success_rate_filtered.update(all_to_target_succ_rate_filtered[0].item(), pert_output_non_target_class.size(0)) if log: print_log('\n\t#######################', log) print_log('\tClean model accuracy: {:.3f}'.format(clean_acc.avg), log) print_log('\tPerturbed model accuracy: {:.3f}'.format(perturbed_acc.avg), log) print_log('\tAbsolute Accuracy Drop: {:.3f}'.format(aad_source), log) print_log('\tRelative Accuracy Drop: {:.3f}'.format(rad_source), log) print_log('\tAttack Success Rate: {:.3f}'.format(100-attack_success_rate.avg), log) print_log('\tFooling Ratio: {:.3f}'.format(fooling_ratio), log) if targeted: print_log('\tAll --> Target Class {} Prec@1 {:.3f}'.format(target_class, all_to_target_success_rate.avg), log) print_log('\tAll (w/o sink samples) --> Sink {} Prec@1 {:.3f}'.format(target_class, all_to_target_success_rate_filtered.avg), log)
def train(data_loader, model, criterion, optimizer, epsilon, num_iterations, targeted, target_class, log, print_freq=200, use_cuda=True): # train function (forward, backward, update) batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter() top1 = AverageMeter() top5 = AverageMeter() # switch to train mode model.module.generator.train() model.module.target_model.eval() end = time.time() data_iterator = iter(data_loader) iteration=0 while (iteration<num_iterations): try: input, target = next(data_iterator) except StopIteration: # StopIteration is thrown if dataset ends # reinitialize data loader data_iterator = iter(data_loader) input, target = next(data_iterator) if targeted: target = torch.ones(input.shape[0], dtype=torch.int64) * target_class # measure data loading time data_time.update(time.time() - end) if use_cuda: target = target.cuda() input = input.cuda() # compute output if model.module._get_name() == "Inception3": output, aux_output = model(input) loss1 = criterion(output, target) loss2 = criterion(aux_output, target) loss = loss1 + 0.4*loss2 else: output = model(input) loss = criterion(output, target) # measure accuracy and record loss if len(target.shape) > 1: target = torch.argmax(target, dim=-1) prec1, prec5 = accuracy(output.data, target, topk=(1, 5)) losses.update(loss.item(), input.size(0)) top1.update(prec1.item(), input.size(0)) top5.update(prec5.item(), input.size(0)) # compute gradient and do SGD step optimizer.zero_grad() loss.backward() optimizer.step() # Projection model.module.generator.uap.data = torch.clamp(model.module.generator.uap.data, -epsilon, epsilon) # measure elapsed time batch_time.update(time.time() - end) end = time.time() if iteration % print_freq == 0: print_log(' Iteration: [{:03d}/{:03d}] ' 'Time {batch_time.val:.3f} ({batch_time.avg:.3f}) ' 'Data {data_time.val:.3f} ({data_time.avg:.3f}) ' 'Loss {loss.val:.4f} ({loss.avg:.4f}) ' 'Prec@1 {top1.val:.3f} ({top1.avg:.3f}) ' 'Prec@5 {top5.val:.3f} ({top5.avg:.3f}) '.format( iteration, num_iterations, batch_time=batch_time, data_time=data_time, loss=losses, top1=top1, top5=top5) + time_string(), log) iteration+=1 print_log(' **Train** Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f} Error@1 {error1:.3f}'.format(top1=top1, top5=top5, error1=100-top1.avg), log)
def train_half_half(sources_data_loader, others_data_loader, model, target_model, criterion, optimizer, epsilon, num_iterations, log, print_freq=200, use_cuda=True, patch=False): # train function (forward, backward, update) batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter() top1 = AverageMeter() top5 = AverageMeter() # switch to train mode model.module.generator.train() model.module.target_model.eval() target_model.eval() end = time.time() sources_data_iterator = iter(sources_data_loader) others_data_iterator = iter(others_data_loader) iteration=0 while (iteration<num_iterations): try: sources_input, sources_target = next(sources_data_iterator) except StopIteration: # StopIteration is thrown if dataset ends # reinitialize data loader sources_data_iterator = iter(sources_data_loader) sources_input, sources_target = next(sources_data_iterator) try: others_input, others_target = next(others_data_iterator) except StopIteration: # StopIteration is thrown if dataset ends # reinitialize data loader others_data_iterator = iter(others_data_loader) others_input, others_target = next(others_data_iterator) # Concat the two batches input = torch.cat([sources_input, others_input], dim=0) target = torch.cat([sources_target, others_target], dim=0) # measure data loading time data_time.update(time.time() - end) if use_cuda: target = target.cuda() input = input.cuda() # compute output output = model(input) target_model_output = target_model(input) loss = criterion(output, target_model_output, target) # measure accuracy and record loss if len(target.shape) > 1: target = torch.argmax(target, dim=-1) prec1, prec5 = accuracy(output.data, target, topk=(1, 5)) losses.update(loss.item(), input.size(0)) top1.update(prec1.item(), input.size(0)) top5.update(prec5.item(), input.size(0)) # compute gradient and do SGD step optimizer.zero_grad() loss.backward() optimizer.step() # Project to l-infinity ball if patch: model.module.generator.uap.data = torch.clamp(model.module.generator.uap.data, 0, epsilon) else: model.module.generator.uap.data = torch.clamp(model.module.generator.uap.data, -epsilon, epsilon) # measure elapsed time batch_time.update(time.time() - end) end = time.time() if iteration % print_freq == 0: print_log(' Iteration: [{:03d}/{:03d}] ' 'Time {batch_time.val:.3f} ({batch_time.avg:.3f}) ' 'Data {data_time.val:.3f} ({data_time.avg:.3f}) ' 'Loss {loss.val:.4f} ({loss.avg:.4f}) ' 'Prec@1 {top1.val:.3f} ({top1.avg:.3f}) ' 'Prec@5 {top5.val:.3f} ({top5.avg:.3f}) '.format( iteration, num_iterations, batch_time=batch_time, data_time=data_time, loss=losses, top1=top1, top5=top5) + time_string(), log) iteration+=1 print_log(' **Train** Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f} Error@1 {error1:.3f}'.format(top1=top1, top5=top5, error1=100-top1.avg), log)
def metrics_evaluate(source_loader, others_loader, target_model, perturbed_model, source_classes, sink_classes, log=None, use_cuda=True): if len(sink_classes)!=0: assert len(source_classes) == len(sink_classes) # switch to evaluate mode target_model.eval() perturbed_model.eval() for loader, loader_name in zip([source_loader, others_loader],["Source", "Others"]): clean_acc = AverageMeter() perturbed_acc = AverageMeter() attack_success_rate = AverageMeter() # Among the correctly classified samples, the ratio of being different from clean prediction (same as gt) if len(sink_classes)!=0: all_to_sink_success_rate = [] source_to_sink_success_rate = [] all_to_sink_success_rate_filtered = [] for i in range(len(sink_classes)): all_to_sink_success_rate.append(AverageMeter()) # The ratio of samples going to the sink classes source_to_sink_success_rate.append(AverageMeter()) all_to_sink_success_rate_filtered.append(AverageMeter()) total_num_samples = 0 num_same_classified = 0 num_diff_classified = 0 if len(loader)>0: # For UAP, all classes will be attacked, so others_loader is empty for input, gt in loader: if use_cuda: gt = gt.cuda() input = input.cuda() # compute output with torch.no_grad(): clean_output = target_model(input) pert_output = perturbed_model(input) correctly_classified_mask = torch.argmax(clean_output, dim=-1).cpu() == gt.cpu() cl_acc = accuracy(clean_output.data, gt, topk=(1,)) clean_acc.update(cl_acc[0].item(), input.size(0)) pert_acc = accuracy(pert_output.data, gt, topk=(1,)) perturbed_acc.update(pert_acc[0].item(), input.size(0)) # Calculating Fooling Ratio params clean_out_class = torch.argmax(clean_output, dim=-1) pert_out_class = torch.argmax(pert_output, dim=-1) total_num_samples += len(clean_out_class) num_same_classified += torch.sum(clean_out_class == pert_out_class).cpu().numpy() num_diff_classified += torch.sum(~(clean_out_class == pert_out_class)).cpu().numpy() if torch.sum(correctly_classified_mask)>0: with torch.no_grad(): pert_output_corr_cl = perturbed_model(input[correctly_classified_mask]) attack_succ_rate = accuracy(pert_output_corr_cl, gt[correctly_classified_mask], topk=(1,)) attack_success_rate.update(attack_succ_rate[0].item(), pert_output_corr_cl.size(0)) # Collect samples from Source go to sink if len(sink_classes)!=0: # Iterate over source class and sink class pairs for cl_idx, (source_cl, sink_cl) in enumerate(zip(source_classes, sink_classes)): # 1. Check how many of the paired source class got to the sink class (Only relevant for source loader) # Filter all idxs which belong to the source class source_cl_idxs = [i == source_cl for i in gt] source_cl_mask = torch.Tensor(source_cl_idxs)==True if torch.sum(source_cl_mask)>0: gt_source_cl = gt[source_cl_mask] pert_output_source_cl = pert_output[source_cl_mask] # Create desired target value target_sink = torch.ones_like(gt_source_cl) * sink_cl source_to_sink_succ_rate = accuracy(pert_output_source_cl, target_sink, topk=(1,)) source_to_sink_success_rate[cl_idx].update(source_to_sink_succ_rate[0].item(), pert_output_source_cl.size(0)) # 2. How many of all samples go the sink class (Only relevant for others loader) target_sink = torch.ones_like(gt) * sink_cl all_to_sink_succ_rate = accuracy(pert_output, target_sink, topk=(1,)) all_to_sink_success_rate[cl_idx].update(all_to_sink_succ_rate[0].item(), pert_output.size(0)) # 3. How many of all samples go the sink class, except gt sink class (Only relevant for others loader) # Filter all idxs which are not belonging to sink class non_sink_class_idxs = [i != sink_cl for i in gt] non_sink_class_mask = torch.Tensor(non_sink_class_idxs)==True if torch.sum(non_sink_class_mask)>0: gt_non_sink_class = gt[non_sink_class_mask] pert_output_non_sink_class = pert_output[non_sink_class_mask] target_sink = torch.ones_like(gt_non_sink_class) * sink_cl all_to_sink_succ_rate_filtered = accuracy(pert_output_non_sink_class, target_sink, topk=(1,)) all_to_sink_success_rate_filtered[cl_idx].update(all_to_sink_succ_rate_filtered[0].item(), pert_output_non_sink_class.size(0)) if log: print_log('\n\t########## {} #############'.format(loader_name), log) if len(sink_classes)!=0: for cl_idx, (source_cl, sink_cl) in enumerate(zip(source_classes, sink_classes)): print_log('\n\tSource {} --> Sink {} Prec@1 {:.3f}'.format(source_cl, sink_cl, source_to_sink_success_rate[cl_idx].avg), log) print_log('\tAll --> Sink {} Prec@1 {:.3f}'.format(sink_cl, all_to_sink_success_rate[cl_idx].avg), log) # Average fooling ratio of the non-source classes into the target label print_log('\tAll (w/o sink samples) --> Sink {} Prec@1 {:.3f}'.format(sink_cl, all_to_sink_success_rate_filtered[cl_idx].avg), log) else: print('\n\t########## {} #############'.format(loader_name)) if len(sink_classes)!=0: for cl_idx, (source_cl, sink_cl) in enumerate(zip(source_classes, sink_classes)): print('\n\tSource {} --> Sink {} Prec@1 {:.3f}'.format(source_cl, sink_cl, source_to_sink_success_rate[cl_idx].avg)) print('\tAll --> Sink {} Prec@1 {:.3f}'.format(sink_cl, all_to_sink_success_rate[cl_idx].avg)) # Average fooling ratio of the non-source classes into the target label print('\tAll (w/o sink samples) --> Sink {} Prec@1 {:.3f}'.format(sink_cl, all_to_sink_success_rate_filtered[cl_idx].avg))
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(): train_data = ACNet_data.SUNRGBD(transform=transforms.Compose([ACNet_data.scaleNorm(), ACNet_data.RandomScale((1.0, 1.4)), ACNet_data.RandomHSV((0.9, 1.1), (0.9, 1.1), (25, 25)), ACNet_data.RandomCrop(image_h, image_w), ACNet_data.RandomFlip(), ACNet_data.ToTensor(), ACNet_data.Normalize()]), phase_train=True, data_dir=args.data_dir) train_loader = DataLoader(train_data, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=False) num_train = len(train_data) if args.last_ckpt: model = ACNet_models_V1.ACNet(num_class=40, pretrained=False) else: model = ACNet_models_V1.ACNet(num_class=40, pretrained=True) if torch.cuda.device_count() > 1: print("Let's use", torch.cuda.device_count(), "GPUs!") model = nn.DataParallel(model) CEL_weighted = utils.CrossEntropyLoss2d(weight=nyuv2_frq) model.train() model.to(device) CEL_weighted.to(device) optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) global_step = 0 if args.last_ckpt: global_step, args.start_epoch = load_ckpt(model, optimizer, args.last_ckpt, device) lr_decay_lambda = lambda epoch: args.lr_decay_rate ** (epoch // args.lr_epoch_per_decay) scheduler = LambdaLR(optimizer, lr_lambda=lr_decay_lambda) writer = SummaryWriter(args.summary_dir) for epoch in range(int(args.start_epoch), args.epochs): scheduler.step(epoch) local_count = 0 last_count = 0 end_time = time.time() if epoch % args.save_epoch_freq == 0 and epoch != args.start_epoch: save_ckpt(args.ckpt_dir, model, optimizer, global_step, epoch, local_count, num_train) for batch_idx, sample in enumerate(train_loader): image = sample['image'].to(device) depth = sample['depth'].to(device) target_scales = [sample[s].to(device) for s in ['label', 'label2', 'label3', 'label4', 'label5']] optimizer.zero_grad() pred_scales = model(image, depth, args.checkpoint) loss = CEL_weighted(pred_scales, target_scales) loss.backward() optimizer.step() local_count += image.data.shape[0] global_step += 1 if global_step % args.print_freq == 0 or global_step == 1: time_inter = time.time() - end_time count_inter = local_count - last_count print_log(global_step, epoch, local_count, count_inter, num_train, loss, time_inter) end_time = time.time() for name, param in model.named_parameters(): writer.add_histogram(name, param.clone().cpu().data.numpy(), global_step, bins='doane') grid_image = make_grid(image[:3].clone().cpu().data, 3, normalize=True) writer.add_image('image', grid_image, global_step) grid_image = make_grid(depth[:3].clone().cpu().data, 3, normalize=True) writer.add_image('depth', grid_image, global_step) grid_image = make_grid(utils.color_label(torch.max(pred_scales[0][:3], 1)[1] + 1), 3, normalize=False, range=(0, 255)) writer.add_image('Predicted label', grid_image, global_step) grid_image = make_grid(utils.color_label(target_scales[0][:3]), 3, normalize=False, range=(0, 255)) writer.add_image('Groundtruth label', grid_image, global_step) writer.add_scalar('CrossEntropyLoss', loss.data, global_step=global_step) writer.add_scalar('Learning rate', scheduler.get_lr()[0], global_step=global_step) last_count = local_count save_ckpt(args.ckpt_dir, model, optimizer, global_step, args.epochs, 0, num_train) print("Training completed ")
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()
# Average PSNR on one epoch train_data train_psnr = sum(psnr_list) / len(psnr_list) train_loss = sum(running_loss) / len(running_loss) print('MSE loss: {0:.4f}, SSIM loss: {1:.4f}, SSIM linear loss: {2:.4f}'. format(np.mean(loss_mse_list), np.mean(loss_ssim_list), np.mean(loss_ssim_lin_list))) # save the network parameters torch.save( net.state_dict(), './weight_checkpoint/rawvsr_{}_{}_{}_{}_{}_raw.pkl'.format( opt['train']['scale'], opt['network']['nframes'], opt['network']['nf'], opt['network']['back_RBs'], epoch)) # use evaluation models during the net evaluating net.eval() test_psnr = test(net, test_data_loader, device, dataset_name) one_epoch_time = time.time() - start_time print_log(epoch + 1, num_epochs, one_epoch_time, train_psnr, test_psnr, dataset_name) if test_psnr >= previous_test_psnr: torch.save( net.state_dict(), './weight_checkpoint/rawvsr_{}_{}_{}_{}_raw_best.pkl'.format( opt['train']['scale'], opt['network']['nframes'], opt['network']['nf'], opt['network']['back_RBs'])) previous_test_psnr = test_psnr