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]