def main(): args = get_args() logger.info(args) np.random.seed(args.seed) torch.manual_seed(args.seed) torch.cuda.manual_seed(args.seed) dataset = cifar10(args.data_dir) test_set = list( zip( transpose(normalise(dataset['valid']['data'].astype(np.float32) / 255, mean=np.array(cifar10_mean, dtype=np.float32), std=np.array(cifar10_std, dtype=np.float32)), source='NHWC', target='NCHW'), dataset['valid']['targets'])) batches = Batches(test_set, args.batch_size, shuffle=False, num_workers=2) epsilon = (args.epsilon / 255.) / std alpha = (args.alpha / 255.) / std model = PreActResNet18().cuda() checkpoint = torch.load(args.fname) model.load_state_dict(checkpoint) model.eval() model.float() total_loss = 0 total_acc = 0 n = 0 if args.attack == 'none': with torch.no_grad(): for batch in batches: X, y = batch['input'], batch['target'] output = model(X) loss = F.cross_entropy(output, y) total_loss += loss.item() * y.size(0) total_acc += (output.max(1)[1] == y).sum().item() n += y.size(0) else: for batch in batches: X, y = batch['input'], batch['target'] if args.attack == 'pgd': delta = attack_pgd(model, X, y, epsilon, alpha, args.attack_iters, args.restarts) elif args.attack == 'fgsm': delta = attack_fgsm(model, X, y, epsilon) with torch.no_grad(): output = model(X + delta) loss = F.cross_entropy(output, y) total_loss += loss.item() * y.size(0) total_acc += (output.max(1)[1] == y).sum().item() n += y.size(0) logger.info('Test Loss: %.4f, Acc: %.4f', total_loss / n, total_acc / n)
def init_models(args): netG = Generator(args).cuda() netD = Discriminator(args).cuda() netL = PreActResNet().cuda() print (netG, netD) optimG = torch.optim.Adam(netG.parameters(), betas=(0.5, 0.999), lr=args.lr) optimD = torch.optim.Adam(netD.parameters(), betas=(0.5, 0.999), lr=args.lr) optimL = torch.optim.Adam(netL.parameters(), betas=(0.5, 0.999), lr=args.lr) if args.resume: Gpath = 'experiments/{}/models/netG_{}.pt'.format(args.name, args.load_step) Dpath = 'experiments/{}/models/netD_{}.pt'.format(args.name, args.load_step) Lpath = 'experiments/{}/models/netL_{}.pt'.format(args.name, args.load_step) netG, optimG = utils.load_model(args, netG, optimG, Gpath) netD, optimD = utils.load_model(args, netD, optimD, Dpath) netL, optimL = utils.load_model(args, netL, optimL, Lpath) return (netG, optimG), (netD, optimD), (netL, optimL)
def main(): args = get_args() logger.info(args) np.random.seed(args.seed) torch.manual_seed(args.seed) torch.cuda.manual_seed(args.seed) start_start_time = time.time() transforms = [Crop(32, 32), FlipLR()] dataset = cifar10(args.data_dir) train_set = list( zip(transpose(normalise(pad(dataset['train']['data'], 4))), dataset['train']['labels'])) train_set_x = Transform(train_set, transforms) train_batches = Batches(train_set_x, args.batch_size, shuffle=True, set_random_choices=True, num_workers=2) epsilon = (args.epsilon / 255.) / std pgd_alpha = (args.pgd_alpha / 255.) / std model = PreActResNet18().cuda() model.apply(initialize_weights) model.train() opt = torch.optim.SGD(model.parameters(), lr=args.lr_max, momentum=0.9, weight_decay=5e-4) model, opt = amp.initialize(model, opt, opt_level="O2", loss_scale=1.0, master_weights=False) criterion = nn.CrossEntropyLoss() if args.attack == 'free': delta = torch.zeros(args.batch_size, 3, 32, 32).cuda() delta.requires_grad = True elif args.attack == 'fgsm' and args.fgsm_init == 'previous': delta = torch.zeros(args.batch_size, 3, 32, 32).cuda() delta.requires_grad = True if args.attack == 'free': assert args.epochs % args.attack_iters == 0 epochs = int(math.ceil(args.epochs / args.attack_iters)) else: epochs = args.epochs if args.lr_schedule == 'cyclic': lr_schedule = lambda t: np.interp([t], [ 0, args.epochs * 2 // 5, args.epochs ], [0, args.lr_max, 0])[0] elif args.lr_schedule == 'piecewise': def lr_schedule(t): if t / args.epochs < 0.5: return args.lr_max elif t / args.epochs < 0.75: return args.lr_max / 10. else: return args.lr_max / 100. prev_robust_acc = 0. logger.info('Epoch \t Time \t LR \t \t Train Loss \t Train Acc') for epoch in range(epochs): start_time = time.time() train_loss = 0 train_acc = 0 train_n = 0 for i, batch in enumerate(train_batches): X, y = batch['input'], batch['target'] if i == 0: first_batch = batch lr = lr_schedule(epoch + (i + 1) / len(train_batches)) opt.param_groups[0].update(lr=lr) if args.attack == 'pgd': delta = attack_pgd(model, X, y, epsilon, pgd_alpha, args.attack_iters, args.restarts, opt) elif args.attack == 'fgsm': if args.fgsm_init == 'zero': delta = torch.zeros_like(X, requires_grad=True) elif args.fgsm_init == 'random': delta = torch.zeros_like(X).cuda() delta[:, 0, :, :].uniform_(-epsilon[0][0][0].item(), epsilon[0][0][0].item()) delta[:, 1, :, :].uniform_(-epsilon[1][0][0].item(), epsilon[1][0][0].item()) delta[:, 2, :, :].uniform_(-epsilon[2][0][0].item(), epsilon[2][0][0].item()) delta.requires_grad = True elif args.fgsm_init == 'previous': delta.requires_grad = True output = model(X + delta[:X.size(0)]) loss = F.cross_entropy(output, y) with amp.scale_loss(loss, opt) as scaled_loss: scaled_loss.backward() grad = delta.grad.detach() delta.data = clamp( delta + args.fgsm_alpha * epsilon * torch.sign(grad), -epsilon, epsilon) delta = delta.detach() elif args.attack == 'free': delta.requires_grad = True for j in range(args.attack_iters): epoch_iters = epoch * args.attack_iters + ( i * args.attack_iters + j + 1) / len(train_batches) lr = lr_schedule(epoch_iters) opt.param_groups[0].update(lr=lr) output = model( clamp(X + delta[:X.size(0)], lower_limit, upper_limit)) loss = F.cross_entropy(output, y) opt.zero_grad() with amp.scale_loss(loss, opt) as scaled_loss: scaled_loss.backward() grad = delta.grad.detach() delta.data = clamp(delta + epsilon * torch.sign(grad), -epsilon, epsilon) nn.utils.clip_grad_norm_(model.parameters(), 0.5) opt.step() delta.grad.zero_() elif args.attack == 'none': delta = torch.zeros_like(X) output = model( clamp(X + delta[:X.size(0)], lower_limit, upper_limit)) loss = criterion(output, y) if args.attack != 'free': opt.zero_grad() with amp.scale_loss(loss, opt) as scaled_loss: scaled_loss.backward() nn.utils.clip_grad_norm_(model.parameters(), 0.5) opt.step() train_loss += loss.item() * y.size(0) train_acc += (output.max(1)[1] == y).sum().item() train_n += y.size(0) if args.overfit_check: # Check current PGD robustness of model using random minibatch X, y = first_batch['input'], first_batch['target'] pgd_delta = attack_pgd(model, X, y, epsilon, pgd_alpha, args.attack_iters, args.restarts, opt) with torch.no_grad(): output = model( clamp(X + pgd_delta[:X.size(0)], lower_limit, upper_limit)) robust_acc = (output.max(1)[1] == y).sum().item() / y.size(0) if robust_acc - prev_robust_acc < -0.5: break prev_robust_acc = robust_acc best_state_dict = copy.deepcopy(model.state_dict()) train_time = time.time() logger.info('%d \t %.1f \t %.4f \t %.4f \t %.4f', epoch, train_time - start_time, lr, train_loss / train_n, train_acc / train_n) torch.save(best_state_dict, args.fname + '.pth') logger.info('Total time: %.4f', train_time - start_start_time)
def main(): np.random.seed(args.seed) torch.manual_seed(args.seed) torch.cuda.manual_seed(args.seed) train_loader, test_loader = get_loaders(args.data_dir, args.batch_size, 2, 32) epsilon = (args.epsilon / 255.) / std if (args.model_name == "PreActResNet18"): model = PreActResNet18().to(device) else: model = models.resnet50().to(device) opt = torch.optim.SGD(model.parameters(), lr=args.lr_max, momentum=args.momentum, weight_decay=args.weight_decay) amp_args = dict(opt_level=args.opt_level, loss_scale=args.loss_scale, verbosity=False) if args.opt_level == 'O2': amp_args['master_weights'] = args.master_weights model, opt = amp.initialize(model, opt, **amp_args) criterion = nn.CrossEntropyLoss() delta = torch.zeros(args.batch_size, 3, 32, 32).cuda() delta.requires_grad = True lr_steps = args.epochs * len(train_loader) * args.minibatch_replays if args.lr_schedule == 'cyclic': scheduler = torch.optim.lr_scheduler.CyclicLR( opt, base_lr=args.lr_min, max_lr=args.lr_max, step_size_up=lr_steps / 2, step_size_down=lr_steps / 2) elif args.lr_schedule == 'multistep': scheduler = torch.optim.lr_scheduler.MultiStepLR( opt, milestones=[lr_steps / 2, lr_steps * 3 / 4], gamma=0.1) # Training total_time = 0 for epoch in range(args.epochs): start_train = time.time() train(train_loader, model, args.minibatch_replays, criterion, epoch, epsilon, delta, opt, scheduler) end_train = time.time() epoch_time = (end_train - start_train) / 60 total_time += epoch_time print("Epoch time: %.4f minutes", epoch_time) # Evaluation best_state_dict = model.state_dict() if (args.model_name == "PreActResNet18"): model_test = PreActResNet18().to(device) else: model_test = models.resnet50().to(device) model_test.load_state_dict(best_state_dict) model_test.float() model_test.eval() # Evaluate standard acc on test set test_loss, test_acc, test_err = evaluate_standard( test_loader, model_test) print("Test acc, err, loss: %.3f, %.3f, %.3f" % (test_acc, test_err, test_loss)) ''' # Evaluate acc against PGD_10 attack pgd_loss, pgd_acc, pgd_err = evaluate_pgd(test_loader, model_test, 10, 1) print("PGD_10 acc, err, loss: %.3f, %.3f, %.3f" %(pgd_acc, pgd_err, pgd_loss)) # Evaluate acc against PGD_20 attack pgd_loss, pgd_acc, pgd_err = evaluate_pgd(test_loader, model_test, 20, 1) print("PGD_20 acc, err, loss: %.3f, %.3f, %.3f" %(pgd_acc, pgd_err, pgd_loss)) ''' # Evaluate acc against PGD_50 attack pgd_loss, pgd_acc, pgd_err = evaluate_pgd(test_loader, model_test, 50, 1) print("PGD_50 acc, err, loss: %.3f, %.3f, %.3f" % (pgd_acc, pgd_err, pgd_loss)) logger.info('Total train time: %.4f minutes', total_time)
def main(): args = get_args() if not os.path.exists(args.out_dir): os.mkdir(args.out_dir) logfile = os.path.join(args.out_dir, 'output.log') if os.path.exists(logfile): os.remove(logfile) logging.basicConfig( format='[%(asctime)s] - %(message)s', datefmt='%Y/%m/%d %H:%M:%S', level=logging.INFO, filename=logfile) logger.info(args) np.random.seed(args.seed) torch.manual_seed(args.seed) torch.cuda.manual_seed(args.seed) train_loader, test_loader = get_loaders(args.data_dir, args.batch_size) epsilon = (args.epsilon / 255.) / std model = PreActResNet18().cuda() model.train() opt = torch.optim.SGD(model.parameters(), lr=args.lr_max, momentum=args.momentum, weight_decay=args.weight_decay) amp_args = dict(opt_level=args.opt_level, loss_scale=args.loss_scale, verbosity=False) if args.opt_level == 'O2': amp_args['master_weights'] = args.master_weights model, opt = amp.initialize(model, opt, **amp_args) criterion = nn.CrossEntropyLoss() delta = torch.zeros(args.batch_size, 3, 32, 32).cuda() delta.requires_grad = True lr_steps = args.epochs * len(train_loader) * args.minibatch_replays if args.lr_schedule == 'cyclic': scheduler = torch.optim.lr_scheduler.CyclicLR(opt, base_lr=args.lr_min, max_lr=args.lr_max, step_size_up=lr_steps / 2, step_size_down=lr_steps / 2) elif args.lr_schedule == 'multistep': scheduler = torch.optim.lr_scheduler.MultiStepLR(opt, milestones=[lr_steps / 2, lr_steps * 3 / 4], gamma=0.1) # Training start_train_time = time.time() logger.info('Epoch \t Seconds \t LR \t \t Train Loss \t Train Acc') for epoch in range(args.epochs): start_epoch_time = time.time() train_loss = 0 train_acc = 0 train_n = 0 for i, (X, y) in enumerate(train_loader): X, y = X.cuda(), y.cuda() for _ in range(args.minibatch_replays): output = model(X + delta[:X.size(0)]) loss = criterion(output, y) opt.zero_grad() with amp.scale_loss(loss, opt) as scaled_loss: scaled_loss.backward() grad = delta.grad.detach() delta.data = clamp(delta + epsilon * torch.sign(grad), -epsilon, epsilon) delta.data[:X.size(0)] = clamp(delta[:X.size(0)], lower_limit - X, upper_limit - X) opt.step() delta.grad.zero_() scheduler.step() train_loss += loss.item() * y.size(0) train_acc += (output.max(1)[1] == y).sum().item() train_n += y.size(0) epoch_time = time.time() lr = scheduler.get_lr()[0] logger.info('%d \t %.1f \t \t %.4f \t %.4f \t %.4f', epoch, epoch_time - start_epoch_time, lr, train_loss/train_n, train_acc/train_n) train_time = time.time() torch.save(model.state_dict(), os.path.join(args.out_dir, 'model.pth')) logger.info('Total train time: %.4f minutes', (train_time - start_train_time)/60) # Evaluation model_test = PreActResNet18().cuda() model_test.load_state_dict(model.state_dict()) model_test.float() model_test.eval() pgd_loss, pgd_acc = evaluate_pgd(test_loader, model_test, 50, 10) test_loss, test_acc = evaluate_standard(test_loader, model_test) logger.info('Test Loss \t Test Acc \t PGD Loss \t PGD Acc') logger.info('%.4f \t \t %.4f \t %.4f \t %.4f', test_loss, test_acc, pgd_loss, pgd_acc)
def main(): # Data Loader (Input Pipeline) print('loading dataset...') train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, num_workers=args.num_workers, drop_last=True, shuffle=True) test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size, num_workers=args.num_workers, drop_last=True, shuffle=False) # Define models print('building model...') if args.dataset == 'mnist': clf1 = MLPNet() if args.dataset == 'cifar10': clf1 = CNN_small(num_classes) if args.dataset == 'cifar100': clf1 = CNN(n_outputs=num_classes) if args.dataset == 'news': clf1 = NewsNet(weights_matrix=train_dataset.weights_matrix, num_classes=num_classes) if args.dataset == 'imagenet_tiny': clf1 = PreActResNet18(num_classes=200) clf1.cuda() print(clf1.parameters) optimizer1 = torch.optim.Adam(clf1.parameters(), lr=learning_rate) if args.dataset == 'mnist': clf2 = MLPNet() if args.dataset == 'cifar10': clf2 = CNN_small(num_classes) if args.dataset == 'cifar100': clf2 = CNN(n_outputs=num_classes) if args.dataset == 'news': clf2 = NewsNet(weights_matrix=train_dataset.weights_matrix, num_classes=num_classes) if args.dataset == 'imagenet_tiny': clf2 = PreActResNet18(num_classes=200) clf2.cuda() print(clf2.parameters) optimizer2 = torch.optim.Adam(clf2.parameters(), lr=learning_rate) with open(txtfile, "a") as myfile: myfile.write('epoch train_acc1 train_acc2 test_acc1 test_acc2\n') epoch = 0 train_acc1 = 0 train_acc2 = 0 # evaluate models with random weights test_acc1, test_acc2 = evaluate(test_loader, clf1, clf2) print( 'Epoch [%d/%d] Test Accuracy on the %s test data: Model1 %.4f %% Model2 %.4f %%' % (epoch + 1, args.n_epoch, len(test_dataset), test_acc1, test_acc2)) # save results with open(txtfile, "a") as myfile: myfile.write( str(int(epoch)) + ' ' + str(train_acc1) + ' ' + str(train_acc2) + ' ' + str(test_acc1) + " " + str(test_acc2) + "\n") # training for epoch in range(1, args.n_epoch): # train models clf1.train() clf2.train() adjust_learning_rate(optimizer1, epoch) adjust_learning_rate(optimizer2, epoch) train_acc1, train_acc2 = train(train_loader, epoch, clf1, optimizer1, clf2, optimizer2) # evaluate models test_acc1, test_acc2 = evaluate(test_loader, clf1, clf2) # save results print( 'Epoch [%d/%d] Test Accuracy on the %s test data: Model1 %.4f %% Model2 %.4f %%' % (epoch + 1, args.n_epoch, len(test_dataset), test_acc1, test_acc2)) with open(txtfile, "a") as myfile: myfile.write( str(int(epoch)) + ' ' + str(train_acc1) + ' ' + str(train_acc2) + ' ' + str(test_acc1) + " " + str(test_acc2) + "\n")
def main(): args = get_args() if not os.path.exists(args.out_dir): os.mkdir(args.out_dir) logfile = os.path.join(args.out_dir, 'output.log') if os.path.exists(logfile): os.remove(logfile) logging.basicConfig(format='[%(asctime)s] - %(message)s', datefmt='%Y/%m/%d %H:%M:%S', level=logging.INFO, filename=os.path.join(args.out_dir, 'output.log')) logger.info(args) np.random.seed(args.seed) torch.manual_seed(args.seed) torch.cuda.manual_seed(args.seed) train_loader, test_loader = get_loaders(args.data_dir, args.batch_size) epsilon = (args.epsilon / 255.) / std alpha = (args.alpha / 255.) / std pgd_alpha = (2 / 255.) / std model = PreActResNet18().cuda() model.train() opt = torch.optim.SGD(model.parameters(), lr=args.lr_max, momentum=args.momentum, weight_decay=args.weight_decay) amp_args = dict(opt_level=args.opt_level, loss_scale=args.loss_scale, verbosity=False) if args.opt_level == 'O2': amp_args['master_weights'] = args.master_weights model, opt = amp.initialize(model, opt, **amp_args) criterion = nn.CrossEntropyLoss() if args.delta_init == 'previous': delta = torch.zeros(args.batch_size, 3, 32, 32).cuda() lr_steps = args.epochs * len(train_loader) if args.lr_schedule == 'cyclic': scheduler = torch.optim.lr_scheduler.CyclicLR( opt, base_lr=args.lr_min, max_lr=args.lr_max, step_size_up=lr_steps / 2, step_size_down=lr_steps / 2) elif args.lr_schedule == 'multistep': scheduler = torch.optim.lr_scheduler.MultiStepLR( opt, milestones=[lr_steps / 2, lr_steps * 3 / 4], gamma=0.1) # Training prev_robust_acc = 0. start_train_time = time.time() logger.info('Epoch \t Seconds \t LR \t \t Train Loss \t Train Acc') for epoch in range(args.epochs): start_epoch_time = time.time() train_loss = 0 train_acc = 0 train_n = 0 for i, (X, y) in enumerate(train_loader): X, y = X.cuda(), y.cuda() if i == 0: first_batch = (X, y) if args.delta_init != 'previous': delta = torch.zeros_like(X).cuda() if args.delta_init == 'random': for i in range(len(epsilon)): delta[:, i, :, :].uniform_(-epsilon[i][0][0].item(), epsilon[0][0][0].item()) delta.data = clamp(delta, lower_limit - X, upper_limit - X) delta.requires_grad = True output = model(X + delta[:X.size(0)]) loss = F.cross_entropy(output, y) with amp.scale_loss(loss, opt) as scaled_loss: scaled_loss.backward() grad = delta.grad.detach() delta.data = clamp(delta + alpha * torch.sign(grad), -epsilon, epsilon) delta.data[:X.size(0)] = clamp(delta[:X.size(0)], lower_limit - X, upper_limit - X) delta = delta.detach() output = model(X + delta[:X.size(0)]) loss = criterion(output, y) opt.zero_grad() with amp.scale_loss(loss, opt) as scaled_loss: scaled_loss.backward() opt.step() train_loss += loss.item() * y.size(0) train_acc += (output.max(1)[1] == y).sum().item() train_n += y.size(0) scheduler.step() if args.early_stop: # Check current PGD robustness of model using random minibatch X, y = first_batch pgd_delta = attack_pgd(model, X, y, epsilon, pgd_alpha, 5, 1, opt) with torch.no_grad(): output = model( clamp(X + pgd_delta[:X.size(0)], lower_limit, upper_limit)) robust_acc = (output.max(1)[1] == y).sum().item() / y.size(0) if robust_acc - prev_robust_acc < -0.2: break prev_robust_acc = robust_acc best_state_dict = copy.deepcopy(model.state_dict()) epoch_time = time.time() lr = scheduler.get_lr()[0] logger.info('%d \t %.1f \t \t %.4f \t %.4f \t %.4f', epoch, epoch_time - start_epoch_time, lr, train_loss / train_n, train_acc / train_n) train_time = time.time() if not args.early_stop: best_state_dict = model.state_dict() torch.save(best_state_dict, os.path.join(args.out_dir, 'model.pth')) logger.info('Total train time: %.4f minutes', (train_time - start_train_time) / 60) # Evaluation model_test = PreActResNet18().cuda() model_test.load_state_dict(best_state_dict) model_test.float() model_test.eval() pgd_loss, pgd_acc = evaluate_pgd(test_loader, model_test, 50, 10) test_loss, test_acc = evaluate_standard(test_loader, model_test) logger.info('Test Loss \t Test Acc \t PGD Loss \t PGD Acc') logger.info('%.4f \t \t %.4f \t %.4f \t %.4f', test_loss, test_acc, pgd_loss, pgd_acc)
def main(): args = get_args() logger.info(args) np.random.seed(args.seed) torch.manual_seed(args.seed) torch.cuda.manual_seed(args.seed) dataset = datasets.CIFAR10(root=args.data_dir, train=False, transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize( cifar10_mean, cifar10_std) ])) dataloader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size, num_workers=2, pin_memory=True, shuffle=False, drop_last=False) epsilon = (args.epsilon / 255.) / std alpha = (args.alpha / 255.) / std model = PreActResNet18().cuda() intvl = [int(x) for x in args.intvl.split(',')] attack = set(args.attack.split(',')) stats = {'epoch': [], 'none': [], 'pgd': []} logger.info('Epoch \t Standard \t PGD') for epoch in range(intvl[0], intvl[1], intvl[2]): epoch_name = args.fname + '_{}.pth'.format(epoch) checkpoint = torch.load(epoch_name) model.load_state_dict(checkpoint) model.eval() model.float() msg = "{} \t ".format(epoch) if 'none' in attack: # test none attack none_total_loss = 0 none_total_acc = 0 none_n = 0 with torch.no_grad(): for X, y in dataloader: X, y = X.to(device), y.to(device) output = model(X) loss = F.cross_entropy(output, y) none_total_loss += loss.item() * y.size(0) none_total_acc += (output.max(1)[1] == y).sum().item() none_n += y.size(0) msg += "{:.4f} \t ".format(none_total_acc / none_n) stats['none'].append(none_total_acc / none_n) else: msg += "- \t " if 'pgd' in attack: # test pgd pgd_total_loss = 0 pgd_total_acc = 0 pgd_n = 0.1 for X, y in dataloader: X, y = X.to(device), y.to(device) delta = attack_pgd(model, X, y, epsilon, alpha, args.attack_iters, args.restarts) with torch.no_grad(): output = model(X + delta) loss = F.cross_entropy(output, y) pgd_total_loss += loss.item() * y.size(0) pgd_total_acc += (output.max(1)[1] == y).sum().item() pgd_n += y.size(0) msg += "{:.4f}".format(pgd_total_acc / pgd_n) stats['pgd'].append(pgd_total_acc / pgd_n) else: msg += "-" stats['epoch'].append(epoch) logger.info(msg) logger.info(stats)
def run(args: DictConfig) -> None: cuda_available = torch.cuda.is_available() torch.manual_seed(args.seed) device = "cuda" if cuda_available and args.device == 'cuda' else "cpu" # n_classes = args.n_classes # classifier = get_model(name=args.classifier_name, n_classes=n_classes).to(device) classifier = PreActResNet18().to(device) logger.info('Classifier: {}, # parameters: {}'.format( args.classifier_name, cal_parameters(classifier))) data_dir = hydra.utils.to_absolute_path(args.data_dir) train_data = get_dataset(data_name=args.dataset, data_dir=data_dir, train=True, crop_flip=True) test_data = get_dataset(data_name=args.dataset, data_dir=data_dir, train=False, crop_flip=False) train_loader = DataLoader(dataset=train_data, batch_size=args.n_batch_train, shuffle=True) test_loader = DataLoader(dataset=test_data, batch_size=args.n_batch_test, shuffle=False) if args.inference is True: classifier.load_state_dict( torch.load('{}_at.pth'.format(args.classifier_name))) logger.info('Load classifier from checkpoint') else: # optimizer = SGD(classifier.parameters(), lr=args.lr_max, momentum=args.momentum, weight_decay=args.weight_decay) # lr_steps = args.n_epochs * len(train_loader) # scheduler = lr_scheduler.CyclicLR(optimizer, base_lr=args.lr_min, max_lr=args.lr_max, # step_size_up=lr_steps/2, step_size_down=lr_steps/2) optimizer = torch.optim.SGD(classifier.parameters(), args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=True) scheduler = torch.optim.lr_scheduler.LambdaLR( optimizer, lr_lambda=lambda step: get_lr( # pylint: disable=g-long-lambda step, args.n_epochs * len(train_loader), 1, # lr_lambda computes multiplicative factor 1e-6 / args.learning_rate)) optimal_loss = 1e5 for epoch in range(1, args.n_epochs + 1): loss, acc = train_epoch(classifier, train_loader, args, optimizer, scheduler=scheduler) lr = scheduler.get_lr()[0] logger.info('Epoch {}, lr:{:.4f}, loss:{:.4f}, Acc:{:.4f}'.format( epoch, lr, loss, acc)) if loss < optimal_loss: optimal_loss = loss torch.save(classifier.state_dict(), '{}_at.pth'.format(args.classifier_name)) clean_loss, clean_acc = eval_epoch(classifier, test_loader, args, adversarial=False) adv_loss, adv_acc = eval_epoch(classifier, test_loader, args, adversarial=True) logger.info('Clean loss: {:.4f}, acc: {:.4f}'.format( clean_loss, clean_acc)) logger.info('Adversarial loss: {:.4f}, acc: {:.4f}'.format( adv_loss, adv_acc))
def main(): # Scale and initialize the parameters best_prec1 = 0 configs.TRAIN.epochs = int( math.ceil(configs.TRAIN.epochs / configs.ADV.n_repeats)) configs.ADV.fgsm_step /= configs.DATA.max_color_value configs.ADV.clip_eps /= configs.DATA.max_color_value # Create output folder if not os.path.isdir(os.path.join('trained_models', configs.output_name)): os.makedirs(os.path.join('trained_models', configs.output_name)) # Log the config details logger.info(pad_str(' ARGUMENTS ')) for k, v in configs.items(): print('{}: {}'.format(k, v)) logger.info(pad_str('')) # Create the model if configs.pretrained: print("=> using pre-trained model '{}'".format(configs.TRAIN.arch)) model = models.__dict__[configs.TRAIN.arch](pretrained=True) else: print("=> creating model '{}'".format(configs.TRAIN.arch)) model = models.__dict__[configs.TRAIN.arch]() # Use GPU or CPU model = model.to(device) # Criterion: criterion = nn.CrossEntropyLoss().cuda() # Optimizer: optimizer = torch.optim.SGD(model.parameters(), configs.TRAIN.lr, momentum=configs.TRAIN.momentum, weight_decay=configs.TRAIN.weight_decay) # Resume if a valid checkpoint path is provided if configs.resume: if os.path.isfile(configs.resume): print("=> loading checkpoint '{}'".format(configs.resume)) checkpoint = torch.load(configs.resume) configs.TRAIN.start_epoch = checkpoint['epoch'] best_prec1 = checkpoint['best_prec1'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) print("=> loaded checkpoint '{}' (epoch {})".format( configs.resume, checkpoint['epoch'])) else: print("=> no checkpoint found at '{}'".format(configs.resume)) # Data print('==> Preparing data..') trainloader, testloader = get_loaders("./data", configs.DATA.batch_size, configs.DATA.workers, configs.DATA.crop_size) # If in evaluate mode: perform validation on PGD attacks as well as clean samples model = PreActResNet18().cuda() # Use GPU or CPU model = model.to(device) # if(configs.load_weights): # logger.info(pad_str("LOADING WEIGHTS")) # model_path = "cifar_model_weights_30_epochs.pth" # state_dict = torch.load(model_path) # model.load_state_dict(state_dict) # model = model.eval() if configs.evaluate: logger.info(pad_str(' Performing PGD Attacks ')) for pgd_param in configs.ADV.pgd_attack: validate_pgd(testloader, model, criterion, pgd_param[0], pgd_param[1], configs, logger) validate(testloader, model, criterion, configs, logger) return start_train_time = time.time() for epoch in range(configs.TRAIN.start_epoch, configs.TRAIN.epochs): adjust_learning_rate(configs.TRAIN.lr, optimizer, epoch, configs.ADV.n_repeats) # train for one epoch train(trainloader, model, criterion, optimizer, epoch) # evaluate on validation set prec1 = validate(testloader, model, criterion, configs, logger) # remember best prec@1 and save checkpoint is_best = prec1 > best_prec1 best_prec1 = max(prec1, best_prec1) save_checkpoint( { 'epoch': epoch + 1, 'arch': configs.TRAIN.arch, 'state_dict': model.state_dict(), 'best_prec1': best_prec1, 'optimizer': optimizer.state_dict(), }, is_best, os.path.join('trained_models', configs.output_name)) train_time = time.time() print('Total train time: %.4f minutes', (train_time - start_train_time) / 60) # Automatically perform PGD Attacks at the end of training logger.info(pad_str(' Performing PGD Attacks ')) ''' for pgd_param in configs.ADV.pgd_attack: validate_pgd(testloader, model, criterion, pgd_param[0], pgd_param[1], configs, logger) ''' train_time = time.time() best_state_dict = model.state_dict() torch.save(best_state_dict, os.path.join(args.out_dir, 'model.pth')) logger.info('Total train time: %.4f minutes', (train_time - start_train_time) / 60) # Evaluation model_test = PreActResNet18().cuda() model_test.load_state_dict(best_state_dict) model_test.float() model_test.eval() test_loss, test_acc = evaluate_standard(testloader, model_test) logger.info('Test Loss \t Test Acc') logger.info('%.4f \t \t %.4f', test_loss, test_acc) for pgd_param in configs.ADV.pgd_attack: logger.info(pad_str("PGD-" + str(pgd_param[0]))) pgd_loss, pgd_acc = evaluate_pgd(testloader, model_test, pgd_param[0], 10) logger.info('PGD Loss \t PGD Acc') logger.info('%.4f \t \t %.4f', pgd_loss, pgd_acc)