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