def main(): if args.tensorboard: configure("runs/%s"%(args.name)) if args.augment: transform_train = transforms.Compose([ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), ]) else: transform_train = transforms.Compose([ transforms.ToTensor(), ]) transform_test = transforms.Compose([ transforms.ToTensor(), ]) kwargs = {'num_workers': 1, 'pin_memory': True} if args.in_dataset == "CIFAR-10": # Data loading code normalizer = transforms.Normalize(mean=[x/255.0 for x in [125.3, 123.0, 113.9]], std=[x/255.0 for x in [63.0, 62.1, 66.7]]) train_loader = torch.utils.data.DataLoader( torchvision.datasets.ImageFolder('./datasets/row_train_data/CIFAR-10', transform=transform_train), batch_size=args.batch_size, shuffle=True, **kwargs) val_loader = torch.utils.data.DataLoader( datasets.CIFAR10('./datasets/cifar10', train=False, transform=transform_test), batch_size=args.batch_size, shuffle=True, **kwargs) num_classes = 10 lr_schedule=[50, 75, 90] elif args.in_dataset == "CIFAR-100": # Data loading code normalizer = transforms.Normalize(mean=[x/255.0 for x in [125.3, 123.0, 113.9]], std=[x/255.0 for x in [63.0, 62.1, 66.7]]) train_loader = torch.utils.data.DataLoader( torchvision.datasets.ImageFolder('./datasets/row_train_data/CIFAR-100', transform=transform_train), batch_size=args.batch_size, shuffle=True, **kwargs) val_loader = torch.utils.data.DataLoader( datasets.CIFAR100('./datasets/cifar100', train=False, transform=transform_test), batch_size=args.batch_size, shuffle=True, **kwargs) num_classes = 100 lr_schedule=[50, 75, 90] elif args.in_dataset == "SVHN": # Data loading code normalizer = None transform = transforms.Compose([transforms.ToTensor(),]) train_loader = torch.utils.data.DataLoader( torchvision.datasets.ImageFolder('./datasets/row_train_data/SVHN', transform=transform), batch_size=args.batch_size, shuffle=True, **kwargs) val_loader = torch.utils.data.DataLoader( svhn.SVHN('datasets/svhn/', split='test', transform=transforms.ToTensor(), download=False), batch_size=args.batch_size, shuffle=False, **kwargs) args.epochs = 20 args.save_epoch = 2 lr_schedule=[10, 15, 18] num_classes = 10 # create model if args.model_arch == 'densenet': model = dn.DenseNet3(args.layers, num_classes + 1, args.growth, reduction=args.reduce, bottleneck=args.bottleneck, dropRate=args.droprate, normalizer=normalizer) elif args.model_arch == 'wideresnet': model = wn.WideResNet(args.depth, num_classes + 1, widen_factor=args.width, dropRate=args.droprate, normalizer=normalizer) else: assert False, 'Not supported model arch: {}'.format(args.model_arch) attack = LinfPGDAttack(model = model, eps=args.epsilon, nb_iter=args.iters, eps_iter=args.iter_size, rand_init=True, targeted=True, num_classes=num_classes+1, loss_func='CE', elementwise_best=True) # get the number of model parameters print('Number of model parameters: {}'.format( sum([p.data.nelement() for p in model.parameters()]))) model = model.cuda() cudnn.benchmark = True # define loss function (criterion) and pptimizer criterion = nn.CrossEntropyLoss().cuda() optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=args.momentum, nesterov=True, weight_decay=args.weight_decay) # optionally resume from a checkpoint if args.resume: if os.path.isfile(args.resume): print("=> loading checkpoint '{}'".format(args.resume)) checkpoint = torch.load(args.resume) args.start_epoch = checkpoint['epoch'] model.load_state_dict(checkpoint['state_dict']) print("=> loaded checkpoint '{}' (epoch {})" .format(args.resume, checkpoint['epoch'])) else: print("=> no checkpoint found at '{}'".format(args.resume)) for epoch in range(args.start_epoch, args.epochs): adjust_learning_rate(optimizer, epoch, lr_schedule) # train for one epoch train_rowl(train_loader, model, criterion, optimizer, epoch, num_classes, attack) # evaluate on validation set prec1 = validate(val_loader, model, criterion, num_classes, epoch) # remember best prec@1 and save checkpoint if (epoch + 1) % args.save_epoch == 0: save_checkpoint({ 'epoch': epoch + 1, 'state_dict': model.state_dict(), }, epoch + 1)
def main(): if args.tensorboard: configure("runs/%s" % (args.name)) # Data loading code normalizer = transforms.Normalize( mean=[x / 255.0 for x in [125.3, 123.0, 113.9]], std=[x / 255.0 for x in [63.0, 62.1, 66.7]]) if args.augment: transform_train = transforms.Compose([ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), ]) else: transform_train = transforms.Compose([ transforms.ToTensor(), ]) transform_test = transforms.Compose([ transforms.ToTensor(), ]) kwargs = {'num_workers': 1, 'pin_memory': True} if args.in_dataset == "CIFAR-10": train_loader = torch.utils.data.DataLoader( datasets.CIFAR10('../../data', train=True, download=True, transform=transform_train), batch_size=args.batch_size, shuffle=True, **kwargs) val_loader = torch.utils.data.DataLoader(datasets.CIFAR10( '../../data', train=False, transform=transform_test), batch_size=args.batch_size, shuffle=True, **kwargs) num_classes = 10 elif args.in_dataset == "CIFAR-100": train_loader = torch.utils.data.DataLoader(datasets.CIFAR100( './datasets/cifar100', train=True, download=True, transform=transform_train), batch_size=args.batch_size, shuffle=True, **kwargs) val_loader = torch.utils.data.DataLoader(datasets.CIFAR100( './datasets/cifar100', train=False, transform=transform_test), batch_size=args.batch_size, shuffle=True, **kwargs) num_classes = 100 # create model model = dn.DenseNet3(args.layers, num_classes, args.growth, reduction=args.reduce, bottleneck=args.bottleneck, dropRate=args.droprate, normalizer=normalizer) if args.adv: attack_in = LinfPGDAttack(model=model, eps=args.epsilon, nb_iter=args.iters, eps_iter=args.iter_size, rand_init=True, loss_func='CE') if args.ood: attack_out = LinfPGDAttack(model=model, eps=args.epsilon, nb_iter=args.iters, eps_iter=args.iter_size, rand_init=True, loss_func='OE') # get the number of model parameters print('Number of model parameters: {}'.format( sum([p.data.nelement() for p in model.parameters()]))) # for training on multiple GPUs. # Use CUDA_VISIBLE_DEVICES=0,1 to specify which GPUs to use # model = torch.nn.DataParallel(model).cuda() model = model.cuda() # optionally resume from a checkpoint if args.resume: if os.path.isfile(args.resume): print("=> loading checkpoint '{}'".format(args.resume)) checkpoint = torch.load(args.resume) args.start_epoch = checkpoint['epoch'] model.load_state_dict(checkpoint['state_dict']) print("=> loaded checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) else: print("=> no checkpoint found at '{}'".format(args.resume)) cudnn.benchmark = True # define loss function (criterion) and pptimizer ! 为NL标签已改为BCELoss criterion = nn.BCELoss().cuda() if args.ood: ood_criterion = OELoss().cuda() optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=args.momentum, nesterov=True, weight_decay=args.weight_decay) if args.lr_scheduler != 'cosine_annealing' and args.lr_scheduler != 'step_decay': assert False, 'Not supported lr_scheduler {}'.format(args.lr_scheduler) if args.lr_scheduler == 'cosine_annealing': scheduler = torch.optim.lr_scheduler.LambdaLR( optimizer, lr_lambda=lambda step: cosine_annealing( step, args.epochs * len(train_loader), 1, # since lr_lambda computes multiplicative factor 1e-6 / args.lr)) else: scheduler = None for epoch in range(args.start_epoch, args.epochs): if args.lr_scheduler == 'step_decay': adjust_learning_rate(optimizer, epoch) # train for one epoch if args.ood: pass else: if args.adv: pass else: train(train_loader, model, criterion, optimizer, scheduler, epoch) # evaluate on validation set prec1 = validate(val_loader, model, criterion, epoch) # remember best prec@1 and save checkpoint if (epoch + 1) % args.save_epoch == 0: save_checkpoint( { 'epoch': epoch + 1, 'state_dict': model.state_dict(), }, epoch + 1)
num_classes = 43 # create model model = dn.DenseNet3(args.layers, num_classes, normalizer=normalizer) model.eval() model = model.cuda() checkpoint = torch.load( "./checkpoints/{name}/checkpoint_{epochs}.pth.tar".format( name=args.name, epochs=args.epochs)) model.load_state_dict(checkpoint['state_dict']) attack = LinfPGDAttack(model=model, eps=args.epsilon, nb_iter=args.iters, eps_iter=args.iter_size, rand_init=True) nat_top1 = AverageMeter() adv_top1 = AverageMeter() for batch_index, (input, target) in enumerate(test_loader): print(batch_index * args.batch_size, '/', 10000) target = target.cuda() nat_input = input.detach().clone() nat_output = model(nat_input) nat_prec1 = accuracy(nat_output.data, target, topk=(1, ))[0]