def main(): # Reproducibility np.random.seed(cfg['seed']) torch.manual_seed(cfg['seed']) # Model & Optimizer model = getModel(cfg['model']) start_epoch = 1 max_epoch = 1 # Load model and optimizer if cfg['load_ckpt'] != '': checkpoint = torch.load(cfg['load_ckpt'], map_location="cpu") model.load_state_dict(checkpoint['model_state']) print("load model on '{}' is complete.".format(cfg['load_ckpt'])) cudnn.benchmark = True # Data Loader in_valid_loader = getDataLoader(ds_cfg=cfg['in_dataset'], dl_cfg=cfg['dataloader'], split="valid") if 'targets' in cfg['in_dataset'].keys(): exp_dir = os.path.join(cfg['exp_root'], cfg['exp_dir'], "logits", cfg['in_dataset']['dataset'], cfg['in_dataset']['targets'][0]) else: exp_dir = os.path.join(cfg['exp_root'], cfg['exp_dir'], "logits", cfg['in_dataset']['dataset']) # Result directory and make tensorboard event file if os.path.exists(exp_dir) is False: os.makedirs(exp_dir) # Outlier detector detector_func = detectors.getDetector(cfg['detector']) global_cfg['detector'] = cfg['detector'] # Outlier detector print("=======================IMPORTANT CONFIG=======================") print(" Model : {}\n \ Detector : {}\n".format(cfg['model']['network_kind'], cfg['detector']['detector'])) print( "========Start logits extraction for GODIN. Result will be saved in {}" .format(exp_dir)) valid_summary = valid_epoch_wo_outlier(model, in_valid_loader, detector_func) summary_log = "Acc [{}]\n".format(valid_summary['classifier_acc']) print(summary_log) torch.save(valid_summary['f_confidences'], os.path.join(exp_dir, 'f_confidences.pt')) torch.save(valid_summary['h_confidences'], os.path.join(exp_dir, 'h_confidences.pt')) torch.save(valid_summary['targets'], os.path.join(exp_dir, 'targets.pt')) torch.save(valid_summary['logits'], os.path.join(exp_dir, 'logits.pt'))
def main(): # Reproducibility np.random.seed(cfg['seed']) torch.manual_seed(cfg['seed']) # Model & Optimizer model = getModel(cfg['model']) start_epoch = 1 max_epoch = 1 # Load model and optimizer if cfg['load_ckpt'] != '': checkpoint = torch.load(cfg['load_ckpt'], map_location="cpu") model.load_state_dict(checkpoint['model_state']) print("load model on '{}' is complete.".format(cfg['load_ckpt'])) cudnn.benchmark = True # Data Loader in_valid_loader = getDataLoader(ds_cfg=cfg['in_dataset'], dl_cfg=cfg['dataloader'], split="valid") exp_dir = os.path.join(cfg['exp_root'], cfg['exp_dir']) # Result directory and make tensorboard event file if os.path.exists(exp_dir) is False: os.makedirs(exp_dir) # Loss function loss_func = losses.getLoss(cfg['loss']) # Outlier detector print("=======================IMPORTANT CONFIG=======================") print(" Model : {}\n \ Loss : {}\n".format(cfg['model']['network_kind'], cfg['loss']['loss'])) print("========Start epsilon regression for GODIN. Result will be saved in {}".format(exp_dir)) logfile = open(os.path.join(exp_dir, "epsilon.txt"), "w") epsilon = 0.0 max_epsilon = 1.0 step = 0.1 while epsilon <= max_epsilon: valid_summary = valid_epoch_wo_outlier(model, in_valid_loader, epsilon) summary_log = "Epsilon [{}] => Score [{}]\n".format(epsilon, valid_summary['score']) print(summary_log) logfile.write(summary_log) epsilon += step logfile.close()
def main(): global global_cfg # Reproducibility np.random.seed(cfg['seed']) torch.manual_seed(cfg['seed']) # Model & Optimizer model = getModel(cfg['model']) model.rot_head = nn.Linear(model.nChannels, 4) model.rot_head.cuda() optimizer = optim.getOptimizer(model, cfg['optim']) start_epoch = 1 # Load model and optimizer if cfg['load_ckpt'] != '': checkpoint = torch.load(cfg['load_ckpt'], map_location="cpu") model.load_state_dict(checkpoint['model_state']) print("load model on '{}' is complete.".format(cfg['load_ckpt'])) if not cfg['finetuning']: optimizer.load_state_dict(checkpoint['optimizer_state']) if 'epoch' in checkpoint.keys() and not cfg['finetuning']: start_epoch = checkpoint['epoch'] print("Restore epoch {}".format(start_epoch)) else: start_epoch = 1 cudnn.benchmark = True # Data Loader in_train_loader = getDataLoader(ds_cfg=cfg['in_dataset'], dl_cfg=cfg['dataloader'], split="train") in_valid_loader = getDataLoader(ds_cfg=cfg['in_dataset'], dl_cfg=cfg['dataloader'], split="valid") attack_in = None if 'PGD' in cfg.keys() and cfg['PGD'] is not None: attack_in = RotPGDAttack(model=model, eps=cfg['PGD']['epsilon'], nb_iter=cfg['PGD']['iters'], eps_iter=cfg['PGD']['iter_size'], rand_init=True, loss_func='CE') if cfg['out_dataset'] is not None: out_train_loader = getDataLoader(ds_cfg=cfg['out_dataset'], dl_cfg=cfg['dataloader'], split="train") out_valid_loader = getDataLoader(ds_cfg=cfg['out_dataset'], dl_cfg=cfg['dataloader'], split="valid") else: out_train_loader = None out_valid_loader = None # Result directory and make tensorboard event file exp_dir = os.path.join(cfg['exp_root'], cfg['exp_dir']) if os.path.exists(exp_dir) is False: os.makedirs(exp_dir) shutil.copy('./config.py', os.path.join(exp_dir, "config.py")) writer_train = SummaryWriter(logdir=os.path.join(exp_dir, 'log', 'train')) writer_valid = SummaryWriter(logdir=os.path.join(exp_dir, 'log', 'valid')) # Stats Meters #train_meter = TrainMeter() #valid_meter = ValidMeter() # Loss function global_cfg['loss'] = cfg['loss'] # Outlier detector detector_func = detectors.getDetector(cfg['detector']) global_cfg['detector'] = cfg['detector'] print("=======================IMPORTANT CONFIG=======================") print(" Model : {}\n \ Loss : {}\n \ Detector : {}\n \ Optimizer: {}\n".format(cfg['model']['network_kind'], cfg['loss']['loss'], cfg['detector']['detector'], cfg['optim']['optimizer'])) print("============Start training. Result will be saved in {}".format(exp_dir)) for cur_epoch in range(start_epoch, cfg['max_epoch'] + 1): if out_train_loader is not None: train_summary = train_epoch_w_outlier(model, optimizer, in_train_loader, out_train_loader, loss_func, detector_func, cur_epoch, cfg['optim'], writer_train) else: train_summary = train_epoch_wo_outlier(model, optimizer, in_train_loader, attack_in, cur_epoch, cfg['optim'], writer_train) summary_write(summary=train_summary, writer=writer_train) print("Training result=========Epoch [{}]/[{}]=========\nlr: {} | loss: {} | acc: {}".format(cur_epoch, cfg['max_epoch'], train_summary['lr'], train_summary['avg_loss'], train_summary['classifier_acc'])) if cur_epoch % cfg['valid_epoch'] == 0: if out_valid_loader is not None: valid_summary = valid_epoch_w_outlier(model, in_valid_loader, out_valid_loader, loss_func, detector_func, cur_epoch) else: valid_summary = valid_epoch_wo_outlier(model, in_valid_loader, cur_epoch) summary_write(summary=valid_summary, writer=writer_valid) print("Validate result=========Epoch [{}]/[{}]=========\nloss: {} | acc: {}".format(cur_epoch, cfg['max_epoch'], valid_summary['avg_loss'], valid_summary['classifier_acc'])) if cur_epoch % cfg['ckpt_epoch'] == 0: ckpt_dir = os.path.join(cfg['exp_root'], cfg['exp_dir'], "ckpt") if os.path.exists(ckpt_dir) is False: os.makedirs(ckpt_dir) model_state = model.module.state_dict() if cfg['ngpu'] > 1 else model.state_dict() checkpoint = { "epoch": cur_epoch, "model_state": model_state, "optimizer_state": optimizer.state_dict(), } ckpt_name = "checkpoint_epoch_{}".format(cur_epoch) ckpt_path = os.path.join(ckpt_dir, ckpt_name + ".pyth") torch.save(checkpoint, ckpt_path)
def main(): # Reproducibility np.random.seed(cfg['seed']) torch.manual_seed(cfg['seed']) # Model & Optimizer model = getModel(cfg['model']) start_epoch = 1 max_epoch = 1 assert len(cfg['load_ckpt']) == len(cfg['in_dataset']['targets']) + 1 # Load model and optimizer for idx, ckpt in enumerate(cfg['load_ckpt']): checkpoint = torch.load(ckpt, map_location="cpu") if idx == 0: model.ava_network.load_state_dict(checkpoint['model_state']) else: model.ova_networks[idx - 1].load_state_dict( checkpoint['model_state']) print("load model on '{}' is complete.".format(ckpt)) cudnn.benchmark = True # Data Loader in_valid_loader = getDataLoader(ds_cfg=cfg['in_dataset'], dl_cfg=cfg['dataloader'], split="valid") if cfg['out_dataset'] is not None: out_valid_loader = getDataLoader(ds_cfg=cfg['out_dataset'], dl_cfg=cfg['dataloader'], split="valid") exp_dir = os.path.join(cfg['exp_root'], cfg['exp_dir'], "valid", cfg['out_dataset']['dataset']) else: out_train_loader = None out_valid_loader = None exp_dir = os.path.join(cfg['exp_root'], cfg['exp_dir'], "valid", "classifier") # Result directory and make tensorboard event file if os.path.exists(exp_dir) is False: os.makedirs(exp_dir) shutil.copy('./config.py', os.path.join(exp_dir, "val_config.py")) # Loss function loss_func = losses.getLoss(cfg['loss']) global_cfg['loss'] = cfg['loss'] # Outlier detector detector_func = detectors.getDetector(cfg['detector']) global_cfg['detector'] = cfg['detector'] print("=======================IMPORTANT CONFIG=======================") print(" Model : {}\n \ Loss : {}\n \ Detector : {}\n".format(cfg['model']['network_kind'], cfg['loss']['loss'], cfg['detector']['detector'])) print( "========Start validation. Result will be saved in {}".format(exp_dir)) logfile = open(os.path.join(exp_dir, "validation_log.txt"), "w") logfile2 = open(os.path.join(exp_dir, "wrong_predict_log.txt"), "w") for cur_epoch in range(start_epoch, max_epoch + 1): if out_valid_loader is not None: valid_summary = valid_epoch_w_outlier(model, in_valid_loader, out_valid_loader, loss_func, detector_func, cur_epoch, logfile2) summary_log = "=============Epoch [{}]/[{}]=============\nloss: {} | acc: {} | acc_w_ood: {}\nAUROC: {} | AUPR: {} | FPR95: {}\nInlier Conf. {} | Outlier Conf. {}\n".format( cur_epoch, max_epoch, valid_summary['avg_loss'], valid_summary['classifier_acc'], valid_summary['acc'], valid_summary['AUROC'], valid_summary['AUPR'], valid_summary['FPR95'], valid_summary['inlier_confidence'], valid_summary['outlier_confidence']) ind_max, ind_min = np.max(valid_summary['inliers']), np.min( valid_summary['inliers']) ood_max, ood_min = np.max(valid_summary['outliers']), np.min( valid_summary['outliers']) ranges = (ind_min if ind_min < ood_min else ood_min, ind_max if ind_max > ood_max else ood_max) fig = plt.figure() sns.distplot(valid_summary['inliers'].ravel(), hist_kws={'range': ranges}, kde=False, bins=50, norm_hist=True, label='In-distribution') sns.distplot(valid_summary['outliers'], hist_kws={'range': ranges}, kde=False, bins=50, norm_hist=True, label='Out-of-distribution') plt.xlabel('Confidence') plt.ylabel('Density') fig.legend() fig.savefig(os.path.join(exp_dir, "confidences.png")) else: valid_summary = valid_epoch_wo_outlier(model, in_valid_loader, loss_func, cur_epoch, logfile2) summary_log = "=============Epoch [{}]/[{}]=============\nloss: {} | acc: {}\n".format( cur_epoch, max_epoch, valid_summary['avg_loss'], valid_summary['classifier_acc']) print(summary_log) logfile.write(summary_log) logfile.close() logfile2.close()
def main(cfg): # set the path to pre-trained model and output pre_trained_net = cfg['load_ckpt'] outf=os.path.join(cfg['exp_root'],cfg['exp_dir'],'valid','MD_ADV') if os.path.isdir(outf) == False: os.makedirs(outf) torch.cuda.manual_seed(cfg['seed']) device = torch.device('cuda:0') # check the in-distribution dataset num_classes=cfg['model']['num_classes'] if cfg['detector']['adv_method'] == 'FGSM': adv_noise = 0.05 elif cfg['detector']['adv_method'] == 'BIM': adv_noise = 0.01 elif cfg['detector']['adv_method'] == 'DeepFool': if cfg['model']['net_type'] == 'resnet': if cfg['in_dataset']['dataset'] == 'cifar10': adv_noise = 0.18 elif cfg['in_dataset']['dataset'] == 'cifar100': adv_noise = 0.03 else: adv_noise = 0.1 else: if cfg['in_dataset']['dataset'] == 'cifar10': adv_noise = 0.6 elif cfg['in_dataset']['dataset'] == 'cifar100': adv_noise = 0.1 else: adv_noise = 0.5 # load networks # if args.net_type == 'densenet': # if args.dataset == 'svhn': model = getModel(cfg['model']) net_type='resnet' ## Should change in the future checkpoint = torch.load(cfg['load_ckpt'], map_location="cpu") model.load_state_dict(checkpoint['model_state']) in_transform = cfg['in_dataset']['valid_transform'] min_pixel = -2.42906570435 max_pixel = 2.75373125076 if cfg['in_dataset']['dataset'] == 'cifar10': if cfg['detector']['adv_method'] == 'FGSM': random_noise_size = 0.25 / 4 elif cfg['detector']['adv_method'] == 'BIM': random_noise_size = 0.13 / 2 elif cfg['detector']['adv_method'] == 'DeepFool': random_noise_size = 0.25 / 4 elif cfg['detector']['adv_method'] == 'CWL2': random_noise_size = 0.05 / 2 elif cfg['in_dataset']['dataset'] == 'cifar100': if cfg['detector']['adv_method'] == 'FGSM': random_noise_size = 0.25 / 8 elif cfg['detector']['adv_method'] == 'BIM': random_noise_size = 0.13 / 4 elif cfg['detector']['adv_method'] == 'DeepFool': random_noise_size = 0.13 / 4 elif cfg['detector']['adv_method'] == 'CWL2': random_noise_size = 0.05 / 2 else: if cfg['detector']['adv_method'] == 'FGSM': random_noise_size = 1 elif cfg['detector']['adv_method'] == 'BIM': random_noise_size = 1 elif cfg['detector']['adv_method'] == 'DeepFool': random_noise_size = 1 elif cfg['detector']['adv_method'] == 'CWL2': random_noise_size = 1 model.cuda() print("load model on '{}' is completed.".format(cfg['load_ckpt'])) # load dataset print('load target data: ', cfg['in_dataset']['dataset']) test_loader = getDataLoader(ds_cfg=cfg['in_dataset'], dl_cfg=cfg['dataloader'], split='valid') print('Attack: ' + cfg['detector']['adv_method'] + ', Dist: ' + cfg['detector']['adv_method'] + '\n') model.eval() adv_data_tot, clean_data_tot, noisy_data_tot = 0, 0, 0 label_tot = 0 correct, adv_correct, noise_correct = 0, 0, 0 total, generated_noise = 0, 0 criterion = nn.CrossEntropyLoss().cuda() selected_list = [] selected_index = 0 for data, target in test_loader: data, target = data.cuda(), target.cuda() with torch.no_grad(): data, target = Variable(data), Variable(target) output = model(data) # compute the accuracy pred = output.data.max(1)[1] equal_flag = pred.eq(target.data).cpu() correct += equal_flag.sum() noisy_data = torch.add(data.data, torch.randn(data.size()).cuda(), alpha = random_noise_size) noisy_data = torch.clamp(noisy_data, min_pixel, max_pixel) if total == 0: clean_data_tot = data.clone().data.cpu() label_tot = target.clone().data.cpu() noisy_data_tot = noisy_data.clone().cpu() else: clean_data_tot = torch.cat((clean_data_tot, data.clone().data.cpu()),0) label_tot = torch.cat((label_tot, target.clone().data.cpu()), 0) noisy_data_tot = torch.cat((noisy_data_tot, noisy_data.clone().cpu()),0) # generate adversarial model.zero_grad() inputs = Variable(data.data, requires_grad=True) output = model(inputs) loss = criterion(output, target) loss.backward() if cfg['detector']['adv_method'] == 'FGSM': gradient = torch.ge(inputs.grad.data, 0) gradient = (gradient.float()-0.5)*2 if cfg['model']['net_type'] == 'densenet': gradient.index_copy_(1, torch.LongTensor([0]).cuda(), \ gradient.index_select(1, torch.LongTensor([0]).cuda()) / (63.0/255.0)) gradient.index_copy_(1, torch.LongTensor([1]).cuda(), \ gradient.index_select(1, torch.LongTensor([1]).cuda()) / (62.1/255.0)) gradient.index_copy_(1, torch.LongTensor([2]).cuda(), \ gradient.index_select(1, torch.LongTensor([2]).cuda()) / (66.7/255.0)) else: gradient.index_copy_(1, torch.LongTensor([0]).cuda(), \ gradient.index_select(1, torch.LongTensor([0]).cuda()) / (0.2023)) gradient.index_copy_(1, torch.LongTensor([1]).cuda(), \ gradient.index_select(1, torch.LongTensor([1]).cuda()) / (0.1994)) gradient.index_copy_(1, torch.LongTensor([2]).cuda(), \ gradient.index_select(1, torch.LongTensor([2]).cuda()) / (0.2010)) elif cfg['detector']['adv_method'] == 'BIM': gradient = torch.sign(inputs.grad.data) for k in range(5): inputs = torch.add(inputs.data, gradient, alpha = adv_noise) inputs = torch.clamp(inputs, min_pixel, max_pixel) inputs = Variable(inputs, requires_grad=True) output = model(inputs) loss = criterion(output, target) loss.backward() gradient = torch.sign(inputs.grad.data) if cfg['model']['net_type'] == 'densenet': gradient.index_copy_(1, torch.LongTensor([0]).cuda(), \ gradient.index_select(1, torch.LongTensor([0]).cuda()) / (63.0/255.0)) gradient.index_copy_(1, torch.LongTensor([1]).cuda(), \ gradient.index_select(1, torch.LongTensor([1]).cuda()) / (62.1/255.0)) gradient.index_copy_(1, torch.LongTensor([2]).cuda(), \ gradient.index_select(1, torch.LongTensor([2]).cuda()) / (66.7/255.0)) else: gradient.index_copy_(1, torch.LongTensor([0]).cuda(), \ gradient.index_select(1, torch.LongTensor([0]).cuda()) / (0.2023)) gradient.index_copy_(1, torch.LongTensor([1]).cuda(), \ gradient.index_select(1, torch.LongTensor([1]).cuda()) / (0.1994)) gradient.index_copy_(1, torch.LongTensor([2]).cuda(), \ gradient.index_select(1, torch.LongTensor([2]).cuda()) / (0.2010)) if cfg['detector']['adv_method'] == 'DeepFool': _, adv_data = adversary.deepfool(model, data.data.clone(), target.data.cpu(), \ args.num_classes, step_size=adv_noise, train_mode=False) adv_data = adv_data.cuda() elif cfg['detector']['adv_method'] == 'CWL2': _, adv_data = adversary.cw(model, data.data.clone(), target.data.cpu(), 1.0, 'l2', crop_frac=1.0) else: adv_data = torch.add(inputs.data, gradient, alpha = adv_noise) adv_data = torch.clamp(adv_data, min_pixel, max_pixel) # measure the noise temp_noise_max = torch.abs((data.data - adv_data).view(adv_data.size(0), -1)) temp_noise_max, _ = torch.max(temp_noise_max, dim=1) generated_noise += torch.sum(temp_noise_max) if total == 0: flag = 1 adv_data_tot = adv_data.clone().cpu() else: adv_data_tot = torch.cat((adv_data_tot, adv_data.clone().cpu()),0) with torch.no_grad(): output = model(Variable(adv_data)) # compute the accuracy pred = output.data.max(1)[1] equal_flag_adv = pred.eq(target.data).cpu() adv_correct += equal_flag_adv.sum() with torch.no_grad(): output = model(Variable(noisy_data)) # compute the accuracy pred = output.data.max(1)[1] equal_flag_noise = pred.eq(target.data).cpu() noise_correct += equal_flag_noise.sum() for i in range(data.size(0)): if equal_flag[i] == 1 and equal_flag_noise[i] == 1 and equal_flag_adv[i] == 0: selected_list.append(selected_index) selected_index += 1 total += data.size(0) selected_list = torch.LongTensor(selected_list) clean_data_tot = torch.index_select(clean_data_tot, 0, selected_list) adv_data_tot = torch.index_select(adv_data_tot, 0, selected_list) noisy_data_tot = torch.index_select(noisy_data_tot, 0, selected_list) label_tot = torch.index_select(label_tot, 0, selected_list) torch.save(clean_data_tot, '%s/clean_data_%s_%s_%s.pth' % (args.outf, args.net_type, args.dataset, args.adv_type)) torch.save(adv_data_tot, '%s/adv_data_%s_%s_%s.pth' % (args.outf, args.net_type, args.dataset, args.adv_type)) torch.save(noisy_data_tot, '%s/noisy_data_%s_%s_%s.pth' % (args.outf, args.net_type, args.dataset, args.adv_type)) torch.save(label_tot, '%s/label_%s_%s_%s.pth' % (args.outf, args.net_type, args.dataset, args.adv_type)) print('Adversarial Noise:({:.2f})\n'.format(generated_noise / total)) print('Final Accuracy: {}/{} ({:.2f}%)\n'.format(correct, total, 100. * correct / total)) print('Adversarial Accuracy: {}/{} ({:.2f}%)\n'.format(adv_correct, total, 100. * adv_correct / total)) print('Noisy Accuracy: {}/{} ({:.2f}%)\n'.format(noise_correct, total, 100. * noise_correct / total))
writer_valid = SummaryWriter(logdir=os.path.join(outf, 'log', 'valid')) writer_test = SummaryWriter(logdir=os.path.join(outf, 'log', 'test')) writer_train = SummaryWriter(logdir=os.path.join(outf, 'log', 'train')) # check the in-distribution dataset num_classes = cfg['model']['num_classes'] best_auroc, best_result, best_index = 0, 0, 0 temp_file_name_1 = '%s/ADV_and_RANDOM_NOISE_SIZE_TRAIN.txt' % (outf) f = open(temp_file_name_1, 'w') temp_file_name_2 = '%s/ADV_and_RANDOM_NOISE_SIZE_VALID.txt' % (outf) g = open(temp_file_name_2, 'w') temp_file_name_3 = '%s/ADV_and_RANDOM_NOISE_SIZE_TEST.txt' % (outf) h = open(temp_file_name_3, 'w') model = getModel(cfg['model']) checkpoint = torch.load(cfg['load_ckpt'], map_location="cpu") model.load_state_dict(checkpoint['model_state']) min_pixel = cfg['min'] max_pixel = cfg['max'] model.cuda() print("load model on '{}' is completed.".format(cfg['load_ckpt'])) # load dataset test_loader = getDataLoader(ds_cfg=cfg['in_dataset'], dl_cfg=cfg['dataloader'], split='valid') fooled_log = open(os.path.join(outf, 'fooled.txt'), 'w')
def main(): global global_cfg # Reproducibility np.random.seed(cfg['seed']) torch.manual_seed(cfg['seed']) start_epoch = 1 # feature_extractor feature_extractor = getModel(cfg['model']['feature_extractor']) f_checkpoint = torch.load(cfg['f_load_ckpt'], map_location="cpu") feature_extractor.load_state_dict(f_checkpoint['model_state']) print("load feature_extractor on '{}' is complete.".format( cfg['f_load_ckpt'])) if cfg['in_dataset']['img_size'] == 32: G = Generator32(cfg['in_dataset']['batch_size'], cfg['in_dataset']['img_size'], cfg['model']['z_dim'], cfg['model']['g_conv_dim']).cuda() D = Discriminator32(cfg['in_dataset']['batch_size'], cfg['in_dataset']['img_size'], cfg['model']['d_conv_dim']).cuda() else: G = Generator(cfg['in_dataset']['batch_size'], cfg['in_dataset']['img_size'], cfg['model']['z_dim'], cfg['model']['g_conv_dim']).cuda() D = Discriminator(cfg['in_dataset']['batch_size'], cfg['in_dataset']['img_size'], cfg['model']['d_conv_dim']).cuda() G_optimizer = torch.optim.Adam( filter(lambda p: p.requires_grad, G.parameters()), cfg['optim']['g_lr'], [cfg['optim']['beta1'], cfg['optim']['beta2']]) D_optimizer = torch.optim.Adam( filter(lambda p: p.requires_grad, D.parameters()), cfg['optim']['d_lr'], [cfg['optim']['beta1'], cfg['optim']['beta2']]) if cfg['g_load_ckpt'] != '': g_checkpoint = torch.load(cfg['g_load_ckpt'], map_location="cpu") G.load_state_dict(g_checkpoint['model_state']) print("load model on '{}' is complete.".format(cfg['g_load_ckpt'])) if cfg['d_load_ckpt'] != '': d_checkpoint = torch.load(cfg['d_load_ckpt'], map_location="cpu") D.load_state_dict(d_checkpoint['model_state']) print("load model on '{}' is complete.".format(cfg['d_load_ckpt'])) cudnn.benchmark = True # Data Loader in_train_loader = getDataLoader(ds_cfg=cfg['in_dataset'], dl_cfg=cfg['dataloader'], split="train") in_valid_loader = getDataLoader(ds_cfg=cfg['in_dataset'], dl_cfg=cfg['dataloader'], split="valid") if cfg['out_dataset'] is not None: out_train_loader = getDataLoader(ds_cfg=cfg['out_dataset'], dl_cfg=cfg['dataloader'], split="train") out_valid_loader = getDataLoader(ds_cfg=cfg['out_dataset'], dl_cfg=cfg['dataloader'], split="valid") else: out_train_loader = None out_valid_loader = None # Result directory and make tensorboard event file exp_dir = os.path.join(cfg['exp_root'], cfg['exp_dir']) if os.path.exists(exp_dir) is False: os.makedirs(exp_dir) shutil.copy('./config.py', os.path.join(exp_dir, "config.py")) writer_train = SummaryWriter(logdir=os.path.join(exp_dir, 'log', 'train')) writer_valid = SummaryWriter(logdir=os.path.join(exp_dir, 'log', 'valid')) global_cfg['loss'] = cfg['loss'] # print("=======================IMPORTANT CONFIG=======================") # print(" Model : {}\n \ # Loss : {}\n \ # Detector : {}\n \ # Optimizer: {}\n".format(cfg['model']['network_kind'], cfg['loss']['loss'], cfg['detector']['detector'], cfg['optim']['optimizer'])) print("============Start training. Result will be saved in {}".format( exp_dir)) fixed_z = tensor2var(torch.randn(1, G.z_dim)) for cur_epoch in range(start_epoch, cfg['max_epoch'] + 1): if out_train_loader is not None: train_summary = train_epoch_w_outlier(model, optimizer, in_train_loader, out_train_loader, loss_func, detector_func, cur_epoch, cfg['optim'], writer_train) else: train_summary = train_epoch_wo_outlier(feature_extractor, G, D, G_optimizer, D_optimizer, in_train_loader, cur_epoch, cfg['optim'], writer_train) summary_write(summary=train_summary, writer=writer_train) print( "Training result=========Epoch [{}]/[{}]=========\nlr: {} | loss: {}" .format(cur_epoch, cfg['max_epoch'], train_summary['lr'], train_summary['avg_loss'])) # Sample image G.eval() fake_images, _, _ = G(fixed_z, train_summary['real_feature']) save_image(((fake_images.data + 1) / 2).clamp_(0, 1), os.path.join( cfg['exp_root'], cfg['exp_dir'], '{}_fake_{}.png'.format(cur_epoch, train_summary['real_target']))) # if cur_epoch % cfg['valid_epoch'] == 0: # if out_valid_loader is not None: # valid_summary = valid_epoch_w_outlier(model, in_valid_loader, out_valid_loader, loss_func, detector_func, cur_epoch) # else: # valid_summary = valid_epoch_wo_outlier(model, in_valid_loader, loss_func, cur_epoch) # summary_write(summary=valid_summary, writer=writer_valid) # print("Validate result=========Epoch [{}]/[{}]=========\nloss: {} | acc: {}".format(cur_epoch, cfg['max_epoch'], valid_summary['avg_loss'], valid_summary['classifier_acc'])) if cur_epoch % cfg['ckpt_epoch'] == 0: ckpt_dir = os.path.join(cfg['exp_root'], cfg['exp_dir'], "ckpt") if os.path.exists(ckpt_dir) is False: os.makedirs(ckpt_dir) G_state = G.module.state_dict( ) if cfg['ngpu'] > 1 else G.state_dict() G_checkpoint = { "epoch": cur_epoch, "model_state": G_state, "optimizer_state": G_optimizer.state_dict(), } ckpt_name = "G_checkpoint_epoch_{}".format(cur_epoch) ckpt_path = os.path.join(ckpt_dir, ckpt_name + ".pyth") torch.save(G_checkpoint, ckpt_path) D_state = D.module.state_dict( ) if cfg['ngpu'] > 1 else D.state_dict() D_checkpoint = { "epoch": cur_epoch, "model_state": D_state, "optimizer_state": D_optimizer.state_dict(), } ckpt_name = "D_checkpoint_epoch_{}".format(cur_epoch) ckpt_path = os.path.join(ckpt_dir, ckpt_name + ".pyth") torch.save(D_checkpoint, ckpt_path)