def main():
global args, best_prec1
args = parser.parse_args()
if args.tensorboard: configure("runs/%s" % (args.name))
# Data loading code
normalize = 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.ToTensor(),
transforms.Lambda(lambda x: F.pad(Variable(
x.unsqueeze(0), requires_grad=False, volatile=True),
(4, 4, 4, 4),
mode='reflect').data.squeeze()),
transforms.ToPILImage(),
transforms.RandomCrop(32),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
])
else:
transform_train = transforms.Compose([
transforms.ToTensor(),
normalize,
])
transform_test = transforms.Compose([transforms.ToTensor(), normalize])
kwargs = {'num_workers': 1, 'pin_memory': True}
assert (args.dataset == 'cifar10' or args.dataset == 'cifar100')
train_set = datasets.ImageFolder(root="../classes_data_with_black/train",
transform=transform_train)
val_set = datasets.ImageFolder(root="../classes_data_with_black/val",
transform=transform_test)
train_loader = torch.utils.data.DataLoader(dataset=train_set,
batch_size=args.batch_size,
shuffle=True,
**kwargs)
val_loader = torch.utils.data.DataLoader(dataset=val_set,
batch_size=args.batch_size,
shuffle=True,
**kwargs)
# create model
model = WideResNet(args.layers,
args.dataset == 'cifar10' and 10 or 100,
args.widen_factor,
dropRate=args.droprate)
# 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']
best_prec1 = checkpoint['best_prec1']
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 optimizer
criterion = nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
nesterov=args.nesterov,
weight_decay=args.weight_decay)
for epoch in range(args.start_epoch, args.epochs):
adjust_learning_rate(optimizer, epoch + 1)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch)
# evaluate on validation set
prec1 = validate(val_loader, model, criterion, epoch)
# remember best prec@1 and save checkpoint
is_best = prec1 > best_prec1
best_prec1 = max(prec1, best_prec1)
save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
}, is_best)
print('Best accuracy: ', best_prec1)
torch.save(model, 'model_original_mnist.pt')
def main():
global args, best_prec1, data
if args.tensorboard: configure("runs/%s" % (args.name))
# Data loading code
normalize = 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.ToTensor(),
transforms.Lambda(lambda x: F.pad(Variable(
x.unsqueeze(0), requires_grad=False, volatile=True),
(4, 4, 4, 4),
mode='reflect').data.squeeze()),
transforms.ToPILImage(),
transforms.RandomCrop(32),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
])
else:
transform_train = transforms.Compose([
transforms.ToTensor(),
normalize,
])
transform_test = transforms.Compose([transforms.ToTensor(), normalize])
kwargs = {'num_workers': args.num_workers, 'pin_memory': use_cuda}
assert (args.dataset == 'cifar10' or args.dataset == 'cifar100')
print("Creating the DataLoader...")
train_loader = torch.utils.data.DataLoader(
datasets.__dict__[args.dataset.upper()]('../data',
train=True,
download=True,
transform=transform_train),
batch_size=args.batch_size,
shuffle=True,
**kwargs)
val_loader = torch.utils.data.DataLoader(
datasets.__dict__[args.dataset.upper()]('../data',
train=False,
transform=transform_test),
batch_size=args.batch_size,
shuffle=True,
**kwargs)
# create model
print("Creating the model...")
model = WideResNet(args.layers,
args.dataset == 'cifar10' and 10 or 100,
args.widen_factor,
dropRate=args.droprate)
# 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
if use_cuda:
print("Moving the model to the GPU")
model = torch.nn.DataParallel(model, device_ids=device_ids)
model = 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']
best_prec1 = checkpoint['best_prec1']
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 optimizer
criterion = nn.CrossEntropyLoss()
if use_cuda:
criterion = criterion.cuda()
#optimizer = torch.optim.SGD(model.parameters(), args.lr,
# momentum=args.momentum, nesterov=args.nesterov,
# weight_decay=args.weight_decay)
optimizer = torch.optim.ASGD(model.parameters(), args.lr)
data = []
train_time = 0
for epoch in range(args.start_epoch, args.epochs):
adjust_learning_rate(optimizer, epoch + 1)
# train for one epoch
start = time.time()
train(train_loader, model, criterion, optimizer, epoch)
train_time += time.time() - start
# evaluate on validation set
datum = validate(val_loader, model, criterion, epoch)
data += [{
'train_time': train_time,
'epoch': epoch + 1,
**vars(args),
**datum
}]
df = pd.DataFrame(data)
_write_csv(df, id=f'{args.num_workers}_{args.seed}')
pprint({
k: v
for k, v in data[-1].items() if k in
['train_time', 'num_workers', 'test_loss', 'test_acc', 'epoch']
})
prec1 = datum['test_acc']
# remember best prec@1 and save checkpoint
is_best = prec1 > best_prec1
best_prec1 = max(prec1, best_prec1)
save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
}, is_best)
print('Best accuracy: ', best_prec1)
def main():
args = get_args()
if args.awp_gamma <= 0.0:
args.awp_warmup = np.infty
if not os.path.exists(args.fname):
os.makedirs(args.fname)
logger = logging.getLogger(__name__)
logging.basicConfig(format='[%(asctime)s] - %(message)s',
datefmt='%Y/%m/%d %H:%M:%S',
level=logging.DEBUG,
handlers=[
logging.FileHandler(
os.path.join(args.fname, 'output.log')),
logging.StreamHandler()
])
logger.info(args)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
train_transform = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.RandomRotation(15),
transforms.ToTensor(),
])
test_transform = transforms.Compose([
transforms.ToTensor(),
])
num_workers = 2
train_dataset = datasets.CIFAR100(args.data_dir,
train=True,
transform=train_transform,
download=True)
test_dataset = datasets.CIFAR100(args.data_dir,
train=False,
transform=test_transform,
download=True)
train_loader = torch.utils.data.DataLoader(
dataset=train_dataset,
batch_size=args.batch_size,
shuffle=True,
pin_memory=True,
num_workers=num_workers,
)
test_loader = torch.utils.data.DataLoader(
dataset=test_dataset,
batch_size=args.batch_size,
shuffle=False,
pin_memory=True,
num_workers=2,
)
epsilon = (args.epsilon / 255.)
pgd_alpha = (args.pgd_alpha / 255.)
if args.model == 'PreActResNet18':
model = PreActResNet18(num_classes=100)
proxy = PreActResNet18(num_classes=100)
elif args.model == 'WideResNet':
model = WideResNet(34,
100,
widen_factor=args.width_factor,
dropRate=0.0)
proxy = WideResNet(34,
100,
widen_factor=args.width_factor,
dropRate=0.0)
else:
raise ValueError("Unknown model")
model = model.cuda()
proxy = proxy.cuda()
if args.l2:
decay, no_decay = [], []
for name, param in model.named_parameters():
if 'bn' not in name and 'bias' not in name:
decay.append(param)
else:
no_decay.append(param)
params = [{
'params': decay,
'weight_decay': args.l2
}, {
'params': no_decay,
'weight_decay': 0
}]
else:
params = model.parameters()
opt = torch.optim.SGD(params,
lr=args.lr_max,
momentum=0.9,
weight_decay=5e-4)
proxy_opt = torch.optim.SGD(proxy.parameters(), lr=0.01)
awp_adversary = AdvWeightPerturb(model=model,
proxy=proxy,
proxy_optim=proxy_opt,
gamma=args.awp_gamma)
criterion = nn.CrossEntropyLoss()
epochs = args.epochs
if args.lr_schedule == 'superconverge':
lr_schedule = lambda t: np.interp([t], [
0, args.epochs * 2 // 5, args.epochs
], [0, args.lr_max, 0])[0]
# lr_schedule = lambda t: np.interp([t], [0, args.epochs], [0, args.lr_max])[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.
best_test_robust_acc = 0
if args.resume:
start_epoch = args.resume
model.load_state_dict(
torch.load(os.path.join(args.fname, f'model_{start_epoch-1}.pth')))
opt.load_state_dict(
torch.load(os.path.join(args.fname, f'opt_{start_epoch-1}.pth')))
logger.info(f'Resuming at epoch {start_epoch}')
if os.path.exists(os.path.join(args.fname, f'model_best.pth')):
best_test_robust_acc = torch.load(
os.path.join(args.fname, f'model_best.pth'))['test_robust_acc']
else:
start_epoch = 0
logger.info(
'Epoch \t Train Time \t Test Time \t LR \t \t Train Loss \t Train Acc \t Train Robust Loss \t Train Robust Acc \t Test Loss \t Test Acc \t Test Robust Loss \t Test Robust Acc'
)
for epoch in range(start_epoch, epochs):
model.train()
start_time = time.time()
train_loss = 0
train_acc = 0
train_robust_loss = 0
train_robust_acc = 0
train_n = 0
for i, (X, y) in enumerate(train_loader):
X, y = X.cuda(), y.cuda()
if args.mixup:
X, y_a, y_b, lam = mixup_data(X, y, args.mixup_alpha)
X, y_a, y_b = map(Variable, (X, y_a, y_b))
lr = lr_schedule(epoch + (i + 1) / len(train_loader))
opt.param_groups[0].update(lr=lr)
if args.attack == 'pgd':
# Random initialization
if args.mixup:
delta = attack_pgd(model,
X,
y,
epsilon,
pgd_alpha,
args.attack_iters,
args.restarts,
args.norm,
mixup=True,
y_a=y_a,
y_b=y_b,
lam=lam)
else:
delta = attack_pgd(model, X, y, epsilon, pgd_alpha,
args.attack_iters, args.restarts,
args.norm)
delta = delta.detach()
# Standard training
elif args.attack == 'none':
delta = torch.zeros_like(X)
X_adv = normalize(
torch.clamp(X + delta[:X.size(0)],
min=lower_limit,
max=upper_limit))
model.train()
# calculate adversarial weight perturbation and perturb it
if epoch >= args.awp_warmup:
# not compatible to mixup currently.
assert (not args.mixup)
awp = awp_adversary.calc_awp(inputs_adv=X_adv, targets=y)
awp_adversary.perturb(awp)
robust_output = model(X_adv)
if args.mixup:
robust_loss = mixup_criterion(criterion, robust_output, y_a,
y_b, lam)
else:
robust_loss = criterion(robust_output, y)
if args.l1:
for name, param in model.named_parameters():
if 'bn' not in name and 'bias' not in name:
robust_loss += args.l1 * param.abs().sum()
opt.zero_grad()
robust_loss.backward()
opt.step()
if epoch >= args.awp_warmup:
awp_adversary.restore(awp)
output = model(normalize(X))
if args.mixup:
loss = mixup_criterion(criterion, output, y_a, y_b, lam)
else:
loss = criterion(output, y)
train_robust_loss += robust_loss.item() * y.size(0)
train_robust_acc += (robust_output.max(1)[1] == y).sum().item()
train_loss += loss.item() * y.size(0)
train_acc += (output.max(1)[1] == y).sum().item()
train_n += y.size(0)
train_time = time.time()
model.eval()
test_loss = 0
test_acc = 0
test_robust_loss = 0
test_robust_acc = 0
test_n = 0
for i, (X, y) in enumerate(test_loader):
X, y = X.cuda(), y.cuda()
# Random initialization
if args.attack == 'none':
delta = torch.zeros_like(X)
else:
delta = attack_pgd(model, X, y, epsilon, pgd_alpha,
args.attack_iters_test, args.restarts,
args.norm)
delta = delta.detach()
robust_output = model(
normalize(
torch.clamp(X + delta[:X.size(0)],
min=lower_limit,
max=upper_limit)))
robust_loss = criterion(robust_output, y)
output = model(normalize(X))
loss = criterion(output, y)
test_robust_loss += robust_loss.item() * y.size(0)
test_robust_acc += (robust_output.max(1)[1] == y).sum().item()
test_loss += loss.item() * y.size(0)
test_acc += (output.max(1)[1] == y).sum().item()
test_n += y.size(0)
test_time = time.time()
logger.info(
'%d \t %.1f \t \t %.1f \t \t %.4f \t %.4f \t %.4f \t %.4f \t \t %.4f \t \t %.4f \t %.4f \t %.4f \t \t %.4f',
epoch, train_time - start_time, test_time - train_time, lr,
train_loss / train_n, train_acc / train_n,
train_robust_loss / train_n, train_robust_acc / train_n,
test_loss / test_n, test_acc / test_n, test_robust_loss / test_n,
test_robust_acc / test_n)
if (epoch + 1) % args.chkpt_iters == 0 or epoch + 1 == epochs:
torch.save(model.state_dict(),
os.path.join(args.fname, f'model_{epoch}.pth'))
torch.save(opt.state_dict(),
os.path.join(args.fname, f'opt_{epoch}.pth'))
# save best
if test_robust_acc / test_n > best_test_robust_acc:
torch.save(
{
'state_dict': model.state_dict(),
'test_robust_acc': test_robust_acc / test_n,
'test_robust_loss': test_robust_loss / test_n,
'test_loss': test_loss / test_n,
'test_acc': test_acc / test_n,
}, os.path.join(args.fname, f'model_best.pth'))
best_test_robust_acc = test_robust_acc / test_n
def main_worker(gpu, ngpus_per_node, args):
global best_acc1
args.gpu = gpu
if args.gpu is not None:
print("Use GPU: {} for training".format(args.gpu))
if args.distributed:
if args.dist_url == "env://" and args.rank == -1:
args.rank = int(os.environ["RANK"])
if args.multiprocessing_distributed:
# For multiprocessing distributed training, rank needs to be the
# global rank among all the processes
args.rank = args.rank * ngpus_per_node + gpu
dist.init_process_group(backend=args.dist_backend,
init_method=args.dist_url,
world_size=args.world_size,
rank=args.rank)
# create model
if args.arch == 'wrn':
num_classes = 1000
size = 64
config = [args.depth, args.k, args.drop, num_classes, size]
model = WideResNet(config)
elif args.pretrained:
print("=> using pre-trained model '{}'".format(args.arch))
model = models.__dict__[args.arch](pretrained=True)
else:
print("=> creating model '{}'".format(args.arch))
model = models.__dict__[args.arch]()
if args.distributed:
# For multiprocessing distributed, DistributedDataParallel constructor
# should always set the single device scope, otherwise,
# DistributedDataParallel will use all available devices.
if args.gpu is not None:
torch.cuda.set_device(args.gpu)
model.cuda(args.gpu)
# When using a single GPU per process and per
# DistributedDataParallel, we need to divide the batch size
# ourselves based on the total number of GPUs we have
args.batch_size = int(args.batch_size / ngpus_per_node)
args.workers = int(
(args.workers + ngpus_per_node - 1) / ngpus_per_node)
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.gpu])
else:
model.cuda()
# DistributedDataParallel will divide and allocate batch_size to all
# available GPUs if device_ids are not set
model = torch.nn.parallel.DistributedDataParallel(model)
elif args.gpu is not None:
torch.cuda.set_device(args.gpu)
model = model.cuda(args.gpu)
else:
# DataParallel will divide and allocate batch_size to all available GPUs
if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
model.features = torch.nn.DataParallel(model.features)
model.cuda()
else:
model = torch.nn.DataParallel(model).cuda()
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda(args.gpu)
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
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']
best_acc1 = checkpoint['best_acc1']
if args.gpu is not None:
# best_acc1 may be from a checkpoint from a different GPU
best_acc1 = best_acc1.to(args.gpu)
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, [50, 65, 75],
gamma=0.1)
cudnn.benchmark = True
# Data loading code
traindir = os.path.join(args.data, 'train')
valdir = os.path.join(args.data, 'val')
# normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
# std=[0.229, 0.224, 0.225])
normalize = transforms.Normalize(mean=[0.482, 0.458, 0.408],
std=[0.269, 0.261, 0.276])
train_dataset = datasets.ImageFolder(
traindir,
transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
if args.distributed:
train_sampler = torch.utils.data.distributed.DistributedSampler(
train_dataset)
else:
train_sampler = None
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=(train_sampler is None),
num_workers=args.workers,
pin_memory=True,
sampler=train_sampler)
val_loader = torch.utils.data.DataLoader(datasets.ImageFolder(
valdir, transforms.Compose([
transforms.ToTensor(),
normalize,
])),
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
if args.evaluate:
validate(val_loader, model, criterion, args)
return
begin = time.time()
for epoch in trange(args.start_epoch, args.epochs):
if args.distributed:
train_sampler.set_epoch(epoch)
# adjust_learning_rate(optimizer, epoch, args)
scheduler.step()
epoch_begin = time.time()
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch, args)
# evaluate on validation set
acc1 = validate(val_loader, model, criterion, args)
epoch_end = time.time()
# remember best acc@1 and save checkpoint
is_best = acc1 > best_acc1
best_acc1 = max(acc1, best_acc1)
print("epoch %d, time %.2f" % (epoch, epoch_end - epoch_begin))
if not args.multiprocessing_distributed or (
args.multiprocessing_distributed
and args.rank % ngpus_per_node == 0):
fname = "%s.pth.tar" % args.arch
if args.arch == "wrn":
fname = "%s-%d-%s.pth.tar" % (args.arch, args.depth, str(
args.k))
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_acc1': best_acc1,
'optimizer': optimizer.state_dict(),
},
is_best,
filename=fname)
print("num epoch %d, time %.2f" % (args.epochs, epoch_end - begin))
def main():
args = get_args()
if not os.path.exists(args.fname):
os.makedirs(args.fname)
logger = logging.getLogger(__name__)
logging.basicConfig(
format='[%(asctime)s] - %(message)s',
datefmt='%Y/%m/%d %H:%M:%S',
level=logging.DEBUG,
handlers=[
logging.FileHandler(
os.path.join(args.fname,
'eval.log' if args.eval else 'output.log')),
logging.StreamHandler()
])
logger.info(args)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
transforms = [Crop(32, 32), FlipLR()]
# transforms = [Crop(32, 32)]
if args.cutout:
transforms.append(Cutout(args.cutout_len, args.cutout_len))
if args.val:
try:
dataset = torch.load("cifar10_validation_split.pth")
except:
print(
"Couldn't find a dataset with a validation split, did you run "
"generate_validation.py?")
return
val_set = list(
zip(transpose(dataset['val']['data'] / 255.),
dataset['val']['labels']))
val_batches = Batches(val_set,
args.batch_size,
shuffle=False,
num_workers=2)
else:
dataset = cifar10(args.data_dir)
train_set = list(
zip(transpose(pad(dataset['train']['data'], 4) / 255.),
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)
test_set = list(
zip(transpose(dataset['test']['data'] / 255.),
dataset['test']['labels']))
test_batches = Batches(test_set,
args.batch_size,
shuffle=False,
num_workers=2)
trn_epsilon = (args.trn_epsilon / 255.)
trn_pgd_alpha = (args.trn_pgd_alpha / 255.)
tst_epsilon = (args.tst_epsilon / 255.)
tst_pgd_alpha = (args.tst_pgd_alpha / 255.)
if args.model == 'PreActResNet18':
model = PreActResNet18()
elif args.model == 'WideResNet':
model = WideResNet(34,
10,
widen_factor=args.width_factor,
dropRate=0.0)
elif args.model == 'DenseNet121':
model = DenseNet121()
elif args.model == 'ResNet18':
model = ResNet18()
else:
raise ValueError("Unknown model")
### temp testing ###
model = model.cuda()
# model = nn.DataParallel(model).cuda()
model.train()
##################################
# load pretrained model if needed
if args.trn_adv_models != 'None':
if args.trn_adv_arch == 'PreActResNet18':
trn_adv_model = PreActResNet18()
elif args.trn_adv_arch == 'WideResNet':
trn_adv_model = WideResNet(34,
10,
widen_factor=args.width_factor,
dropRate=0.0)
elif args.trn_adv_arch == 'DenseNet121':
trn_adv_model = DenseNet121()
elif args.trn_adv_arch == 'ResNet18':
trn_adv_model = ResNet18()
trn_adv_model = nn.DataParallel(trn_adv_model).cuda()
trn_adv_model.load_state_dict(
torch.load(
os.path.join('./adv_models', args.trn_adv_models,
'model_best.pth'))['state_dict'])
logger.info(f'loaded adv_model: {args.trn_adv_models}')
else:
trn_adv_model = None
if args.tst_adv_models != 'None':
if args.tst_adv_arch == 'PreActResNet18':
tst_adv_model = PreActResNet18()
elif args.tst_adv_arch == 'WideResNet':
tst_adv_model = WideResNet(34,
10,
widen_factor=args.width_factor,
dropRate=0.0)
elif args.tst_adv_arch == 'DenseNet121':
tst_adv_model = DenseNet121()
elif args.tst_adv_arch == 'ResNet18':
tst_adv_model = ResNet18()
### temp testing ###
tst_adv_model = tst_adv_model.cuda()
tst_adv_model.load_state_dict(
torch.load(
os.path.join('./adv_models', args.tst_adv_models,
'model_best.pth')))
# tst_adv_model = nn.DataParallel(tst_adv_model).cuda()
# tst_adv_model.load_state_dict(torch.load(os.path.join('./adv_models',args.tst_adv_models, 'model_best.pth'))['state_dict'])
logger.info(f'loaded adv_model: {args.tst_adv_models}')
else:
tst_adv_model = None
##################################
if args.l2:
decay, no_decay = [], []
for name, param in model.named_parameters():
if 'bn' not in name and 'bias' not in name:
decay.append(param)
else:
no_decay.append(param)
params = [{
'params': decay,
'weight_decay': args.l2
}, {
'params': no_decay,
'weight_decay': 0
}]
else:
params = model.parameters()
opt = torch.optim.SGD(params,
lr=args.lr_max,
momentum=0.9,
weight_decay=5e-4)
criterion = nn.CrossEntropyLoss()
if args.trn_attack == 'free':
delta = torch.zeros(args.batch_size, 3, 32, 32).cuda()
delta.requires_grad = True
elif args.trn_attack == 'fgsm' and args.trn_fgsm_init == 'previous':
delta = torch.zeros(args.batch_size, 3, 32, 32).cuda()
delta.requires_grad = True
if args.trn_attack == 'free':
epochs = int(math.ceil(args.epochs / args.trn_attack_iters))
else:
epochs = args.epochs
if args.lr_schedule == 'superconverge':
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.
elif args.lr_schedule == 'linear':
lr_schedule = lambda t: np.interp([t], [
0, args.epochs // 3, args.epochs * 2 // 3, args.epochs
], [args.lr_max, args.lr_max, args.lr_max / 10, args.lr_max / 100])[0]
elif args.lr_schedule == 'onedrop':
def lr_schedule(t):
if t < args.lr_drop_epoch:
return args.lr_max
else:
return args.lr_one_drop
elif args.lr_schedule == 'multipledecay':
def lr_schedule(t):
return args.lr_max - (t //
(args.epochs // 10)) * (args.lr_max / 10)
elif args.lr_schedule == 'cosine':
def lr_schedule(t):
return args.lr_max * 0.5 * (1 + np.cos(t / args.epochs * np.pi))
best_test_robust_acc = 0
best_val_robust_acc = 0
if args.resume:
### temp testing ###
model.load_state_dict(
torch.load(os.path.join(args.fname, 'model_best.pth')))
start_epoch = args.resume
# model.load_state_dict(torch.load(os.path.join(args.fname, f'model_{start_epoch-1}.pth')))
# opt.load_state_dict(torch.load(os.path.join(args.fname, f'opt_{start_epoch-1}.pth')))
# logger.info(f'Resuming at epoch {start_epoch}')
# best_test_robust_acc = torch.load(os.path.join(args.fname, f'model_best.pth'))['test_robust_acc']
if args.val:
best_val_robust_acc = torch.load(
os.path.join(args.fname, f'model_val.pth'))['val_robust_acc']
else:
start_epoch = 0
if args.eval:
if not args.resume:
logger.info(
"No model loaded to evaluate, specify with --resume FNAME")
return
logger.info("[Evaluation mode]")
logger.info(
'Epoch \t Train Time \t Test Time \t LR \t \t Train Loss \t Train Acc \t Train Robust Loss \t Train Robust Acc \t Test Loss \t Test Acc \t Test Robust Loss \t Test Robust Acc'
)
for epoch in range(start_epoch, epochs):
model.train()
start_time = time.time()
train_loss = 0
train_acc = 0
train_robust_loss = 0
train_robust_acc = 0
train_n = 0
for i, batch in enumerate(train_batches):
if args.eval:
break
X, y = batch['input'], batch['target']
if args.mixup:
X, y_a, y_b, lam = mixup_data(X, y, args.mixup_alpha)
X, y_a, y_b = map(Variable, (X, y_a, y_b))
lr = lr_schedule(epoch + (i + 1) / len(train_batches))
opt.param_groups[0].update(lr=lr)
if args.trn_attack == 'pgd':
# Random initialization
if args.mixup:
delta = attack_pgd(model,
X,
y,
trn_epsilon,
trn_pgd_alpha,
args.trn_attack_iters,
args.trn_restarts,
args.trn_norm,
mixup=True,
y_a=y_a,
y_b=y_b,
lam=lam,
adv_models=trn_adv_model)
else:
delta = attack_pgd(model,
X,
y,
trn_epsilon,
trn_pgd_alpha,
args.trn_attack_iters,
args.trn_restarts,
args.trn_norm,
adv_models=trn_adv_model)
delta = delta.detach()
elif args.trn_attack == 'fgsm':
delta = attack_pgd(model,
X,
y,
trn_epsilon,
args.trn_fgsm_alpha * trn_epsilon,
1,
1,
args.trn_norm,
adv_models=trn_adv_model,
rand_init=args.trn_fgsm_init)
delta = delta.detach()
# Standard training
elif args.trn_attack == 'none':
delta = torch.zeros_like(X)
# The Momentum Iterative Attack
elif args.trn_attack == 'tmim':
if trn_adv_model is None:
adversary = MomentumIterativeAttack(
model,
nb_iter=args.trn_attack_iters,
eps=trn_epsilon,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
eps_iter=trn_pgd_alpha,
clip_min=0,
clip_max=1,
targeted=False)
else:
trn_adv_model = nn.Sequential(
NormalizeByChannelMeanStd(CIFAR10_MEAN, CIFAR10_STD),
trn_adv_model)
adversary = MomentumIterativeAttack(
trn_adv_model,
nb_iter=args.trn_attack_iters,
eps=trn_epsilon,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
eps_iter=trn_pgd_alpha,
clip_min=0,
clip_max=1,
targeted=False)
data_adv = adversary.perturb(X, y)
delta = data_adv - X
delta = delta.detach()
robust_output = model(
normalize(
torch.clamp(X + delta[:X.size(0)],
min=lower_limit,
max=upper_limit)))
if args.mixup:
robust_loss = mixup_criterion(criterion, robust_output, y_a,
y_b, lam)
else:
robust_loss = criterion(robust_output, y)
if args.l1:
for name, param in model.named_parameters():
if 'bn' not in name and 'bias' not in name:
robust_loss += args.l1 * param.abs().sum()
opt.zero_grad()
robust_loss.backward()
opt.step()
output = model(normalize(X))
if args.mixup:
loss = mixup_criterion(criterion, output, y_a, y_b, lam)
else:
loss = criterion(output, y)
train_robust_loss += robust_loss.item() * y.size(0)
train_robust_acc += (robust_output.max(1)[1] == y).sum().item()
train_loss += loss.item() * y.size(0)
train_acc += (output.max(1)[1] == y).sum().item()
train_n += y.size(0)
train_time = time.time()
model.eval()
test_loss = 0
test_acc = 0
test_robust_loss = 0
test_robust_acc = 0
test_n = 0
for i, batch in enumerate(test_batches):
X, y = batch['input'], batch['target']
# Random initialization
if args.tst_attack == 'none':
delta = torch.zeros_like(X)
elif args.tst_attack == 'pgd':
delta = attack_pgd(model,
X,
y,
tst_epsilon,
tst_pgd_alpha,
args.tst_attack_iters,
args.tst_restarts,
args.tst_norm,
adv_models=tst_adv_model,
rand_init=args.tst_fgsm_init)
elif args.tst_attack == 'fgsm':
delta = attack_pgd(model,
X,
y,
tst_epsilon,
tst_epsilon,
1,
1,
args.tst_norm,
rand_init=args.tst_fgsm_init,
adv_models=tst_adv_model)
# The Momentum Iterative Attack
elif args.tst_attack == 'tmim':
if tst_adv_model is None:
adversary = MomentumIterativeAttack(
model,
nb_iter=args.tst_attack_iters,
eps=tst_epsilon,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
eps_iter=tst_pgd_alpha,
clip_min=0,
clip_max=1,
targeted=False)
else:
tmp_model = nn.Sequential(
NormalizeByChannelMeanStd(cifar10_mean, cifar10_std),
tst_adv_model).to(device)
adversary = MomentumIterativeAttack(
tmp_model,
nb_iter=args.tst_attack_iters,
eps=tst_epsilon,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
eps_iter=tst_pgd_alpha,
clip_min=0,
clip_max=1,
targeted=False)
data_adv = adversary.perturb(X, y)
delta = data_adv - X
# elif args.tst_attack == 'pgd':
# if tst_adv_model is None:
# tmp_model = nn.Sequential(NormalizeByChannelMeanStd(cifar10_mean, cifar10_std), model).to(device)
# adversary = PGDAttack(tmp_model, nb_iter=args.tst_attack_iters,
# eps = tst_epsilon,
# loss_fn=nn.CrossEntropyLoss(reduction="sum"),
# eps_iter=tst_pgd_alpha, clip_min = 0, clip_max = 1, targeted=False)
# else:
# tmp_model = nn.Sequential(NormalizeByChannelMeanStd(cifar10_mean, cifar10_std), tst_adv_model).to(device)
# adversary = PGDAttack(tmp_model, nb_iter=args.tst_attack_iters,
# eps = tst_epsilon,
# loss_fn=nn.CrossEntropyLoss(reduction="sum"),
# eps_iter=tst_pgd_alpha, clip_min = 0, clip_max = 1, targeted=False)
# data_adv = adversary.perturb(X, y)
# delta = data_adv - X
delta = delta.detach()
robust_output = model(
normalize(
torch.clamp(X + delta[:X.size(0)],
min=lower_limit,
max=upper_limit)))
robust_loss = criterion(robust_output, y)
output = model(normalize(X))
loss = criterion(output, y)
test_robust_loss += robust_loss.item() * y.size(0)
test_robust_acc += (robust_output.max(1)[1] == y).sum().item()
test_loss += loss.item() * y.size(0)
test_acc += (output.max(1)[1] == y).sum().item()
test_n += y.size(0)
test_time = time.time()
if args.val:
val_loss = 0
val_acc = 0
val_robust_loss = 0
val_robust_acc = 0
val_n = 0
for i, batch in enumerate(val_batches):
X, y = batch['input'], batch['target']
# Random initialization
if args.tst_attack == 'none':
delta = torch.zeros_like(X)
elif args.tst_attack == 'pgd':
delta = attack_pgd(model,
X,
y,
tst_epsilon,
tst_pgd_alpha,
args.tst_attack_iters,
args.tst_restarts,
args.tst_norm,
early_stop=args.eval)
elif args.tst_attack == 'fgsm':
delta = attack_pgd(model,
X,
y,
tst_epsilon,
tst_epsilon,
1,
1,
args.tst_norm,
early_stop=args.eval,
rand_init=args.tst_fgsm_init)
delta = delta.detach()
robust_output = model(
normalize(
torch.clamp(X + delta[:X.size(0)],
min=lower_limit,
max=upper_limit)))
robust_loss = criterion(robust_output, y)
output = model(normalize(X))
loss = criterion(output, y)
val_robust_loss += robust_loss.item() * y.size(0)
val_robust_acc += (robust_output.max(1)[1] == y).sum().item()
val_loss += loss.item() * y.size(0)
val_acc += (output.max(1)[1] == y).sum().item()
val_n += y.size(0)
if not args.eval:
logger.info(
'%d \t %.1f \t \t %.1f \t \t %.4f \t %.4f \t %.4f \t %.4f \t \t %.4f \t \t %.4f \t %.4f \t %.4f \t \t %.4f',
epoch, train_time - start_time, test_time - train_time, lr,
train_loss / train_n, train_acc / train_n,
train_robust_loss / train_n, train_robust_acc / train_n,
test_loss / test_n, test_acc / test_n,
test_robust_loss / test_n, test_robust_acc / test_n)
if args.val:
logger.info('validation %.4f \t %.4f \t %.4f \t %.4f',
val_loss / val_n, val_acc / val_n,
val_robust_loss / val_n, val_robust_acc / val_n)
if val_robust_acc / val_n > best_val_robust_acc:
torch.save(
{
'state_dict': model.state_dict(),
'test_robust_acc': test_robust_acc / test_n,
'test_robust_loss': test_robust_loss / test_n,
'test_loss': test_loss / test_n,
'test_acc': test_acc / test_n,
'val_robust_acc': val_robust_acc / val_n,
'val_robust_loss': val_robust_loss / val_n,
'val_loss': val_loss / val_n,
'val_acc': val_acc / val_n,
}, os.path.join(args.fname, f'model_val.pth'))
best_val_robust_acc = val_robust_acc / val_n
# save checkpoint
if (epoch + 1) % args.chkpt_iters == 0 or epoch + 1 == epochs:
torch.save(model.state_dict(),
os.path.join(args.fname, f'model_{epoch}.pth'))
torch.save(opt.state_dict(),
os.path.join(args.fname, f'opt_{epoch}.pth'))
# save best
if test_robust_acc / test_n > best_test_robust_acc:
torch.save(
{
'state_dict': model.state_dict(),
'test_robust_acc': test_robust_acc / test_n,
'test_robust_loss': test_robust_loss / test_n,
'test_loss': test_loss / test_n,
'test_acc': test_acc / test_n,
}, os.path.join(args.fname, f'model_best.pth'))
best_test_robust_acc = test_robust_acc / test_n
else:
logger.info(
'%d \t %.1f \t \t %.1f \t \t %.4f \t %.4f \t %.4f \t %.4f \t \t %.4f \t \t %.4f \t %.4f \t %.4f \t \t %.4f',
epoch, train_time - start_time, test_time - train_time, -1, -1,
-1, -1, -1, test_loss / test_n, test_acc / test_n,
test_robust_loss / test_n, test_robust_acc / test_n)
return
def main():
args = get_args()
if args.awp_gamma <= 0.0:
args.awp_warmup = np.infty
fname = None
if args.sample == 100:
fname = args.output_dir + "/default/" + args.attack + "/"
else:
fname = args.output_dir + "/" + str(
args.sample) + "sample/" + args.attack + "/"
if args.attack == "combine":
fname += args.list + "/"
if not os.path.exists(fname):
os.makedirs(fname)
logger = logging.getLogger(__name__)
logging.basicConfig(
format='[%(asctime)s] - %(message)s',
datefmt='%Y/%m/%d %H:%M:%S',
level=logging.DEBUG,
handlers=[
logging.FileHandler(
os.path.join(fname,
'eval.log' if args.eval else 'output.log')),
logging.StreamHandler()
])
logger.info(args)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
shuffle = False
# setup data loader
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
])
transform_test = transforms.Compose([
transforms.ToTensor(),
])
train_set = torchvision.datasets.CIFAR10(root='../data',
train=True,
download=True,
transform=transform_train)
test_set = torchvision.datasets.CIFAR10(root='../data',
train=False,
download=True,
transform=transform_test)
if args.attack == "all":
train_data = np.array(train_set.data) / 255.
train_data = transpose(train_data).astype(np.float32)
train_labels = np.array(train_set.targets)
oversampled_train_data = np.tile(train_data, (11, 1, 1, 1))
oversampled_train_labels = np.tile(train_labels, (11))
train_set = list(
zip(torch.from_numpy(oversampled_train_data),
torch.from_numpy(oversampled_train_labels)))
elif args.attack == "combine":
train_data = np.array(train_set.data) / 255.
train_data = transpose(train_data).astype(np.float32)
train_labels = np.array(train_set.targets)
oversampled_train_data = train_data.copy()
oversampled_train_labels = train_labels.copy()
logger.info("Attacks")
attacks = args.list.split("_")
logger.info(attacks)
oversampled_train_data = np.tile(train_data, (len(attacks), 1, 1, 1))
oversampled_train_labels = np.tile(train_labels, (len(attacks)))
train_set = list(
zip(torch.from_numpy(oversampled_train_data),
torch.from_numpy(oversampled_train_labels)))
else:
train_data = np.array(train_set.data) / 255.
train_data = transpose(train_data).astype(np.float32)
train_labels = np.array(train_set.targets)
train_set = list(
zip(torch.from_numpy(train_data), torch.from_numpy(train_labels)))
test_data = np.array(test_set.data) / 255.
test_data = transpose(test_data).astype(np.float32)
test_labels = np.array(test_set.targets)
test_set = list(
zip(torch.from_numpy(test_data), torch.from_numpy(test_labels)))
if args.sample != 100:
n = len(train_set)
n_sample = int(n * args.sample / 100)
np.random.shuffle(train_set)
train_set = train_set[:n_sample]
train_batches = Batches(train_set,
args.batch_size,
shuffle=shuffle,
set_random_choices=False,
num_workers=2)
test_batches = Batches(test_set,
args.batch_size,
shuffle=False,
num_workers=0)
train_adv_images = None
train_adv_labels = None
test_adv_images = None
test_adv_labels = None
adv_dir = "adv_examples/{}/".format(args.attack)
train_path = adv_dir + "train.pth"
test_path = adv_dir + "test.pth"
if args.attack in TOOLBOX_ADV_ATTACK_LIST:
adv_train_data = torch.load(train_path)
train_adv_images = adv_train_data["adv"]
train_adv_labels = adv_train_data["label"]
adv_test_data = torch.load(test_path)
test_adv_images = adv_test_data["adv"]
test_adv_labels = adv_test_data["label"]
elif args.attack in ["ffgsm", "mifgsm", "tpgd"]:
adv_data = {}
adv_data["adv"], adv_data["label"] = torch.load(train_path)
train_adv_images = adv_data["adv"].numpy()
train_adv_labels = adv_data["label"].numpy()
adv_data = {}
adv_data["adv"], adv_data["label"] = torch.load(test_path)
test_adv_images = adv_data["adv"].numpy()
test_adv_labels = adv_data["label"].numpy()
elif args.attack == "combine":
print("Attacks")
attacks = args.list.split("_")
print(attacks)
for i in range(len(attacks)):
_adv_dir = "adv_examples/{}/".format(attacks[i])
train_path = _adv_dir + "train.pth"
test_path = _adv_dir + "test.pth"
adv_train_data = torch.load(train_path)
adv_test_data = torch.load(test_path)
if i == 0:
train_adv_images = adv_train_data["adv"]
train_adv_labels = adv_train_data["label"]
test_adv_images = adv_test_data["adv"]
test_adv_labels = adv_test_data["label"]
else:
train_adv_images = np.concatenate(
(train_adv_images, adv_train_data["adv"]))
train_adv_labels = np.concatenate(
(train_adv_labels, adv_train_data["label"]))
test_adv_images = np.concatenate(
(test_adv_images, adv_test_data["adv"]))
test_adv_labels = np.concatenate(
(test_adv_labels, adv_test_data["label"]))
elif args.attack == "all":
print("Loading attacks")
ATTACK_LIST = [
"autoattack", "autopgd", "bim", "cw", "deepfool", "fgsm",
"newtonfool", "pgd", "pixelattack", "spatialtransformation",
"squareattack"
]
for i in range(len(ATTACK_LIST)):
print("Attack: ", ATTACK_LIST[i])
_adv_dir = "adv_examples/{}/".format(ATTACK_LIST[i])
train_path = _adv_dir + "train.pth"
test_path = _adv_dir + "test.pth"
adv_train_data = torch.load(train_path)
adv_test_data = torch.load(test_path)
if i == 0:
train_adv_images = adv_train_data["adv"]
train_adv_labels = adv_train_data["label"]
test_adv_images = adv_test_data["adv"]
test_adv_labels = adv_test_data["label"]
else:
train_adv_images = np.concatenate(
(train_adv_images, adv_train_data["adv"]))
train_adv_labels = np.concatenate(
(train_adv_labels, adv_train_data["label"]))
test_adv_images = np.concatenate(
(test_adv_images, adv_test_data["adv"]))
test_adv_labels = np.concatenate(
(test_adv_labels, adv_test_data["label"]))
else:
raise ValueError("Unknown adversarial data")
print("")
print("Train Adv Attack Data: ", args.attack)
print("Dataset shape: ", train_adv_images.shape)
print("Dataset type: ", type(train_adv_images))
print("Label shape: ", len(train_adv_labels))
print("")
train_adv_set = list(zip(train_adv_images, train_adv_labels))
if args.sample != 100:
n = len(train_adv_set)
n_sample = int(n * args.sample / 100)
np.random.shuffle(train_adv_set)
train_adv_set = train_adv_set[:n_sample]
train_adv_batches = Batches(train_adv_set,
args.batch_size,
shuffle=shuffle,
set_random_choices=False,
num_workers=0)
test_adv_set = list(zip(test_adv_images, test_adv_labels))
test_adv_batches = Batches(test_adv_set,
args.batch_size,
shuffle=False,
num_workers=0)
epsilon = (args.epsilon / 255.)
pgd_alpha = (args.pgd_alpha / 255.)
if args.model == "ResNet18":
model = resnet18(pretrained=True)
proxy = resnet18(pretrained=True)
elif args.model == 'PreActResNet18':
model = PreActResNet18()
proxy = PreActResNet18()
elif args.model == 'WideResNet':
model = WideResNet(34,
10,
widen_factor=args.width_factor,
dropRate=0.0)
proxy = WideResNet(34,
10,
widen_factor=args.width_factor,
dropRate=0.0)
else:
raise ValueError("Unknown model")
model = nn.DataParallel(model).cuda()
proxy = nn.DataParallel(proxy).cuda()
if args.l2:
decay, no_decay = [], []
for name, param in model.named_parameters():
if 'bn' not in name and 'bias' not in name:
decay.append(param)
else:
no_decay.append(param)
params = [{
'params': decay,
'weight_decay': args.l2
}, {
'params': no_decay,
'weight_decay': 0
}]
else:
params = model.parameters()
opt = torch.optim.SGD(params,
lr=args.lr_max,
momentum=0.9,
weight_decay=5e-4)
proxy_opt = torch.optim.SGD(proxy.parameters(), lr=0.01)
awp_adversary = AdvWeightPerturb(model=model,
proxy=proxy,
proxy_optim=proxy_opt,
gamma=args.awp_gamma)
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':
epochs = int(math.ceil(args.epochs / args.attack_iters))
else:
epochs = args.epochs
if args.lr_schedule == 'superconverge':
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.
elif args.lr_schedule == 'linear':
lr_schedule = lambda t: np.interp([t], [
0, args.epochs // 3, args.epochs * 2 // 3, args.epochs
], [args.lr_max, args.lr_max, args.lr_max / 10, args.lr_max / 100])[0]
elif args.lr_schedule == 'onedrop':
def lr_schedule(t):
if t < args.lr_drop_epoch:
return args.lr_max
else:
return args.lr_one_drop
elif args.lr_schedule == 'multipledecay':
def lr_schedule(t):
return args.lr_max - (t //
(args.epochs // 10)) * (args.lr_max / 10)
elif args.lr_schedule == 'cosine':
def lr_schedule(t):
return args.lr_max * 0.5 * (1 + np.cos(t / args.epochs * np.pi))
elif args.lr_schedule == 'cyclic':
lr_schedule = lambda t: np.interp(
[t], [0, 0.4 * args.epochs, args.epochs], [0, args.lr_max, 0])[0]
best_test_robust_acc = 0
best_val_robust_acc = 0
if args.resume:
start_epoch = args.resume
model.load_state_dict(
torch.load(os.path.join(fname, f'model_{start_epoch-1}.pth')))
opt.load_state_dict(
torch.load(os.path.join(fname, f'opt_{start_epoch-1}.pth')))
logger.info(f'Resuming at epoch {start_epoch}')
if os.path.exists(os.path.join(fname, f'model_best.pth')):
best_test_robust_acc = torch.load(
os.path.join(fname, f'model_best.pth'))['test_robust_acc']
if args.val:
best_val_robust_acc = torch.load(
os.path.join(args.output_dir,
f'model_val.pth'))['val_robust_acc']
else:
start_epoch = 0
if args.eval:
if not args.resume:
logger.info(
"No model loaded to evaluate, specify with --resume FNAME")
return
logger.info("[Evaluation mode]")
model.cuda()
model.eval()
# Evaluate on original test data
test_acc = 0
test_n = 0
for i, batch in enumerate(test_batches):
X, y = batch['input'], batch['target']
clean_input = normalize(X)
output = model(clean_input)
test_acc += (output.max(1)[1] == y).sum().item()
test_n += y.size(0)
logger.info('Intial Accuracy on Original Test Data: %.4f (Test Acc)',
test_acc / test_n)
test_adv_acc = 0
test_adv_n = 0
for i, batch in enumerate(test_adv_batches):
adv_input = normalize(batch['input'])
y = batch['target']
robust_output = model(adv_input)
test_adv_acc += (robust_output.max(1)[1] == y).sum().item()
test_adv_n += y.size(0)
logger.info(
'Intial Accuracy on Adversarial Test Data: %.4f (Test Robust Acc)',
test_adv_acc / test_adv_n)
model.train()
logger.info(
'Epoch \t Train Acc \t Train Robust Acc \t Test Acc \t Test Robust Acc'
)
for epoch in range(start_epoch, epochs):
start_time = time.time()
train_loss = 0
train_acc = 0
train_robust_loss = 0
train_robust_acc = 0
train_n = 0
for i, (batch,
adv_batch) in enumerate(zip(train_batches, train_adv_batches)):
if args.eval:
break
X, y = batch['input'], batch['target']
if args.mixup:
X, y_a, y_b, lam = mixup_data(X, y, args.mixup_alpha)
X, y_a, y_b = map(Variable, (X, y_a, y_b))
lr = lr_schedule(epoch + (i + 1) / len(train_batches))
opt.param_groups[0].update(lr=lr)
X_adv = normalize(adv_batch["input"])
y_adv = adv_batch["target"]
model.train()
# calculate adversarial weight perturbation and perturb it
if epoch >= args.awp_warmup:
# not compatible to mixup currently.
assert (not args.mixup)
awp = awp_adversary.calc_awp(inputs_adv=X_adv, targets=y_adv)
awp_adversary.perturb(awp)
robust_output = model(X_adv)
if args.mixup:
robust_loss = mixup_criterion(criterion, robust_output, y_a,
y_b, lam)
else:
robust_loss = criterion(robust_output, y_adv)
if args.l1:
for name, param in model.named_parameters():
if 'bn' not in name and 'bias' not in name:
robust_loss += args.l1 * param.abs().sum()
opt.zero_grad()
robust_loss.backward()
opt.step()
if epoch >= args.awp_warmup:
awp_adversary.restore(awp)
output = model(normalize(X))
if args.mixup:
loss = mixup_criterion(criterion, output, y_a, y_b, lam)
else:
loss = criterion(output, y)
train_robust_loss += robust_loss.item() * y_adv.size(0)
train_robust_acc += (robust_output.max(1)[1] == y_adv).sum().item()
train_loss += loss.item() * y.size(0)
train_acc += (output.max(1)[1] == y).sum().item()
train_n += y.size(0)
train_time = time.time()
model.eval()
test_loss = 0
test_acc = 0
test_n = 0
test_robust_loss = 0
test_robust_acc = 0
test_robust_n = 0
for i, batch in enumerate(test_batches):
X, y = normalize(batch['input']), batch['target']
output = model(X)
loss = criterion(output, y)
test_loss += loss.item() * y.size(0)
test_acc += (output.max(1)[1] == y).sum().item()
test_n += y.size(0)
for i, adv_batch in enumerate(test_adv_batches):
X_adv, y_adv = normalize(adv_batch["input"]), adv_batch["target"]
robust_output = model(X_adv)
robust_loss = criterion(robust_output, y_adv)
test_robust_loss += robust_loss.item() * y_adv.size(0)
test_robust_acc += (robust_output.max(1)[1] == y_adv).sum().item()
test_robust_n += y_adv.size(0)
test_time = time.time()
logger.info('%d \t %.3f \t\t %.3f \t\t\t %.3f \t\t %.3f', epoch,
train_acc / train_n, train_robust_acc / train_n,
test_acc / test_n, test_robust_acc / test_robust_n)
# save checkpoint
if (epoch + 1) % args.chkpt_iters == 0 or epoch + 1 == epochs:
torch.save(model.state_dict(),
os.path.join(fname, f'model_{epoch}.pth'))
torch.save(opt.state_dict(), os.path.join(fname,
f'opt_{epoch}.pth'))
# save best
if test_robust_acc / test_n > best_test_robust_acc:
torch.save(
{
'state_dict': model.state_dict(),
'test_robust_acc': test_robust_acc / test_robust_n,
'test_robust_loss': test_robust_loss / test_robust_n,
'test_loss': test_loss / test_n,
'test_acc': test_acc / test_n,
}, os.path.join(fname, f'model_best.pth'))
best_test_robust_acc = test_robust_acc / test_n
def main():
args = get_args()
if not os.path.exists(args.fname):
os.makedirs(args.fname)
logger = logging.getLogger(__name__)
logging.basicConfig(
format='[%(asctime)s] - %(message)s',
datefmt='%Y/%m/%d %H:%M:%S',
level=logging.DEBUG,
handlers=[
logging.FileHandler(
os.path.join(args.fname,
'eval.log' if args.eval else 'output.log')),
logging.StreamHandler()
])
logger.info(args)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
transforms = [Crop(32, 32), FlipLR()]
dataset = cifar10(args.data_dir)
train_set = list(
zip(transpose(pad(dataset['train']['data'], 4) / 255.),
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)
test_set = list(
zip(transpose(dataset['test']['data'] / 255.),
dataset['test']['labels']))
test_batches = Batches(test_set,
args.batch_size,
shuffle=False,
num_workers=2)
epsilon = (args.epsilon / 255.)
pgd_alpha = (args.pgd_alpha / 255.)
if args.model == 'PreActResNet18':
model = PreActResNet18()
elif args.model == 'WideResNet':
model = WideResNet(34,
10,
widen_factor=args.width_factor,
dropRate=0.0)
else:
raise ValueError("Unknown model")
model = nn.DataParallel(model).cuda()
model.train()
params = model.parameters()
opt = torch.optim.SGD(params,
lr=args.lr_max,
momentum=0.9,
weight_decay=5e-4)
criterion = nn.CrossEntropyLoss()
epochs = args.epochs
if args.lr_schedule == 'cyclic':
lr_schedule = lambda t: np.interp(
[t], [0, args.epochs // 2, args.epochs], [0, args.lr_max, 0])[0]
elif args.lr_schedule == 'onedrop':
def lr_schedule(t):
if t < args.lr_drop:
return args.lr_max
else:
return args.lr_max / 10.
best_test_robust_acc = 0
best_val_robust_acc = 0
if args.resume:
start_epoch = args.resume
model.load_state_dict(
torch.load(os.path.join(args.fname, f'model_{start_epoch-1}.pth')))
opt.load_state_dict(
torch.load(os.path.join(args.fname, f'opt_{start_epoch-1}.pth')))
logger.info(f'Resuming at epoch {start_epoch}')
best_test_robust_acc = torch.load(
os.path.join(args.fname, f'model_best.pth'))['test_robust_acc']
else:
start_epoch = 0
if args.eval:
if not args.resume:
logger.info(
"No model loaded to evaluate, specify with --resume FNAME")
return
logger.info("[Evaluation mode]")
logger.info(
'Epoch \t Train Time \t Test Time \t LR \t \t Train Robust Loss \t Train Robust Acc \t Test Loss \t Test Acc \t Test Robust Loss \t Test Robust Acc \t Defence Mean \t Attack Mean'
)
for epoch in range(start_epoch, epochs):
model.train()
start_time = time.time()
train_robust_loss = 0
train_robust_acc = 0
train_n = 0
defence_mean = 0
for i, batch in enumerate(train_batches):
if args.eval:
break
X, y = batch['input'], batch['target']
lr = lr_schedule(epoch + (i + 1) / len(train_batches))
opt.param_groups[0].update(lr=lr)
if args.attack == 'fgsm':
delta = attack_pgd(model,
X,
y,
epsilon,
args.fgsm_alpha * epsilon,
1,
1,
'l_inf',
fgsm_init=args.fgsm_init)
elif args.attack == 'multigrad':
delta = multi_grad(model, X, y, epsilon,
args.fgsm_alpha * epsilon, args.multi_samps,
args.multi_th, args.multi_parallel)
elif args.attack == 'zerograd':
delta = zero_grad(model, X, y, epsilon,
args.fgsm_alpha * epsilon, args.zero_qval,
args.zero_iters, args.fgsm_init)
elif args.attack == 'pgd':
delta = attack_pgd(model, X, y, epsilon, pgd_alpha, 10, 1,
'l_inf')
delta = delta.detach()
robust_output = model(
normalize(
torch.clamp(X + delta[:X.size(0)],
min=lower_limit,
max=upper_limit)))
robust_loss = criterion(robust_output, y)
opt.zero_grad()
robust_loss.backward()
opt.step()
train_robust_loss += robust_loss.item() * y.size(0)
train_robust_acc += (robust_output.max(1)[1] == y).sum().item()
train_n += y.size(0)
defence_mean += torch.mean(torch.abs(delta)) * y.size(0)
train_time = time.time()
model.eval()
test_loss = 0
test_acc = 0
test_robust_loss = 0
test_robust_acc = 0
test_n = 0
attack_mean = 0
for i, batch in enumerate(test_batches):
if not epoch + 1 == epochs and not args.full_test and i > len(
test_batches) / 10:
break
X, y = batch['input'], batch['target']
delta = attack_pgd(model,
X,
y,
epsilon,
pgd_alpha,
args.test_iters,
args.test_restarts,
'l_inf',
early_stop=args.eval)
delta = delta.detach()
robust_output = model(
normalize(
torch.clamp(X + delta[:X.size(0)],
min=lower_limit,
max=upper_limit)))
robust_loss = criterion(robust_output, y)
output = model(normalize(X))
loss = criterion(output, y)
test_robust_loss += robust_loss.item() * y.size(0)
test_robust_acc += (robust_output.max(1)[1] == y).sum().item()
test_loss += loss.item() * y.size(0)
test_acc += (output.max(1)[1] == y).sum().item()
test_n += y.size(0)
attack_mean += torch.mean(torch.abs(delta)) * y.size(0)
test_time = time.time()
if not args.eval:
logger.info(
'%d \t %.1f \t \t %.1f \t \t %.4f \t %.4f \t \t %.4f \t \t %.4f \t %.4f \t %.4f \t \t %.4f \t \t %.4f \t %.4f',
epoch, train_time - start_time, test_time - train_time, lr,
train_robust_loss / train_n, train_robust_acc / train_n,
test_loss / test_n, test_acc / test_n,
test_robust_loss / test_n, test_robust_acc / test_n,
defence_mean * 255 / train_n, attack_mean * 255 / test_n)
# save checkpoint
if (epoch + 1) % args.chkpt_iters == 0 or epoch + 1 == epochs:
torch.save(model.state_dict(),
os.path.join(args.fname, f'model_{epoch}.pth'))
torch.save(opt.state_dict(),
os.path.join(args.fname, f'opt_{epoch}.pth'))
# save best
if test_robust_acc / test_n > best_test_robust_acc:
torch.save(
{
'state_dict': model.state_dict(),
'test_robust_acc': test_robust_acc / test_n,
'test_robust_loss': test_robust_loss / test_n,
'test_loss': test_loss / test_n,
'test_acc': test_acc / test_n,
}, os.path.join(args.fname, f'model_best.pth'))
best_test_robust_acc = test_robust_acc / test_n
else:
logger.info(
'%d \t %.1f \t \t %.1f \t \t %.4f \t %.4f \t \t %.4f \t \t %.4f \t %.4f \t %.4f \t \t %.4f \t \t %.4f \t %.4f',
epoch, train_time - start_time, test_time - train_time, -1, -1,
-1, test_loss / test_n, test_acc / test_n,
test_robust_loss / test_n, test_robust_acc / test_n, -1,
attack_mean * 255 / test_n)
if args.eval or epoch + 1 == epochs:
start_test_time = time.time()
test_loss = 0
test_acc = 0
test_robust_loss = 0
test_robust_acc = 0
test_n = 0
for i, batch in enumerate(test_batches):
X, y = batch['input'], batch['target']
delta = attack_pgd(model,
X,
y,
epsilon,
pgd_alpha,
50,
10,
'l_inf',
early_stop=True)
delta = delta.detach()
robust_output = model(
normalize(
torch.clamp(X + delta[:X.size(0)],
min=lower_limit,
max=upper_limit)))
robust_loss = criterion(robust_output, y)
output = model(normalize(X))
loss = criterion(output, y)
test_robust_loss += robust_loss.item() * y.size(0)
test_robust_acc += (robust_output.max(1)[1] == y).sum().item()
test_loss += loss.item() * y.size(0)
test_acc += (output.max(1)[1] == y).sum().item()
test_n += y.size(0)
logger.info(
'PGD50 \t time: %.1f,\t clean loss: %.4f,\t clean acc: %.4f,\t robust loss: %.4f,\t robust acc: %.4f',
time.time() - start_test_time, test_loss / test_n,
test_acc / test_n, test_robust_loss / test_n,
test_robust_acc / test_n)
return
def main():
global args, best_prec1, exp_dir
best_prec1 = 0
args = parser.parse_args()
print(args.lr_decay_at)
assert args.normalization in ['GCN_ZCA',
'GCN'], 'normalization {} unknown'.format(
args.normalization)
global zca
if 'ZCA' in args.normalization:
zca_params = torch.load('./data/cifar-10-batches-py/zca_params.pth')
zca = ZCA(zca_params)
else:
zca = None
exp_dir = os.path.join('experiments', args.name)
if os.path.exists(exp_dir):
print("same experiment exist...")
#return
else:
os.makedirs(exp_dir)
# DATA SETTINGS
global dataset_train, dataset_test
if args.dataset == 'cifar10':
import cifar
dataset_train = cifar.CIFAR10(args, train=True)
dataset_test = cifar.CIFAR10(args, train=False)
if args.UDA:
# loader for UDA
dataset_train_uda = cifar.CIFAR10(args, True, True)
uda_loader = torch.utils.data.DataLoader(
dataset_train_uda,
batch_size=args.batch_size_unsup,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
iter_uda = iter(uda_loader)
else:
iter_uda = None
train_loader, test_loader = initialize_loader()
# MODEL SETTINGS
if args.arch == 'WRN-28-2':
model = WideResNet(28, [100, 10][int(args.dataset == 'cifar10')],
2,
dropRate=args.dropout_rate)
model = torch.nn.DataParallel(model.cuda())
else:
raise NotImplementedError('arch {} is not implemented'.format(
args.arch))
if args.optimizer == 'Adam':
print("use Adam optimizer")
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=args.l2_reg)
elif args.optimizer == 'SGD':
print("use SGD optimizer")
optimizer = torch.optim.SGD(model.parameters(),
lr=args.lr,
momentum=0.9,
weight_decay=args.l2_reg,
nesterov=args.nesterov)
if args.lr_decay == 'cosine':
print("use cosine lr scheduler")
global scheduler
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer, args.max_iter, eta_min=args.final_lr)
global batch_time, losses_sup, losses_unsup, top1, losses_l1, losses_unsup
batch_time, losses_sup, losses_unsup, top1, losses_l1, losses_unsup = AverageMeter(
), AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter(
), AverageMeter()
t = time.time()
model.train()
iter_sup = iter(train_loader)
for train_iter in range(args.max_iter):
# TRAIN
lr = adjust_learning_rate(optimizer, train_iter + 1)
train(model,
iter_sup,
optimizer,
train_iter,
data_iterator_uda=iter_uda)
# LOGGING
if (train_iter + 1) % args.print_freq == 0:
print('ITER: [{0}/{1}]\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
'L1 Loss {l1_loss.val:.4f} ({l1_loss.avg:.4f})\t'
'Unsup Loss {unsup_loss.val:.4f} ({unsup_loss.avg:.4f})\t'
'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
'Learning rate {2} TSA th {3}'.format(
train_iter,
args.max_iter,
lr,
TSA_th(train_iter),
batch_time=batch_time,
loss=losses_sup,
l1_loss=losses_l1,
unsup_loss=losses_unsup,
top1=top1))
if (train_iter +
1) % args.eval_iter == 0 or train_iter + 1 == args.max_iter:
# EVAL
print("evaluation at iter {}".format(train_iter))
prec1 = test(model, test_loader)
is_best = prec1 > best_prec1
best_prec1 = max(prec1, best_prec1)
save_checkpoint(
{
'iter': train_iter + 1,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
}, is_best)
print("* Best accuracy: {}".format(best_prec1))
eval_interval_time = time.time() - t
t = time.time()
print("total {} sec for {} iterations".format(
eval_interval_time, args.eval_iter))
seconds_remaining = eval_interval_time / float(
args.eval_iter) * (args.max_iter - train_iter)
print("{}:{}:{} remaining".format(
int(seconds_remaining // 3600),
int((seconds_remaining % 3600) // 60),
int(seconds_remaining % 60)))
model.train()
batch_time, losses_sup, losses_unsup, top1, losses_l1, losses_unsup = AverageMeter(
), AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter(
), AverageMeter()
iter_sup = iter(train_loader)
if iter_uda is not None:
iter_uda = iter(uda_loader)
def main():
args = get_args()
if args.fname == 'auto':
names = get_auto_fname(args)
args.fname = 'noise_predictor/' + names
else:
args.fname = 'noise_predictor/' + args.fname
if not os.path.exists(args.fname):
os.makedirs(args.fname)
logger = logging.getLogger(__name__)
logging.basicConfig(
format='[%(asctime)s] - %(message)s',
datefmt='%Y/%m/%d %H:%M:%S',
level=logging.DEBUG,
handlers=[
logging.FileHandler(
os.path.join(args.fname,
'eval.log' if args.eval else 'output.log')),
logging.StreamHandler()
])
logger.info(args)
# Set seed
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
shuffle = True
train_adv_images = None
train_adv_labels = None
test_adv_images = None
test_adv_labels = None
print("Attacks")
attacks = args.list.split("_")
print(attacks)
for i in range(len(attacks)):
_adv_dir = "adv_examples/{}/".format(attacks[i])
train_path = _adv_dir + "train.pth"
test_path = _adv_dir + "test.pth"
adv_train_data = torch.load(train_path)
adv_test_data = torch.load(test_path)
if i == 0:
train_adv_images = adv_train_data["adv"]
test_adv_images = adv_test_data["adv"]
train_adv_labels = [i] * len(adv_train_data["label"])
test_adv_labels = [i] * len(adv_test_data["label"])
else:
train_adv_images = np.concatenate(
(train_adv_images, adv_train_data["adv"]))
test_adv_images = np.concatenate(
(test_adv_images, adv_test_data["adv"]))
train_adv_labels = np.concatenate(
(train_adv_labels, [i] * len(adv_train_data["label"])))
test_adv_labels = np.concatenate(
(test_adv_labels, [i] * len(adv_test_data["label"])))
train_adv_set = list(zip(train_adv_images, train_adv_labels))
print("")
print("Train Adv Attack Data: ", attacks)
print("Len: ", len(train_adv_set))
print("")
train_adv_batches = Batches(train_adv_set,
args.batch_size,
shuffle=shuffle,
set_random_choices=False,
num_workers=4)
test_adv_set = list(zip(test_adv_images, test_adv_labels))
test_adv_batches = Batches(test_adv_set,
args.batch_size,
shuffle=False,
num_workers=4)
# Set perturbations
epsilon = (args.epsilon / 255.)
test_epsilon = (args.test_epsilon / 255.)
pgd_alpha = (args.pgd_alpha / 255.)
test_pgd_alpha = (args.test_pgd_alpha / 255.)
# Set models
model = None
if args.model == "resnet18":
model = resnet18(num_classes=args.num_classes)
elif args.model == "resnet20":
model = resnet20()
elif args.model == "vgg16bn":
model = vgg16_bn(num_classes=args.num_classes)
elif args.model == "densenet121":
model = densenet121(pretrained=True)
elif args.model == "googlenet":
model = googlenet(pretrained=True)
elif args.model == "inceptionv3":
model = inception_v3(num_classes=3)
elif args.model == 'WideResNet':
model = WideResNet(34,
10,
widen_factor=10,
dropRate=0.0,
normalize=args.use_FNandWN,
activation=args.activation,
softplus_beta=args.softplus_beta)
elif args.model == 'WideResNet_20':
model = WideResNet(34,
10,
widen_factor=20,
dropRate=0.0,
normalize=args.use_FNandWN,
activation=args.activation,
softplus_beta=args.softplus_beta)
else:
raise ValueError("Unknown model")
# Set training hyperparameters
if args.l2:
decay, no_decay = [], []
for name, param in model.named_parameters():
if 'bn' not in name and 'bias' not in name:
decay.append(param)
else:
no_decay.append(param)
params = [{
'params': decay,
'weight_decay': args.l2
}, {
'params': no_decay,
'weight_decay': 0
}]
else:
params = model.parameters()
if args.lr_schedule == 'cyclic':
opt = torch.optim.Adam(params,
lr=args.lr_max,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=args.weight_decay)
else:
if args.optimizer == 'momentum':
opt = torch.optim.SGD(params,
lr=args.lr_max,
momentum=0.9,
weight_decay=args.weight_decay)
elif args.optimizer == 'Nesterov':
opt = torch.optim.SGD(params,
lr=args.lr_max,
momentum=0.9,
weight_decay=args.weight_decay,
nesterov=True)
elif args.optimizer == 'SGD_GC':
opt = SGD_GC(params,
lr=args.lr_max,
momentum=0.9,
weight_decay=args.weight_decay)
elif args.optimizer == 'SGD_GCC':
opt = SGD_GCC(params,
lr=args.lr_max,
momentum=0.9,
weight_decay=args.weight_decay)
elif args.optimizer == 'Adam':
opt = torch.optim.Adam(params,
lr=args.lr_max,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=args.weight_decay)
elif args.optimizer == 'AdamW':
opt = torch.optim.AdamW(params,
lr=args.lr_max,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=args.weight_decay)
# Cross-entropy (mean)
if args.labelsmooth:
criterion = LabelSmoothingLoss(smoothing=args.labelsmoothvalue)
else:
criterion = nn.CrossEntropyLoss()
epochs = args.epochs
# Set lr schedule
if args.lr_schedule == 'superconverge':
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, warm_up_lr=args.warmup_lr):
if t < 100:
if warm_up_lr and t < args.warmup_lr_epoch:
return (t + 1.) / args.warmup_lr_epoch * args.lr_max
else:
return args.lr_max
if args.lrdecay == 'lineardecay':
if t < 105:
return args.lr_max * 0.02 * (105 - t)
else:
return 0.
elif args.lrdecay == 'intenselr':
if t < 102:
return args.lr_max / 10.
else:
return args.lr_max / 100.
elif args.lrdecay == 'looselr':
if t < 150:
return args.lr_max / 10.
else:
return args.lr_max / 100.
elif args.lrdecay == 'base':
if t < 105:
return args.lr_max / 10.
else:
return args.lr_max / 100.
elif args.lr_schedule == 'linear':
lr_schedule = lambda t: np.interp([t], [
0, args.epochs // 3, args.epochs * 2 // 3, args.epochs
], [args.lr_max, args.lr_max, args.lr_max / 10, args.lr_max / 100])[0]
elif args.lr_schedule == 'onedrop':
def lr_schedule(t):
if t < args.lr_drop_epoch:
return args.lr_max
else:
return args.lr_one_drop
elif args.lr_schedule == 'multipledecay':
def lr_schedule(t):
return args.lr_max - (t //
(args.epochs // 10)) * (args.lr_max / 10)
elif args.lr_schedule == 'cosine':
def lr_schedule(t):
return args.lr_max * 0.5 * (1 + np.cos(t / args.epochs * np.pi))
elif args.lr_schedule == 'cyclic':
def lr_schedule(t, stepsize=18, min_lr=1e-5, max_lr=args.lr_max):
# Scaler: we can adapt this if we do not want the triangular CLR
scaler = lambda x: 1.
# Additional function to see where on the cycle we are
cycle = math.floor(1 + t / (2 * stepsize))
x = abs(t / stepsize - 2 * cycle + 1)
relative = max(0, (1 - x)) * scaler(cycle)
return min_lr + (max_lr - min_lr) * relative
#### Set stronger adv attacks when decay the lr ####
def eps_alpha_schedule(
t,
warm_up_eps=args.warmup_eps,
if_use_stronger_adv=args.use_stronger_adv,
stronger_index=args.stronger_index): # Schedule number 0
if stronger_index == 0:
epsilon_s = [epsilon * 1.5, epsilon * 2]
pgd_alpha_s = [pgd_alpha, pgd_alpha]
elif stronger_index == 1:
epsilon_s = [epsilon * 1.5, epsilon * 2]
pgd_alpha_s = [pgd_alpha * 1.25, pgd_alpha * 1.5]
elif stronger_index == 2:
epsilon_s = [epsilon * 2, epsilon * 2.5]
pgd_alpha_s = [pgd_alpha * 1.5, pgd_alpha * 2]
else:
print('Undefined stronger index')
if if_use_stronger_adv:
if t < 100:
if t < args.warmup_eps_epoch and warm_up_eps:
return (
t + 1.
) / args.warmup_eps_epoch * epsilon, pgd_alpha, args.restarts
else:
return epsilon, pgd_alpha, args.restarts
elif t < 105:
return epsilon_s[0], pgd_alpha_s[0], args.restarts
else:
return epsilon_s[1], pgd_alpha_s[1], args.restarts
else:
if t < args.warmup_eps_epoch and warm_up_eps:
return (
t + 1.
) / args.warmup_eps_epoch * epsilon, pgd_alpha, args.restarts
else:
return epsilon, pgd_alpha, args.restarts
#### Set the counter for the early stop of PGD ####
def early_stop_counter_schedule(t):
if t < args.earlystopPGDepoch1:
return 1
elif t < args.earlystopPGDepoch2:
return 2
else:
return 3
best_test_adv_acc = 0
best_val_robust_acc = 0
if args.resume:
start_epoch = args.resume
model.load_state_dict(
torch.load(os.path.join(args.fname, f'model_{start_epoch-1}.pth')))
opt.load_state_dict(
torch.load(os.path.join(args.fname, f'opt_{start_epoch-1}.pth')))
logger.info(f'Resuming at epoch {start_epoch}')
best_test_adv_acc = torch.load(
os.path.join(args.fname, f'model_best.pth'))['test_adv_acc']
if args.val:
best_val_robust_acc = torch.load(
os.path.join(args.fname, f'model_val.pth'))['val_robust_acc']
else:
start_epoch = 0
if args.eval:
if not args.resume:
logger.info(
"No model loaded to evaluate, specify with --resume FNAME")
return
logger.info("[Evaluation mode]")
model.cuda()
model.train()
logger.info('Epoch \t Train Robust Acc \t Test Robust Acc')
# Records per epoch for savetxt
train_loss_record = []
train_acc_record = []
train_robust_loss_record = []
train_robust_acc_record = []
train_grad_record = []
test_loss_record = []
test_acc_record = []
test_adv_loss_record = []
test_adv_acc_record = []
test_grad_record = []
for epoch in range(start_epoch, epochs):
model.train()
start_time = time.time()
train_robust_loss = 0
train_robust_acc = 0
train_n = 0
train_grad = 0
record_iter = torch.tensor([])
for i, adv_batch in enumerate(train_adv_batches):
if args.eval:
break
adv_input = normalize(adv_batch['input'])
adv_y = adv_batch['target']
adv_input.requires_grad = True
robust_output = model(adv_input)
robust_loss = criterion(robust_output, adv_y)
if args.l1:
for name, param in model.named_parameters():
if 'bn' not in name and 'bias' not in name:
robust_loss += args.l1 * param.abs().sum()
opt.zero_grad()
robust_loss.backward()
opt.step()
# Record the statstic values
train_robust_loss += robust_loss.item() * adv_y.size(0)
train_robust_acc += (robust_output.max(1)[1] == adv_y).sum().item()
train_n += adv_y.size(0)
train_time = time.time()
if args.earlystopPGD:
print('Iter mean: ',
record_iter.mean().item(), ' Iter std: ',
record_iter.std().item())
# Evaluate on test data
model.eval()
test_adv_loss = 0
test_adv_acc = 0
test_adv_n = 0
for i, batch in enumerate(test_adv_batches):
adv_input = normalize(batch['input'])
y = batch['target']
robust_output = model(adv_input)
robust_loss = criterion(robust_output, y)
test_adv_loss += robust_loss.item() * y.size(0)
test_adv_acc += (robust_output.max(1)[1] == y).sum().item()
test_adv_n += y.size(0)
test_time = time.time()
logger.info('%d \t %.4f \t\t %.4f', epoch + 1,
train_robust_acc / train_n, test_adv_acc / test_adv_n)
# Save results
train_robust_loss_record.append(train_robust_loss / train_n)
train_robust_acc_record.append(train_robust_acc / train_n)
np.savetxt(args.fname + '/train_robust_loss_record.txt',
np.array(train_robust_loss_record))
np.savetxt(args.fname + '/train_robust_acc_record.txt',
np.array(train_robust_acc_record))
test_adv_loss_record.append(test_adv_loss / train_n)
test_adv_acc_record.append(test_adv_acc / train_n)
np.savetxt(args.fname + '/test_adv_loss_record.txt',
np.array(test_adv_loss_record))
np.savetxt(args.fname + '/test_adv_acc_record.txt',
np.array(test_adv_acc_record))
# save checkpoint
if epoch > 99 or (epoch +
1) % args.chkpt_iters == 0 or epoch + 1 == epochs:
torch.save(model.state_dict(),
os.path.join(args.fname, f'model_{epoch}.pth'))
torch.save(opt.state_dict(),
os.path.join(args.fname, f'opt_{epoch}.pth'))
# save best
if test_adv_acc / test_adv_n > best_test_adv_acc:
torch.save(
{
'state_dict': model.state_dict(),
'test_adv_acc': test_adv_acc / test_adv_n,
'test_adv_loss': test_adv_loss / test_adv_n,
}, os.path.join(args.fname, f'model_best.pth'))
best_test_adv_acc = test_adv_acc / test_adv_n
def main():
args = get_args()
if args.fname == 'auto':
names = get_auto_fname(args)
args.fname = '../../trained_models/' + names
else:
args.fname = '../../trained_models/' + args.fname
if not os.path.exists(args.fname):
os.makedirs(args.fname)
logger = logging.getLogger(__name__)
logging.basicConfig(
format='[%(asctime)s] - %(message)s',
datefmt='%Y/%m/%d %H:%M:%S',
level=logging.DEBUG,
handlers=[
logging.FileHandler(
os.path.join(args.fname,
'eval.log' if args.eval else 'output.log')),
logging.StreamHandler()
])
logger.info(args)
# Set seed
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
# setup data loader
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
])
transform_test = transforms.Compose([
transforms.ToTensor(),
])
train_set = torchvision.datasets.CIFAR10(root='../data',
train=True,
download=True,
transform=transform_train)
test_set = torchvision.datasets.CIFAR10(root='../data',
train=False,
download=True,
transform=transform_test)
if args.attack == "all":
train_data = np.array(train_set.data) / 255.
train_data = transpose(train_data).astype(np.float32)
train_labels = np.array(train_set.targets)
oversampled_train_data = np.tile(train_data, (11, 1, 1, 1))
oversampled_train_labels = np.tile(train_labels, (11))
train_set = list(
zip(torch.from_numpy(oversampled_train_data),
torch.from_numpy(oversampled_train_labels)))
elif args.attack == "combine":
train_data = np.array(train_set.data) / 255.
train_data = transpose(train_data).astype(np.float32)
train_labels = np.array(train_set.targets)
oversampled_train_data = train_data.copy()
oversampled_train_labels = train_labels.copy()
logger.info("Attacks")
attacks = args.list.split("_")
logger.info(attacks)
oversampled_train_data = np.tile(train_data, (len(attacks), 1, 1, 1))
oversampled_train_labels = np.tile(train_labels, (len(attacks)))
train_set = list(
zip(torch.from_numpy(oversampled_train_data),
torch.from_numpy(oversampled_train_labels)))
else:
train_data = np.array(train_set.data) / 255.
train_data = transpose(train_data).astype(np.float32)
train_labels = np.array(train_set.targets)
train_set = list(
zip(torch.from_numpy(train_data), torch.from_numpy(train_labels)))
test_data = np.array(test_set.data) / 255.
test_data = transpose(test_data).astype(np.float32)
test_labels = np.array(test_set.targets)
test_set = list(
zip(torch.from_numpy(test_data), torch.from_numpy(test_labels)))
if args.sample != 100:
n = len(train_set)
n_sample = int(n * args.sample / 100)
np.random.shuffle(train_set)
train_set = train_set[:n_sample]
print("")
print("Train Original Data: ")
print("Len: ", len(train_set))
print("")
shuffle = False
train_batches = Batches(train_set, args.batch_size, shuffle=shuffle)
test_batches = Batches(test_set, args.batch_size, shuffle=False)
train_adv_images = None
train_adv_labels = None
test_adv_images = None
test_adv_labels = None
adv_dir = "adv_examples/{}/".format(args.attack)
train_path = adv_dir + "train.pth"
test_path = adv_dir + "test.pth"
# ATTACK_LIST = ["autoattack", "autopgd", "bim", "cw", "deepfool", "fgsm", "newtonfool", "pgd", "pixelattack", "spatialtransformation", "squareattack"]
ATTACK_LIST = [
"pixelattack", "spatialtransformation", "squareattack", "fgsm",
"deepfool", "bim", "cw", "pgd", "autoattack", "autopgd", "newtonfool"
]
if args.attack in TOOLBOX_ADV_ATTACK_LIST:
adv_train_data = torch.load(train_path)
train_adv_images = adv_train_data["adv"]
train_adv_labels = adv_train_data["label"]
adv_test_data = torch.load(test_path)
test_adv_images = adv_test_data["adv"]
test_adv_labels = adv_test_data["label"]
elif args.attack in ["ffgsm", "mifgsm", "tpgd"]:
adv_data = {}
adv_data["adv"], adv_data["label"] = torch.load(train_path)
train_adv_images = adv_data["adv"].numpy()
train_adv_labels = adv_data["label"].numpy()
adv_data = {}
adv_data["adv"], adv_data["label"] = torch.load(test_path)
test_adv_images = adv_data["adv"].numpy()
test_adv_labels = adv_data["label"].numpy()
elif args.attack == "all":
for i in range(len(ATTACK_LIST)):
_adv_dir = "adv_examples/{}/".format(ATTACK_LIST[i])
train_path = _adv_dir + "train.pth"
test_path = _adv_dir + "test.pth"
adv_train_data = torch.load(train_path)
adv_test_data = torch.load(test_path)
if i == 0:
train_adv_images = adv_train_data["adv"]
train_adv_labels = adv_train_data["label"]
test_adv_images = adv_test_data["adv"]
test_adv_labels = adv_test_data["label"]
else:
# print(train_adv_images.shape)
# print(adv_train_data["adv"].shape)
train_adv_images = np.concatenate(
(train_adv_images, adv_train_data["adv"]))
train_adv_labels = np.concatenate(
(train_adv_labels, adv_train_data["label"]))
test_adv_images = np.concatenate(
(test_adv_images, adv_test_data["adv"]))
test_adv_labels = np.concatenate(
(test_adv_labels, adv_test_data["label"]))
elif args.attack == "combine":
print("Attacks")
attacks = args.list.split("_")
print(attacks)
if args.balanced == None:
for i in range(len(attacks)):
_adv_dir = "adv_examples/{}/".format(attacks[i])
train_path = _adv_dir + "train.pth"
test_path = _adv_dir + "test.pth"
adv_train_data = torch.load(train_path)
adv_test_data = torch.load(test_path)
if i == 0:
train_adv_images = adv_train_data["adv"]
train_adv_labels = adv_train_data["label"]
test_adv_images = adv_test_data["adv"]
test_adv_labels = adv_test_data["label"]
else:
# print(train_adv_images.shape)
# print(adv_train_data["adv"].shape)
train_adv_images = np.concatenate(
(train_adv_images, adv_train_data["adv"]))
train_adv_labels = np.concatenate(
(train_adv_labels, adv_train_data["label"]))
test_adv_images = np.concatenate(
(test_adv_images, adv_test_data["adv"]))
test_adv_labels = np.concatenate(
(test_adv_labels, adv_test_data["label"]))
else:
proportion_str = args.balanced.split("_")
proportion = [int(x) for x in proportion_str]
sum_proportion = sum(proportion)
proportion = [float(x) / float(sum_proportion) for x in proportion]
sum_samples = 0
for i in range(len(attacks)):
_adv_dir = "adv_examples/{}/".format(attacks[i])
train_path = _adv_dir + "train.pth"
test_path = _adv_dir + "test.pth"
adv_train_data = torch.load(train_path)
adv_test_data = torch.load(test_path)
random_state = 0
num_samples = 0
total = 50000
if i != len(attacks) - 1:
n_samples = int(proportion[i] * total)
sum_samples += n_samples
else:
n_samples = total - sum_samples
print("Sample")
print(n_samples)
if i == 0:
train_adv_images = resample(adv_train_data["adv"],
n_samples=n_samples,
random_state=random_state)
train_adv_labels = resample(adv_train_data["label"],
n_samples=n_samples,
random_state=random_state)
test_adv_images = resample(adv_test_data["adv"],
n_samples=n_samples,
random_state=random_state)
test_adv_labels = resample(adv_test_data["label"],
n_samples=n_samples,
random_state=random_state)
else:
train_adv_images = np.concatenate(
(train_adv_images,
resample(adv_train_data["adv"],
n_samples=n_samples,
random_state=random_state)))
train_adv_labels = np.concatenate(
(train_adv_labels,
resample(adv_train_data["label"],
n_samples=n_samples,
random_state=random_state)))
test_adv_images = np.concatenate(
(test_adv_images,
resample(adv_test_data["adv"],
n_samples=n_samples,
random_state=random_state)))
test_adv_labels = np.concatenate(
(test_adv_labels,
resample(adv_test_data["label"],
n_samples=n_samples,
random_state=random_state)))
else:
raise ValueError("Unknown adversarial data")
train_adv_set = list(zip(train_adv_images, train_adv_labels))
if args.sample != 100:
n = len(train_adv_set)
n_sample = int(n * args.sample / 100)
np.random.shuffle(train_adv_set)
train_adv_set = train_adv_set[:n_sample]
print("")
print("Train Adv Attack Data: ", args.attack)
print("Len: ", len(train_adv_set))
print("")
train_adv_batches = Batches(train_adv_set,
args.batch_size,
shuffle=shuffle,
set_random_choices=False,
num_workers=4)
test_adv_set = list(zip(test_adv_images, test_adv_labels))
test_adv_batches = Batches(test_adv_set,
args.batch_size,
shuffle=False,
num_workers=4)
# Set perturbations
epsilon = (args.epsilon / 255.)
test_epsilon = (args.test_epsilon / 255.)
pgd_alpha = (args.pgd_alpha / 255.)
test_pgd_alpha = (args.test_pgd_alpha / 255.)
# Set models
model = None
if args.model == "resnet18":
model = resnet18(pretrained=True)
elif args.model == "resnet20":
model = resnet20()
elif args.model == "vgg16bn":
model = vgg16_bn(pretrained=True)
elif args.model == "densenet121":
model = densenet121(pretrained=True)
elif args.model == "googlenet":
model = googlenet(pretrained=True)
elif args.model == "inceptionv3":
model = inception_v3(pretrained=True)
elif args.model == 'WideResNet':
model = WideResNet(34,
10,
widen_factor=10,
dropRate=0.0,
normalize=args.use_FNandWN,
activation=args.activation,
softplus_beta=args.softplus_beta)
elif args.model == 'WideResNet_20':
model = WideResNet(34,
10,
widen_factor=20,
dropRate=0.0,
normalize=args.use_FNandWN,
activation=args.activation,
softplus_beta=args.softplus_beta)
else:
raise ValueError("Unknown model")
# Set training hyperparameters
if args.l2:
decay, no_decay = [], []
for name, param in model.named_parameters():
if 'bn' not in name and 'bias' not in name:
decay.append(param)
else:
no_decay.append(param)
params = [{
'params': decay,
'weight_decay': args.l2
}, {
'params': no_decay,
'weight_decay': 0
}]
else:
params = model.parameters()
if args.lr_schedule == 'cyclic':
opt = torch.optim.Adam(params,
lr=args.lr_max,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=args.weight_decay)
else:
if args.optimizer == 'momentum':
opt = torch.optim.SGD(params,
lr=args.lr_max,
momentum=0.9,
weight_decay=args.weight_decay)
elif args.optimizer == 'Nesterov':
opt = torch.optim.SGD(params,
lr=args.lr_max,
momentum=0.9,
weight_decay=args.weight_decay,
nesterov=True)
elif args.optimizer == 'SGD_GC':
opt = SGD_GC(params,
lr=args.lr_max,
momentum=0.9,
weight_decay=args.weight_decay)
elif args.optimizer == 'SGD_GCC':
opt = SGD_GCC(params,
lr=args.lr_max,
momentum=0.9,
weight_decay=args.weight_decay)
elif args.optimizer == 'Adam':
opt = torch.optim.Adam(params,
lr=args.lr_max,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=args.weight_decay)
elif args.optimizer == 'AdamW':
opt = torch.optim.AdamW(params,
lr=args.lr_max,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=args.weight_decay)
# Cross-entropy (mean)
if args.labelsmooth:
criterion = LabelSmoothingLoss(smoothing=args.labelsmoothvalue)
else:
criterion = nn.CrossEntropyLoss()
# If we use freeAT or fastAT with previous init
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':
epochs = int(math.ceil(args.epochs / args.attack_iters))
else:
epochs = args.epochs
# Set lr schedule
if args.lr_schedule == 'superconverge':
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, warm_up_lr=args.warmup_lr):
if t < 100:
if warm_up_lr and t < args.warmup_lr_epoch:
return (t + 1.) / args.warmup_lr_epoch * args.lr_max
else:
return args.lr_max
if args.lrdecay == 'lineardecay':
if t < 105:
return args.lr_max * 0.02 * (105 - t)
else:
return 0.
elif args.lrdecay == 'intenselr':
if t < 102:
return args.lr_max / 10.
else:
return args.lr_max / 100.
elif args.lrdecay == 'looselr':
if t < 150:
return args.lr_max / 10.
else:
return args.lr_max / 100.
elif args.lrdecay == 'base':
if t < 105:
return args.lr_max / 10.
else:
return args.lr_max / 100.
elif args.lr_schedule == 'linear':
lr_schedule = lambda t: np.interp([t], [
0, args.epochs // 3, args.epochs * 2 // 3, args.epochs
], [args.lr_max, args.lr_max, args.lr_max / 10, args.lr_max / 100])[0]
elif args.lr_schedule == 'onedrop':
def lr_schedule(t):
if t < args.lr_drop_epoch:
return args.lr_max
else:
return args.lr_one_drop
elif args.lr_schedule == 'multipledecay':
def lr_schedule(t):
return args.lr_max - (t //
(args.epochs // 10)) * (args.lr_max / 10)
elif args.lr_schedule == 'cosine':
def lr_schedule(t):
return args.lr_max * 0.5 * (1 + np.cos(t / args.epochs * np.pi))
elif args.lr_schedule == 'cyclic':
def lr_schedule(t, stepsize=18, min_lr=1e-5, max_lr=args.lr_max):
# Scaler: we can adapt this if we do not want the triangular CLR
scaler = lambda x: 1.
# Additional function to see where on the cycle we are
cycle = math.floor(1 + t / (2 * stepsize))
x = abs(t / stepsize - 2 * cycle + 1)
relative = max(0, (1 - x)) * scaler(cycle)
return min_lr + (max_lr - min_lr) * relative
#### Set stronger adv attacks when decay the lr ####
def eps_alpha_schedule(
t,
warm_up_eps=args.warmup_eps,
if_use_stronger_adv=args.use_stronger_adv,
stronger_index=args.stronger_index): # Schedule number 0
if stronger_index == 0:
epsilon_s = [epsilon * 1.5, epsilon * 2]
pgd_alpha_s = [pgd_alpha, pgd_alpha]
elif stronger_index == 1:
epsilon_s = [epsilon * 1.5, epsilon * 2]
pgd_alpha_s = [pgd_alpha * 1.25, pgd_alpha * 1.5]
elif stronger_index == 2:
epsilon_s = [epsilon * 2, epsilon * 2.5]
pgd_alpha_s = [pgd_alpha * 1.5, pgd_alpha * 2]
else:
print('Undefined stronger index')
if if_use_stronger_adv:
if t < 100:
if t < args.warmup_eps_epoch and warm_up_eps:
return (
t + 1.
) / args.warmup_eps_epoch * epsilon, pgd_alpha, args.restarts
else:
return epsilon, pgd_alpha, args.restarts
elif t < 105:
return epsilon_s[0], pgd_alpha_s[0], args.restarts
else:
return epsilon_s[1], pgd_alpha_s[1], args.restarts
else:
if t < args.warmup_eps_epoch and warm_up_eps:
return (
t + 1.
) / args.warmup_eps_epoch * epsilon, pgd_alpha, args.restarts
else:
return epsilon, pgd_alpha, args.restarts
#### Set the counter for the early stop of PGD ####
def early_stop_counter_schedule(t):
if t < args.earlystopPGDepoch1:
return 1
elif t < args.earlystopPGDepoch2:
return 2
else:
return 3
best_test_adv_acc = 0
best_val_robust_acc = 0
if args.resume:
start_epoch = args.resume
model.load_state_dict(
torch.load(os.path.join(args.fname, f'model_{start_epoch-1}.pth')))
opt.load_state_dict(
torch.load(os.path.join(args.fname, f'opt_{start_epoch-1}.pth')))
logger.info(f'Resuming at epoch {start_epoch}')
best_test_adv_acc = torch.load(
os.path.join(args.fname, f'model_best.pth'))['test_adv_acc']
if args.val:
best_val_robust_acc = torch.load(
os.path.join(args.fname, f'model_val.pth'))['val_robust_acc']
else:
start_epoch = 0
if args.eval:
if not args.resume:
logger.info(
"No model loaded to evaluate, specify with --resume FNAME")
return
logger.info("[Evaluation mode]")
model.cuda()
model.eval()
# Evaluate on original test data
test_acc = 0
test_n = 0
for i, batch in enumerate(test_batches):
X, y = batch['input'], batch['target']
clean_input = normalize(X)
output = model(clean_input)
test_acc += (output.max(1)[1] == y).sum().item()
test_n += y.size(0)
logger.info('Intial Accuracy on Original Test Data: %.4f (Test Acc)',
test_acc / test_n)
test_adv_acc = 0
test_adv_n = 0
for i, batch in enumerate(test_adv_batches):
adv_input = normalize(batch['input'])
y = batch['target']
robust_output = model(adv_input)
test_adv_acc += (robust_output.max(1)[1] == y).sum().item()
test_adv_n += y.size(0)
logger.info(
'Intial Accuracy on Adversarial Test Data: %.4f (Test Robust Acc)',
test_adv_acc / test_adv_n)
model.train()
logger.info(
'Epoch \t Train Acc \t Train Robust Acc \t Test Acc \t Test Robust Acc'
)
# Records per epoch for savetxt
train_loss_record = []
train_acc_record = []
train_robust_loss_record = []
train_robust_acc_record = []
train_grad_record = []
test_loss_record = []
test_acc_record = []
test_adv_loss_record = []
test_adv_acc_record = []
test_grad_record = []
for epoch in range(start_epoch, epochs):
model.train()
start_time = time.time()
train_loss = 0
train_acc = 0
train_robust_loss = 0
train_robust_acc = 0
train_n = 0
train_grad = 0
record_iter = torch.tensor([])
for i, (batch,
adv_batch) in enumerate(zip(train_batches, train_adv_batches)):
if args.eval:
break
X, y = batch['input'], batch['target']
adv_input = normalize(adv_batch['input'])
adv_y = adv_batch['target']
adv_input.requires_grad = True
robust_output = model(adv_input)
# Training losses
if args.mixup:
clean_input = normalize(X)
clean_input.requires_grad = True
output = model(clean_input)
robust_loss = mixup_criterion(criterion, robust_output, y_a,
y_b, lam)
elif args.mixture:
clean_input = normalize(X)
clean_input.requires_grad = True
output = model(clean_input)
robust_loss = args.mixture_alpha * criterion(
robust_output,
adv_y) + (1 - args.mixture_alpha) * criterion(output, y)
else:
clean_input = normalize(X)
clean_input.requires_grad = True
output = model(clean_input)
if args.focalloss:
criterion_nonreduct = nn.CrossEntropyLoss(reduction='none')
robust_confidence = F.softmax(robust_output,
dim=1)[:, adv_y].detach()
robust_loss = (criterion_nonreduct(robust_output, adv_y) *
((1. - robust_confidence)**
args.focallosslambda)).mean()
elif args.use_DLRloss:
beta_ = 0.8 * epoch_now / args.epochs
robust_loss = (1. - beta_) * F.cross_entropy(
robust_output, adv_y) + beta_ * dlr_loss(
robust_output, adv_y)
elif args.use_CWloss:
beta_ = 0.8 * epoch_now / args.epochs
robust_loss = (1. - beta_) * F.cross_entropy(
robust_output, adv_y) + beta_ * CW_loss(
robust_output, adv_y)
elif args.use_FNandWN:
#print('use FN and WN with margin')
robust_loss = criterion(
args.s_FN * robust_output -
onehot_target_withmargin_HE, adv_y)
else:
robust_loss = criterion(robust_output, adv_y)
if args.l1:
for name, param in model.named_parameters():
if 'bn' not in name and 'bias' not in name:
robust_loss += args.l1 * param.abs().sum()
opt.zero_grad()
robust_loss.backward()
opt.step()
clean_input = normalize(X)
clean_input.requires_grad = True
output = model(clean_input)
if args.mixup:
loss = mixup_criterion(criterion, output, y_a, y_b, lam)
else:
loss = criterion(output, y)
# Get the gradient norm values
input_grads = torch.autograd.grad(loss,
clean_input,
create_graph=False)[0]
# Record the statstic values
train_robust_loss += robust_loss.item() * adv_y.size(0)
train_robust_acc += (robust_output.max(1)[1] == adv_y).sum().item()
train_loss += loss.item() * y.size(0)
train_acc += (output.max(1)[1] == y).sum().item()
train_n += y.size(0)
train_grad += input_grads.abs().sum()
train_time = time.time()
if args.earlystopPGD:
print('Iter mean: ',
record_iter.mean().item(), ' Iter std: ',
record_iter.std().item())
# Evaluate on test data
model.eval()
test_loss = 0
test_acc = 0
test_n = 0
test_adv_loss = 0
test_adv_acc = 0
test_adv_n = 0
for i, batch in enumerate(test_batches):
X, y = batch['input'], batch['target']
clean_input = normalize(X)
output = model(clean_input)
loss = criterion(output, y)
test_loss += loss.item() * y.size(0)
test_acc += (output.max(1)[1] == y).sum().item()
test_n += y.size(0)
for i, batch in enumerate(test_adv_batches):
adv_input = normalize(batch['input'])
y = batch['target']
robust_output = model(adv_input)
robust_loss = criterion(robust_output, y)
test_adv_loss += robust_loss.item() * y.size(0)
test_adv_acc += (robust_output.max(1)[1] == y).sum().item()
test_adv_n += y.size(0)
test_time = time.time()
logger.info('%d \t %.4f \t %.4f \t\t %.4f \t %.4f', epoch + 1,
train_acc / train_n, train_robust_acc / train_n,
test_acc / test_n, test_adv_acc / test_adv_n)
# Save results
train_loss_record.append(train_loss / train_n)
train_acc_record.append(train_acc / train_n)
train_robust_loss_record.append(train_robust_loss / train_n)
train_robust_acc_record.append(train_robust_acc / train_n)
np.savetxt(args.fname + '/train_loss_record.txt',
np.array(train_loss_record))
np.savetxt(args.fname + '/train_acc_record.txt',
np.array(train_acc_record))
np.savetxt(args.fname + '/train_robust_loss_record.txt',
np.array(train_robust_loss_record))
np.savetxt(args.fname + '/train_robust_acc_record.txt',
np.array(train_robust_acc_record))
test_loss_record.append(test_loss / train_n)
test_acc_record.append(test_acc / train_n)
test_adv_loss_record.append(test_adv_loss / train_n)
test_adv_acc_record.append(test_adv_acc / train_n)
np.savetxt(args.fname + '/test_loss_record.txt',
np.array(test_loss_record))
np.savetxt(args.fname + '/test_acc_record.txt',
np.array(test_acc_record))
np.savetxt(args.fname + '/test_adv_loss_record.txt',
np.array(test_adv_loss_record))
np.savetxt(args.fname + '/test_adv_acc_record.txt',
np.array(test_adv_acc_record))
# save checkpoint
if epoch > 99 or (epoch +
1) % args.chkpt_iters == 0 or epoch + 1 == epochs:
torch.save(model.state_dict(),
os.path.join(args.fname, f'model_{epoch}.pth'))
torch.save(opt.state_dict(),
os.path.join(args.fname, f'opt_{epoch}.pth'))
# save best
if test_adv_acc / test_n > best_test_adv_acc:
torch.save(
{
'state_dict': model.state_dict(),
'test_adv_acc': test_adv_acc / test_n,
'test_adv_loss': test_adv_loss / test_n,
'test_loss': test_loss / test_n,
'test_acc': test_acc / test_n,
}, os.path.join(args.fname, f'model_best.pth'))
best_test_adv_acc = test_adv_acc / test_n
num_workers=args.n_workers)
testloader = DataLoader(dataset=testset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.n_workers)
print('[init cache]')
partition_size, _ = divmod(len(unlabeledloader.sampler), args.n_partitions)
assert _ == 0
cache = Cache(n_entries=partition_size,
entry_size=args.sparsity).to(args.output_device)
print('[init model]')
model = WideResNet(num_classes=args.n_classes).to(args.output_device)
model_ema = WideResNet(num_classes=args.n_classes).to(args.output_device)
model_ema.load_state_dict(model.state_dict())
print('[init optimizer]')
optimizer = optim.AdamW(model.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
print('[init critera]')
criterion_labeled = CrossEntropyLoss()
criterion_unlabeled = MatchingLoss()
criterion_val = nn.CrossEntropyLoss()
print('[start training]')
with SummaryWriter(log_dir=args.tensorboard_dir) as tblogger:
run(labeledloader=labeledloader,
unlabeledloader=unlabeledloader,
def main():
global args, best_prec1
args = parser.parse_args()
if args.tensorboard:
configure("runs/%s" % (args.name))
# Data loading code
# normalize = 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.Resize(256,256),
# transforms.ToTensor(),
# transforms.Lambda(lambda x: F.pad(
# Variable(x.unsqueeze(0), requires_grad=False, volatile=True),
# (4,4,4,4),mode='reflect').data.squeeze()),
# transforms.ToPILImage(),
# transforms.RandomCrop(32),
# transforms.RandomHorizontalFlip(),
# transforms.ToTensor(),
# normalize,
# ])
# else:
# transform_train = transforms.Compose([
# transforms.ToTensor(),
# normalize,
# ])
# transform_test = transforms.Compose([
# transforms.ToTensor(),
# normalize
# ])
#kwargs = {'num_workers': 1, 'pin_memory': True}
#assert(args.dataset == 'cifar10' or args.dataset == 'cifar100')
train_data_path = "/home/mil/gupta/ifood18/data/training_set/train_set/"
val_data_path = "/home/mil/gupta/ifood18/data/val_set/"
train_label = "/home/mil/gupta/ifood18/data/labels/train_info.csv"
val_label = "/home/mil/gupta/ifood18/data/labels/val_info.csv"
#transformations = transforms.Compose([transforms.ToTensor()])
# train_h5 = "/home/mil/gupta/ifood18/data/h5data/train_data.h5py"
train_h5 = "/home/mil/gupta/ifood18/data/h5data/train_data_partial.h5py"
train_dataset = H5Dataset(train_h5)
#val_dataset = H5Dataset(val_data_path, val_label, transform= None)
#custom_mnist_from_csv = \
# CustomDatasetFromCSV('../data/mnist_in_csv.csv',
# 28, 28,
# transformations)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
batch_size=64,
shuffle=True,
num_workers=1)
#val_loader = torch.utils.data.DataLoader(dataset=val_dataset,batch_size=256,num_workers=1,shuffle=True)
# train_labels = pd.read_csv('./data/labels/train_info.csv')
# train_loader = torch.utils.data.DataLoader(
# datasets.__dict__[args.dataset.upper()](train_data_path, train=True, download=True,
# transform=transform_train),
# batch_size=args.batch_size, shuffle=True, **kwargs)
# val_loader = torch.utils.data.DataLoader(
# datasets.__dict__[args.dataset.upper()](val_data_path, train=False, transform=transform_test),
# batch_size=args.batch_size, shuffle=True, **kwargs)
# create model
model = WideResNet(args.layers,
211,
args.widen_factor,
dropRate=args.droprate)
# 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']
best_prec1 = checkpoint['best_prec1']
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 optimizer
criterion = nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
nesterov=args.nesterov,
weight_decay=args.weight_decay)
for epoch in range(args.start_epoch, args.epochs):
adjust_learning_rate(optimizer, epoch + 1)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch)
#if epoch % args.validate_freq == 0:
# evaluate on validation set
prec3 = validate(val_loader, model, criterion, epoch)
# remember best prec@1 and save checkpoint
is_best = prec3 > best_prec3
best_prec1 = max(prec3, best_prec3)
save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'best_prec1': best_prec3,
}, is_best)
print('Best accuracy: ', best_prec1)
def main():
args = get_args()
if args.awp_gamma <= 0.0:
args.awp_warmup = np.infty
if not os.path.exists(args.fname):
os.makedirs(args.fname)
logger = logging.getLogger(__name__)
logging.basicConfig(
format='[%(asctime)s] - %(message)s',
datefmt='%Y/%m/%d %H:%M:%S',
level=logging.DEBUG,
handlers=[
logging.FileHandler(os.path.join(args.fname, 'eval.log' if args.eval else 'output.log')),
logging.StreamHandler()
])
logger.info(args)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
transforms = [Crop(32, 32), FlipLR()]
if args.cutout:
transforms.append(Cutout(args.cutout_len, args.cutout_len))
if args.val:
try:
dataset = torch.load("cifar10_validation_split.pth")
except:
print("Couldn't find a dataset with a validation split, did you run "
"generate_validation.py?")
return
val_set = list(zip(transpose(dataset['val']['data']/255.), dataset['val']['labels']))
val_batches = Batches(val_set, args.batch_size, shuffle=False, num_workers=2)
else:
dataset = cifar10(args.data_dir)
train_set = list(zip(transpose(pad(dataset['train']['data'], 4)/255.),
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)
test_set = list(zip(transpose(dataset['test']['data']/255.), dataset['test']['labels']))
test_batches = Batches(test_set, args.batch_size_test, shuffle=False, num_workers=2)
epsilon = (args.epsilon / 255.)
pgd_alpha = (args.pgd_alpha / 255.)
if args.model == 'PreActResNet18':
model = PreActResNet18()
proxy = PreActResNet18()
elif args.model == 'WideResNet':
model = WideResNet(34, 10, widen_factor=args.width_factor, dropRate=0.0)
proxy = WideResNet(34, 10, widen_factor=args.width_factor, dropRate=0.0)
else:
raise ValueError("Unknown model")
model = nn.DataParallel(model).cuda()
proxy = nn.DataParallel(proxy).cuda()
if args.l2:
decay, no_decay = [], []
for name,param in model.named_parameters():
if 'bn' not in name and 'bias' not in name:
decay.append(param)
else:
no_decay.append(param)
params = [{'params':decay, 'weight_decay':args.l2},
{'params':no_decay, 'weight_decay': 0 }]
else:
params = model.parameters()
opt = torch.optim.SGD(params, lr=args.lr_max, momentum=0.9, weight_decay=5e-4)
proxy_opt = torch.optim.SGD(proxy.parameters(), lr=0.01)
awp_adversary = AdvWeightPerturb(model=model, proxy=proxy, proxy_optim=proxy_opt, gamma=args.awp_gamma)
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':
epochs = int(math.ceil(args.epochs / args.attack_iters))
else:
epochs = args.epochs
if args.lr_schedule == 'superconverge':
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.
elif args.lr_schedule == 'linear':
lr_schedule = lambda t: np.interp([t], [0, args.epochs // 3, args.epochs * 2 // 3, args.epochs], [args.lr_max, args.lr_max, args.lr_max / 10, args.lr_max / 100])[0]
elif args.lr_schedule == 'onedrop':
def lr_schedule(t):
if t < args.lr_drop_epoch:
return args.lr_max
else:
return args.lr_one_drop
elif args.lr_schedule == 'multipledecay':
def lr_schedule(t):
return args.lr_max - (t//(args.epochs//10))*(args.lr_max/10)
elif args.lr_schedule == 'cosine':
def lr_schedule(t):
return args.lr_max * 0.5 * (1 + np.cos(t / args.epochs * np.pi))
elif args.lr_schedule == 'cyclic':
lr_schedule = lambda t: np.interp([t], [0, 0.4 * args.epochs, args.epochs], [0, args.lr_max, 0])[0]
best_test_robust_acc = 0
best_val_robust_acc = 0
if args.resume:
start_epoch = args.resume
model.load_state_dict(torch.load(os.path.join(args.fname, f'model_{start_epoch-1}.pth')))
opt.load_state_dict(torch.load(os.path.join(args.fname, f'opt_{start_epoch-1}.pth')))
logger.info(f'Resuming at epoch {start_epoch}')
if os.path.exists(os.path.join(args.fname, f'model_best.pth')):
best_test_robust_acc = torch.load(os.path.join(args.fname, f'model_best.pth'))['test_robust_acc']
if args.val:
best_val_robust_acc = torch.load(os.path.join(args.fname, f'model_val.pth'))['val_robust_acc']
else:
start_epoch = 0
if args.eval:
if not args.resume:
logger.info("No model loaded to evaluate, specify with --resume FNAME")
return
logger.info("[Evaluation mode]")
logger.info('Epoch \t Train Time \t Test Time \t LR \t \t Train Loss \t Train Acc \t Train Robust Loss \t Train Robust Acc \t Test Loss \t Test Acc \t Test Robust Loss \t Test Robust Acc')
for epoch in range(start_epoch, epochs):
start_time = time.time()
train_loss = 0
train_acc = 0
train_robust_loss = 0
train_robust_acc = 0
train_n = 0
for i, batch in enumerate(train_batches):
if args.eval:
break
X, y = batch['input'], batch['target']
if args.mixup:
X, y_a, y_b, lam = mixup_data(X, y, args.mixup_alpha)
X, y_a, y_b = map(Variable, (X, y_a, y_b))
lr = lr_schedule(epoch + (i + 1) / len(train_batches))
opt.param_groups[0].update(lr=lr)
if args.attack == 'pgd':
# Random initialization
if args.mixup:
delta = attack_pgd(model, X, y, epsilon, pgd_alpha, args.attack_iters, args.restarts, args.norm, mixup=True, y_a=y_a, y_b=y_b, lam=lam)
else:
delta = attack_pgd(model, X, y, epsilon, pgd_alpha, args.attack_iters, args.restarts, args.norm)
delta = delta.detach()
elif args.attack == 'fgsm':
delta = attack_pgd(model, X, y, epsilon, args.fgsm_alpha*epsilon, 1, 1, args.norm)
# Standard training
elif args.attack == 'none':
delta = torch.zeros_like(X)
X_adv = normalize(torch.clamp(X + delta[:X.size(0)], min=lower_limit, max=upper_limit))
model.train()
# calculate adversarial weight perturbation and perturb it
if epoch >= args.awp_warmup:
# not compatible to mixup currently.
assert (not args.mixup)
awp = awp_adversary.calc_awp(inputs_adv=X_adv,
targets=y)
awp_adversary.perturb(awp)
robust_output = model(X_adv)
if args.mixup:
robust_loss = mixup_criterion(criterion, robust_output, y_a, y_b, lam)
else:
robust_loss = criterion(robust_output, y)
if args.l1:
for name,param in model.named_parameters():
if 'bn' not in name and 'bias' not in name:
robust_loss += args.l1*param.abs().sum()
opt.zero_grad()
robust_loss.backward()
opt.step()
if epoch >= args.awp_warmup:
awp_adversary.restore(awp)
output = model(normalize(X))
if args.mixup:
loss = mixup_criterion(criterion, output, y_a, y_b, lam)
else:
loss = criterion(output, y)
train_robust_loss += robust_loss.item() * y.size(0)
train_robust_acc += (robust_output.max(1)[1] == y).sum().item()
train_loss += loss.item() * y.size(0)
train_acc += (output.max(1)[1] == y).sum().item()
train_n += y.size(0)
train_time = time.time()
model.eval()
test_loss = 0
test_acc = 0
test_robust_loss = 0
test_robust_acc = 0
test_n = 0
for i, batch in enumerate(test_batches):
X, y = batch['input'], batch['target']
# Random initialization
if args.attack == 'none':
delta = torch.zeros_like(X)
else:
delta = attack_pgd(model, X, y, epsilon, pgd_alpha, args.attack_iters_test, args.restarts, args.norm, early_stop=args.eval)
delta = delta.detach()
robust_output = model(normalize(torch.clamp(X + delta[:X.size(0)], min=lower_limit, max=upper_limit)))
robust_loss = criterion(robust_output, y)
output = model(normalize(X))
loss = criterion(output, y)
test_robust_loss += robust_loss.item() * y.size(0)
test_robust_acc += (robust_output.max(1)[1] == y).sum().item()
test_loss += loss.item() * y.size(0)
test_acc += (output.max(1)[1] == y).sum().item()
test_n += y.size(0)
test_time = time.time()
if args.val:
val_loss = 0
val_acc = 0
val_robust_loss = 0
val_robust_acc = 0
val_n = 0
for i, batch in enumerate(val_batches):
X, y = batch['input'], batch['target']
# Random initialization
if args.attack == 'none':
delta = torch.zeros_like(X)
else:
delta = attack_pgd(model, X, y, epsilon, pgd_alpha, args.attack_iters_test, args.restarts, args.norm, early_stop=args.eval)
delta = delta.detach()
robust_output = model(normalize(torch.clamp(X + delta[:X.size(0)], min=lower_limit, max=upper_limit)))
robust_loss = criterion(robust_output, y)
output = model(normalize(X))
loss = criterion(output, y)
val_robust_loss += robust_loss.item() * y.size(0)
val_robust_acc += (robust_output.max(1)[1] == y).sum().item()
val_loss += loss.item() * y.size(0)
val_acc += (output.max(1)[1] == y).sum().item()
val_n += y.size(0)
if not args.eval:
logger.info('%d \t %.1f \t \t %.1f \t \t %.4f \t %.4f \t %.4f \t %.4f \t \t %.4f \t \t %.4f \t %.4f \t %.4f \t \t %.4f',
epoch, train_time - start_time, test_time - train_time, lr,
train_loss/train_n, train_acc/train_n, train_robust_loss/train_n, train_robust_acc/train_n,
test_loss/test_n, test_acc/test_n, test_robust_loss/test_n, test_robust_acc/test_n)
if args.val:
logger.info('validation %.4f \t %.4f \t %.4f \t %.4f',
val_loss/val_n, val_acc/val_n, val_robust_loss/val_n, val_robust_acc/val_n)
if val_robust_acc/val_n > best_val_robust_acc:
torch.save({
'state_dict':model.state_dict(),
'test_robust_acc':test_robust_acc/test_n,
'test_robust_loss':test_robust_loss/test_n,
'test_loss':test_loss/test_n,
'test_acc':test_acc/test_n,
'val_robust_acc':val_robust_acc/val_n,
'val_robust_loss':val_robust_loss/val_n,
'val_loss':val_loss/val_n,
'val_acc':val_acc/val_n,
}, os.path.join(args.fname, f'model_val.pth'))
best_val_robust_acc = val_robust_acc/val_n
# save checkpoint
if (epoch+1) % args.chkpt_iters == 0 or epoch+1 == epochs:
torch.save(model.state_dict(), os.path.join(args.fname, f'model_{epoch}.pth'))
torch.save(opt.state_dict(), os.path.join(args.fname, f'opt_{epoch}.pth'))
# save best
if test_robust_acc/test_n > best_test_robust_acc:
torch.save({
'state_dict':model.state_dict(),
'test_robust_acc':test_robust_acc/test_n,
'test_robust_loss':test_robust_loss/test_n,
'test_loss':test_loss/test_n,
'test_acc':test_acc/test_n,
}, os.path.join(args.fname, f'model_best.pth'))
best_test_robust_acc = test_robust_acc/test_n
else:
logger.info('%d \t %.1f \t \t %.1f \t \t %.4f \t %.4f \t %.4f \t %.4f \t \t %.4f \t \t %.4f \t %.4f \t %.4f \t \t %.4f',
epoch, train_time - start_time, test_time - train_time, -1,
-1, -1, -1, -1,
test_loss/test_n, test_acc/test_n, test_robust_loss/test_n, test_robust_acc/test_n)
return
def main():
args = get_args()
if not os.path.exists(args.fname):
os.makedirs(args.fname)
logger = logging.getLogger(__name__)
logging.basicConfig(
format='[%(asctime)s] - %(message)s',
datefmt='%Y/%m/%d %H:%M:%S',
level=logging.DEBUG,
handlers=[
logging.FileHandler(os.path.join(args.fname, 'eval.log' if args.eval else 'output.log')),
logging.StreamHandler()
])
logger.info(args)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
train_transform = transforms.Compose([
transforms.ToTensor(),
# transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])
test_transform = transforms.Compose([
transforms.ToTensor(),
# transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])
num_workers = 2
train_dataset = datasets.SVHN(
args.data_dir, split='train', transform=train_transform, download=True)
test_dataset = datasets.SVHN(
args.data_dir, split='test', transform=test_transform, download=True)
train_loader = torch.utils.data.DataLoader(
dataset=train_dataset,
batch_size=args.batch_size,
shuffle=True,
pin_memory=True,
num_workers=num_workers,
)
test_loader = torch.utils.data.DataLoader(
dataset=test_dataset,
batch_size=args.batch_size,
shuffle=False,
pin_memory=True,
num_workers=2,
)
epsilon = (args.epsilon / 255.)
pgd_alpha = (args.pgd_alpha / 255.)
# model = models_dict[args.architecture]().cuda()
# model.apply(initialize_weights)
if args.model == 'PreActResNet18':
model = PreActResNet18()
elif args.model == 'WideResNet':
model = WideResNet(34, 10, widen_factor=args.width_factor, dropRate=0.0)
else:
raise ValueError("Unknown model")
model = model.cuda()
model.train()
if args.l2:
decay, no_decay = [], []
for name,param in model.named_parameters():
if 'bn' not in name and 'bias' not in name:
decay.append(param)
else:
no_decay.append(param)
params = [{'params':decay, 'weight_decay':args.l2},
{'params':no_decay, 'weight_decay': 0 }]
else:
params = model.parameters()
opt = torch.optim.SGD(params, lr=args.lr_max, momentum=0.9, weight_decay=5e-4)
criterion = nn.CrossEntropyLoss()
epochs = args.epochs
if args.lr_schedule == 'superconverge':
lr_schedule = lambda t: np.interp([t], [0, args.epochs * 2 // 5, args.epochs], [0, args.lr_max, 0])[0]
# lr_schedule = lambda t: np.interp([t], [0, args.epochs], [0, args.lr_max])[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.
if args.resume:
start_epoch = args.resume
model.load_state_dict(torch.load(os.path.join(args.fname, f'model_{start_epoch-1}.pth')))
opt.load_state_dict(torch.load(os.path.join(args.fname, f'opt_{start_epoch-1}.pth')))
logger.info(f'Resuming at epoch {start_epoch}')
else:
start_epoch = 0
if args.eval:
if not args.resume:
logger.info("No model loaded to evaluate, specify with --resume FNAME")
return
logger.info("[Evaluation mode]")
logger.info('Epoch \t Train Time \t Test Time \t LR \t \t Train Loss \t Train Acc \t Train Robust Loss \t Train Robust Acc \t Test Loss \t Test Acc \t Test Robust Loss \t Test Robust Acc')
for epoch in range(start_epoch, epochs):
model.train()
start_time = time.time()
train_loss = 0
train_acc = 0
train_robust_loss = 0
train_robust_acc = 0
train_n = 0
for i, (X, y) in enumerate(train_loader):
if args.eval:
break
X, y = X.cuda(), y.cuda()
if args.mixup:
X, y_a, y_b, lam = mixup_data(X, y, args.mixup_alpha)
X, y_a, y_b = map(Variable, (X, y_a, y_b))
lr = lr_schedule(epoch + (i + 1) / len(train_loader))
opt.param_groups[0].update(lr=lr)
if args.attack == 'pgd':
# Random initialization
if args.mixup:
delta = attack_pgd(model, X, y, epsilon, pgd_alpha, args.attack_iters, args.restarts, args.norm, mixup=True, y_a=y_a, y_b=y_b, lam=lam)
else:
delta = attack_pgd(model, X, y, epsilon, pgd_alpha, args.attack_iters, args.restarts, args.norm)
delta = delta.detach()
# Standard training
elif args.attack == 'none':
delta = torch.zeros_like(X)
robust_output = model(normalize(torch.clamp(X + delta[:X.size(0)], min=lower_limit, max=upper_limit)))
if args.mixup:
robust_loss = mixup_criterion(criterion, robust_output, y_a, y_b, lam)
else:
robust_loss = criterion(robust_output, y)
if args.l1:
for name,param in model.named_parameters():
if 'bn' not in name and 'bias' not in name:
robust_loss += args.l1*param.abs().sum()
opt.zero_grad()
robust_loss.backward()
opt.step()
output = model(normalize(X))
if args.mixup:
loss = mixup_criterion(criterion, output, y_a, y_b, lam)
else:
loss = criterion(output, y)
train_robust_loss += robust_loss.item() * y.size(0)
train_robust_acc += (robust_output.max(1)[1] == y).sum().item()
train_loss += loss.item() * y.size(0)
train_acc += (output.max(1)[1] == y).sum().item()
train_n += y.size(0)
train_time = time.time()
model.eval()
test_loss = 0
test_acc = 0
test_robust_loss = 0
test_robust_acc = 0
test_n = 0
for i, (X, y) in enumerate(test_loader):
X, y = X.cuda(), y.cuda()
# Random initialization
if args.attack == 'none':
delta = torch.zeros_like(X)
else:
delta = attack_pgd(model, X, y, epsilon, pgd_alpha, args.attack_iters, args.restarts, args.norm, early_stop=args.eval)
delta = delta.detach()
robust_output = model(normalize(torch.clamp(X + delta[:X.size(0)], min=lower_limit, max=upper_limit)))
robust_loss = criterion(robust_output, y)
output = model(normalize(X))
loss = criterion(output, y)
test_robust_loss += robust_loss.item() * y.size(0)
test_robust_acc += (robust_output.max(1)[1] == y).sum().item()
test_loss += loss.item() * y.size(0)
test_acc += (output.max(1)[1] == y).sum().item()
test_n += y.size(0)
test_time = time.time()
if not args.eval:
logger.info('%d \t %.1f \t \t %.1f \t \t %.4f \t %.4f \t %.4f \t %.4f \t \t %.4f \t \t %.4f \t %.4f \t %.4f \t \t %.4f',
epoch, train_time - start_time, test_time - train_time, lr,
train_loss/train_n, train_acc/train_n, train_robust_loss/train_n, train_robust_acc/train_n,
test_loss/test_n, test_acc/test_n, test_robust_loss/test_n, test_robust_acc/test_n)
if (epoch+1) % args.chkpt_iters == 0 or epoch+1 == epochs:
torch.save(model.state_dict(), os.path.join(args.fname, f'model_{epoch}.pth'))
torch.save(opt.state_dict(), os.path.join(args.fname, f'opt_{epoch}.pth'))
else:
logger.info('%d \t %.1f \t \t %.1f \t \t %.4f \t %.4f \t %.4f \t %.4f \t \t %.4f \t \t %.4f \t %.4f \t %.4f \t \t %.4f',
epoch, train_time - start_time, test_time - train_time, -1,
-1, -1, -1, -1,
test_loss/test_n, test_acc/test_n, test_robust_loss/test_n, test_robust_acc/test_n)
return
def train_fusionWRN6(
self,
epochs1=120,
epochs2=3,
device="cuda:0"
): # https://github.com/xternalz/WideResNet-pytorch.git #120 80
sys.path.append('/media/rene/code/WideResNet-pytorch')
from wideresnet import WideResNet
epochs1, epochs2 = int(epochs1), int(epochs2)
num_workers = 4
PATH = Path('/media/rene/data/')
save_path = Path('/media/rene/code/WideResNet-pytorch/runs')
model_name_list = [
'WideResNet-28-10_0/model_best.pth.tar',
'WideResNet-28-10_1/model_best.pth.tar',
'WideResNet-28-10_2/model_best.pth.tar',
'WideResNet-28-10_3/model_best.pth.tar',
'WideResNet-28-10_4/model_best.pth.tar',
'WideResNet-28-10_5/model_best.pth.tar'
]
batch_size = 8
dataloaders, dataset_sizes = make_batch_gen_cifar(str(PATH),
batch_size,
num_workers,
valid_name='valid')
pretrained_model_list = []
# First trained model was with DATA PARALLEL
model = WideResNet(28, 10, 20)
model = model.to(device)
state_dict = torch.load(
os.path.join(
save_path,
'WideResNet-28-10_0/model_best.pth.tar'))['state_dict']
# create new OrderedDict that does not contain `module.`
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] # remove `module.`
new_state_dict[name] = v
model.load_state_dict(new_state_dict)
pretrained_model_list.append(model)
# get all the models
for i, model_name in enumerate(model_name_list[1:]):
print('------------loading model: ', model_name)
model = WideResNet(28, 10, 20)
model = model.to(device)
# original saved file with DataParallel
state_dict = torch.load(os.path.join(save_path,
model_name))['state_dict']
model.load_state_dict(state_dict)
pretrained_model_list.append(model)
model = Fusion6(pretrained_model_list, num_input=60, num_output=10)
###################### TRAIN LAST FEW LAYERS
# print('training last few layers')
model_name = 'Fusion6_WRN_1'
for p in model.parameters():
p.requires_grad = True
for p in model.model1.parameters():
p.requires_grad = False
for p in model.model2.parameters():
p.requires_grad = False
for p in model.model3.parameters():
p.requires_grad = False
for p in model.model4.parameters():
p.requires_grad = False
for p in model.model5.parameters():
p.requires_grad = False
for p in model.model6.parameters():
p.requires_grad = False
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(filter(lambda p: p.requires_grad,
model.parameters()),
lr=.005,
momentum=0.9,
weight_decay=5e-4)
scheduler = lr_scheduler.StepLR(optimizer,
step_size=int(epochs1 / 3),
gamma=0.1)
best_acc, model = train_model(model,
criterion,
optimizer,
scheduler,
epochs1,
dataloaders,
dataset_sizes,
device=device)
torch.save(model.state_dict(), str(save_path / model_name))
######################## TRAIN ALL LAYERS
# net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count()))
# model.load_state_dict(torch.load(save_path / 'Fusion2_WRN_1'))
model_name = 'Fusion6_WRN_2'
batch_size = 1
print('---------', batch_size)
dataloaders, dataset_sizes = make_batch_gen_cifar(str(PATH),
batch_size,
num_workers,
valid_name='valid')
for p in model.parameters():
p.requires_grad = True
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(filter(lambda p: p.requires_grad,
model.parameters()),
lr=.0001,
momentum=0.9,
weight_decay=5e-4)
scheduler = lr_scheduler.StepLR(optimizer,
step_size=int(epochs2 / 3),
gamma=0.1)
best_acc, model = train_model(model,
criterion,
optimizer,
scheduler,
epochs2,
dataloaders,
dataset_sizes,
device=device)
torch.save(model.state_dict(), str(save_path / model_name))
def main():
global args, best_prec1
args = parser.parse_args()
if args.tensorboard: configure(args.checkpoint_dir+"/%s"%(args.name))
# Data loading code
normalize = 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.ToTensor(),
transforms.Lambda(lambda x: F.pad(x.unsqueeze(0),
(4,4,4,4),mode='reflect').squeeze()),
transforms.ToPILImage(),
transforms.RandomCrop(32),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
])
else:
transform_train = transforms.Compose([
transforms.ToTensor(),
normalize,
])
transform_test = transforms.Compose([
transforms.ToTensor(),
normalize
])
kwargs = {'num_workers': 1, 'pin_memory': True}
assert(args.dataset == 'cifar10' or args.dataset == 'cifar100')
train_loader = torch.utils.data.DataLoader(
datasets.__dict__[args.dataset.upper()](args.dataset_path, train=True, download=True,
transform=transform_train),
batch_size=args.batch_size, shuffle=True, **kwargs)
val_loader = torch.utils.data.DataLoader(
datasets.__dict__[args.dataset.upper()](args.dataset_path, train=False, transform=transform_test),
batch_size=args.batch_size, shuffle=True, **kwargs)
device = "cuda" if torch.cuda.is_available() else "cpu"
# create model
model = WideResNet(args.layers, args.dataset == 'cifar10' and 10 or 100,
args.widen_factor, dropRate=args.droprate,
semantic_loss=args.sloss, device=device)
# 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.to(device)
# 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, map_location=torch.device('cpu'))
args.start_epoch = checkpoint['epoch']
best_prec1 = checkpoint['best_prec1']
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 optimizer
criterion = nn.CrossEntropyLoss().to(device)
optimizer = torch.optim.SGD(model.parameters(), args.lr,
momentum=args.momentum, nesterov=args.nesterov,
weight_decay=args.weight_decay)
# cosine learning rate
# scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, len(train_loader)*args.epochs, eta_min=1e-6)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=40, gamma=.2)
calc_logic, logic_net, examples, logic_optimizer, decoder_optimizer, logic_scheduler, decoder_scheduler = None,None,None,None,None,None,None
if args.dataset == "cifar100":
examples, logic_fn, group_precision = get_cifar100_experiment_params(train_loader.dataset)
assert logic_fn(torch.arange(100), examples).all()
else:
examples, logic_fn, group_precision = get_cifar10_experiment_params(train_loader.dataset)
assert logic_fn(torch.arange(10), examples).all()
if args.sloss:
examples = examples.to(device)
logic_net = LogicNet(num_classes=len(train_loader.dataset.classes))
logic_net.to(device)
# logic_optimizer = torch.optim.Adam(logic_net.parameters(), 1e-1*args.lr)
logic_optimizer = torch.optim.SGD(logic_net.parameters(),
args.lr,
momentum=args.momentum,
nesterov=args.nesterov,
weight_decay=args.weight_decay)
# logic_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(logic_optimizer, len(train_loader) * args.epochs, eta_min=1e-8)
logic_scheduler = torch.optim.lr_scheduler.StepLR(logic_optimizer, step_size=25, gamma=.2)
# decoder_optimizer = torch.optim.Adam(model.global_paramters, args.lr)
decoder_optimizer = torch.optim.SGD(model.global_paramters,
args.lr,
momentum=args.momentum,
nesterov=args.nesterov,
weight_decay=args.weight_decay)
# decoder_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(decoder_optimizer, len(train_loader) * args.epochs, eta_min=1e-8)
decoder_scheduler = torch.optim.lr_scheduler.StepLR(decoder_optimizer, step_size=25, gamma=.2)
calc_logic = lambda predictions, targets: calc_logic_loss(predictions, targets, logic_net, logic_fn, num_classes=model.num_classes, device=device)
# override the oprimizer from above
optimizer = torch.optim.SGD(model.local_parameters, # TODO: still might be better for parameters()
args.lr,
momentum=args.momentum,
nesterov=args.nesterov,
weight_decay=args.weight_decay)
# scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, len(train_loader) * args.epochs)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=40, gamma=.2)
name = "_".join([str(getattr(args, source)) for source in ['lr', 'sloss', 'sloss_weight', 'dataset']])
if args.resume:
targets, preds, outs = validate(val_loader, model, criterion, 1, args, group_precision, device=device)
from sklearn.metrics import confusion_matrix
import pickle
confusion_matrix(targets, preds)
group_precision(torch.tensor(targets), torch.tensor(np.concatenate(outs, axis=0)))
dict_ = {"targets": targets, "pred": np.concatenate(outs, axis=0)}
f = open('../semantic_loss/notebooks/results.pickle', 'wb')
pickle.dump(dict_, f)
f.close()
import pdb
pdb.set_trace()
for epoch in range(args.start_epoch, args.epochs):
# train for one epoch
train(train_loader, model, logic_net,
criterion, examples,
optimizer, logic_optimizer, decoder_optimizer,
scheduler, logic_scheduler, decoder_scheduler,
epoch, args, calc_logic, device=device)
# evaluate on validation set
prec1 = validate(val_loader, model, criterion, epoch, args, group_precision, device=device)
# remember best prec@1 and save checkpoint
is_best = prec1 > best_prec1
best_prec1 = max(prec1, best_prec1)
save_checkpoint({
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
}, is_best, filename=f"{name}.checkpoint.pt")
if args.sloss:
logic_scheduler.step()
decoder_scheduler.step()
scheduler.step()
print('Best accuracy: ', best_prec1)
def train_fusionWRN_last3(
self,
epochs1=40,
epochs2=25,
device="cuda:1"
): # https://github.com/xternalz/WideResNet-pytorch.git #120 80
with torch.cuda.device(1):
sys.path.append('/media/rene/code/WideResNet-pytorch')
from wideresnet import WideResNet
epochs1, epochs2 = int(epochs1), int(epochs2)
num_workers = 4
PATH = Path('/media/rene/data/')
save_path = Path('/media/rene/code/WideResNet-pytorch/runs')
model_name_list = [
'WideResNet-28-10_0/model_best.pth.tar',
'WideResNet-28-10_1/model_best.pth.tar',
'WideResNet-28-10_2/model_best.pth.tar',
'WideResNet-28-10_3/model_best.pth.tar',
'WideResNet-28-10_4/model_best.pth.tar',
'WideResNet-28-10_5/model_best.pth.tar'
]
batch_size = 300
transform_test = transforms.Compose([
transforms.ToTensor(),
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]])
])
dataloaders, dataset_sizes = make_batch_gen_cifar(
str(PATH),
batch_size,
num_workers,
valid_name='valid',
transformation=transform_test)
pretrained_model_list = []
# First trained model was with DATA PARALLEL
model = WideResNet(28, 10, 20)
model = model.to(device)
state_dict = torch.load(
os.path.join(
save_path,
'WideResNet-28-10_0/model_best.pth.tar'))['state_dict']
# create new OrderedDict that does not contain `module.`
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] # remove `module.`
new_state_dict[name] = v
model.load_state_dict(new_state_dict)
pretrained_model_list.append(model)
# get all the models
for i, model_name in enumerate(model_name_list[1:3]):
print('------------loading model: ', model_name)
model = WideResNet(28, 10, 20)
model = model.to(device)
# original saved file with DataParallel
state_dict = torch.load(os.path.join(save_path,
model_name))['state_dict']
model.load_state_dict(state_dict)
pretrained_model_list.append(model)
model = Fusion3(pretrained_model_list, num_input=30, num_output=10)
###################### TRAIN LAST FEW LAYERS
print('training last few layers')
model_name = 'fusionWRN_last3_1'
for p in model.parameters():
p.requires_grad = True
for p in model.model1.parameters():
p.requires_grad = False
for p in model.model2.parameters():
p.requires_grad = False
for p in model.model3.parameters():
p.requires_grad = False
# criterion = nn.CrossEntropyLoss()
# optimizer = optim.SGD(filter(lambda p: p.requires_grad,model.parameters()), lr=.005, momentum=0.9, weight_decay=5e-4)
# scheduler = lr_scheduler.StepLR(optimizer, step_size=int(epochs1/3), gamma=0.3)
# best_acc, model = train_model(model, criterion, optimizer, scheduler, epochs1,
# dataloaders, dataset_sizes, device=device)
# torch.save(model.state_dict(), str(save_path / model_name))
######################## TRAIN ALL LAYERS
model.load_state_dict(torch.load(save_path / 'fusionWRN_last3_1'))
model = model.to(device)
model_name = 'fusionWRN_last3_2'
batch_size = 88
dataloaders, dataset_sizes = make_batch_gen_cifar(
str(PATH),
batch_size,
num_workers,
valid_name='valid',
transformation=transform_test)
### ONLY THE LAST BLOCK:
for i, child in enumerate(model.model1.children()):
if (i >= 3):
for p in child.parameters():
p.requires_grad = True
for i, child in enumerate(model.model2.children()):
if (i >= 3):
for p in child.parameters():
p.requires_grad = True
for i, child in enumerate(model.model3.children()):
if (i >= 3):
for p in child.parameters():
p.requires_grad = True
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(filter(lambda p: p.requires_grad,
model.parameters()),
lr=.0001,
momentum=0.9,
weight_decay=5e-4)
scheduler = lr_scheduler.StepLR(optimizer,
step_size=int(epochs2 / 2),
gamma=0.1)
best_acc, model = train_model(model,
criterion,
optimizer,
scheduler,
2,
dataloaders,
dataset_sizes,
device=device,
multi_gpu=False)
optimizer = optim.SGD(filter(lambda p: p.requires_grad,
model.parameters()),
lr=.001,
momentum=0.9,
weight_decay=5e-4)
scheduler = lr_scheduler.StepLR(optimizer,
step_size=int(epochs2 / 2),
gamma=0.2)
best_acc, model = train_model(model,
criterion,
optimizer,
scheduler,
epochs2,
dataloaders,
dataset_sizes,
device=device,
multi_gpu=False)
torch.save(model.state_dict(), str(save_path / model_name))
def main():
global args, best_prec1, data
if args.tensorboard: configure("runs/%s"%(args.name))
# Data loading code
normalize = 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.ToTensor(),
transforms.Lambda(lambda x: F.pad(
Variable(x.unsqueeze(0), requires_grad=False, volatile=True),
(4,4,4,4),mode='reflect').data.squeeze()),
transforms.ToPILImage(),
transforms.RandomCrop(32),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
])
else:
transform_train = transforms.Compose([
transforms.ToTensor(),
normalize,
])
transform_test = transforms.Compose([
transforms.ToTensor(),
normalize
])
# create model
print("Creating the model...")
model = WideResNet(args.layers, args.dataset == 'cifar10' and 10 or 100,
args.widen_factor, dropRate=args.droprate)
rank = MPI.COMM_WORLD.Get_rank()
args.rank = rank
args.seed += rank
_set_seed(args.seed)
print("Creating the DataLoader...")
kwargs = {'num_workers': args.num_workers}#, 'pin_memory': args.use_cuda}
assert(args.dataset == 'cifar10' or args.dataset == 'cifar100')
train_loader = torch.utils.data.DataLoader(
datasets.__dict__[args.dataset.upper()]('../data', train=True,
download=True,
transform=transform_train),
batch_size=args.batch_size, shuffle=True, **kwargs)
val_loader = torch.utils.data.DataLoader(
datasets.__dict__[args.dataset.upper()]('../data', train=False, transform=transform_test),
batch_size=args.batch_size, shuffle=False, **kwargs)
# get the number of model parameters
args.num_parameters = sum([p.data.nelement() for p in model.parameters()])
print('Number of model parameters: {}'.format(args.num_parameters))
# for training on multiple GPUs.
# Use CUDA_VISIBLE_DEVICES=0,1 to specify which GPUs to use
if args.use_cuda:
print("Moving the model to the GPU")
model = torch.nn.DataParallel(model, device_ids=device_ids)
model = 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']
best_prec1 = checkpoint['best_prec1']
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 optimizer
criterion = nn.CrossEntropyLoss()
if args.use_cuda:
criterion = criterion.cuda()
# optimizer = torch.optim.SGD(model.parameters(), args.lr,
# momentum=args.momentum, nesterov=args.nesterov,
# weight_decay=args.weight_decay)
# optimizer = torch.optim.ASGD(model.parameters(), args.lr)
# from distributed_opt import MiniBatchSGD
# import torch.distributed as dist
# rank = np.random.choice('gloo')
# TODO: pass kwargs to code
print('initing MiniBatchSGD')
print(list(model.parameters())[6].view(-1)[:3])
if args.code == 'sgd':
code = codings.svd.SVD(compress=False)
elif args.code == 'svd':
print("train.py, svd_rank =", args.svd_rank)
code = codings.svd.SVD(random_sample=args.svd_rescale, rank=args.svd_rank,
compress=args.compress)
elif args.code == 'qsgd':
code = codings.qsgd.QSGD(scheme='qsgd')
elif args.code == 'terngrad':
code = codings.qsgd.QSGD(scheme='terngrad')
elif args.code == 'qsvd':
code = codings.qsvd.QSVD(scheme=args.scheme, rank=args.svd_rank)
else:
raise ValueError('args.code not recognized')
names = [n for n, p in model.named_parameters()]
assert len(names) == len(set(names))
# optimizer = MPI_PS(model.named_parameters(), model.parameters(), args.lr,
# code=code, optim='adam',
# use_mpi=args.use_mpi, cuda=args.use_cuda)
optimizer = SGD(model.named_parameters(), model.parameters(), args.lr,
code=code, optim='sgd',
use_mpi=args.use_mpi, cuda=args.use_cuda, compress_level=args.compress_level)
data = []
train_data = []
train_time = 0
for epoch in range(args.start_epoch, args.epochs + 1):
print(f"epoch {epoch}")
adjust_learning_rate(optimizer, epoch+1)
# train for one epoch
start = time.time()
if epoch >= 0:
train_d = train(train_loader, model, criterion, optimizer, epoch)
else:
train_d = []
train_time += time.time() - start
train_data += [dict(datum, **vars(args)) for datum in train_d]
# evaluate on validation set
if epoch >= args.epochs:
train_datum = validate(train_loader, model, criterion, epoch)
else:
train_datum = {'acc_test': np.inf, 'loss_test': np.inf}
datum = validate(val_loader, model, criterion, epoch)
# train_datum = {'acc_train': 0.1, 'loss_train': 2.3}
data += [{'train_time': train_time,
'whole_train_acc': train_datum['acc_test'],
'whole_train_loss': train_datum['loss_test'],
'epoch': epoch + 1, **vars(args), **datum}]
if epoch > 0:
if len(data) > 1:
data[-1]['epoch_train_time'] = data[-1]['train_time'] - data[-2]['train_time']
for key in train_data[-1]:
values = [datum[key] for i, datum in enumerate(train_data)]
if 'time' in key:
data[-1]["epoch_" + key] = np.sum(values)
else:
data[-1]["epoch_" + key] = values[0]
df = pd.DataFrame(data)
train_df = pd.DataFrame(train_data)
if True:
time.sleep(1)
# Yes loss_test IS on train data. (look at what validate returns)
print('\n\nmin_train_loss', train_datum['loss_test'],
optimizer.steps, '\n\n')
time.sleep(1)
ids = [str(getattr(args, key)) for key in
['layers', 'lr', 'batch_size', 'compress', 'seed', 'num_workers',
'svd_rank', 'svd_rescale', 'use_mpi', 'qsgd', 'world_size',
'rank', 'code', 'scheme']]
_write_csv(df, id=f'-'.join(ids))
_write_csv(train_df, id=f'-'.join(ids) + '_train')
pprint({k: v for k, v in data[-1].items()
if k in ['svd_rank', 'svd_rescale', 'qsgd', 'compress']})
pprint({k: v for k, v in data[-1].items()
if k in ['train_time', 'num_workers', 'loss_test',
'acc_test', 'epoch', 'compress', 'svd_rank', 'qsgd'] or 'time' in k})
prec1 = datum['acc_test']
# remember best prec@1 and save checkpoint
is_best = prec1 > best_prec1
best_prec1 = max(prec1, best_prec1)
save_checkpoint({
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
}, is_best)
print('Best accuracy: ', best_prec1)
