def main():
args = get_args()
logging.basicConfig(
format='[%(asctime)s] - %(message)s',
datefmt='%Y/%m/%d %H:%M:%S',
filename=f'{args.ename}.log',
level=logging.DEBUG)
logger.info(args)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
mnist_train = datasets.MNIST("../mnist-data", train=True, download=True, transform=transforms.ToTensor())
train_loader = torch.utils.data.DataLoader(mnist_train, batch_size=args.batch_size, shuffle=True)
model = mnist_net().cuda()
model.train()
opt = torch.optim.Adam(model.parameters(), lr=args.lr_max)
if args.lr_type == 'cyclic':
lr_schedule = lambda t: np.interp([t], [0, args.epochs * 2//5, args.epochs], [0, args.lr_max, 0])[0]
elif args.lr_type == 'flat':
lr_schedule = lambda t: args.lr_max
else:
raise ValueError('Unknown lr_type')
criterion = nn.CrossEntropyLoss()
logger.info('Epoch \t Time \t LR \t \t Train Loss \t Train Acc')
for epoch in range(args.epochs):
start_time = time.time()
train_loss = 0
train_acc = 0
train_n = 0
deltas = None
gradients = None
selected_attack = []
for i, (X, y) in enumerate(train_loader):
X, y = X.cuda(), y.cuda()
lr = lr_schedule(epoch + (i+1)/len(train_loader))
opt.param_groups[0].update(lr=lr)
if args.attack == 'fgsm':
if args.attack_type == 'random':
norms_list = ['linf', 'l2-scaled', 'l1']
epsilon_list = [0.3, 6.5, 180.0]
alpha_list = [0.375, 2.5, 225.0]
curr_norm = np.random.randint(len(norms_list))
selected_attack.append(norms_list[curr_norm])
delta = attack_fgsm(model, X, y, epsilon_list[curr_norm], alpha_list[curr_norm], norms_list[curr_norm], args.init)
elif args.attack_type == 'max' or args.attack_type == "avg" or args.attack_type == "avg_loss":
norms_list = ['linf', 'l2-scaled', 'l1']
epsilon_list = [0.3, 6.5, 180.0]
alpha_list = [0.375, 2.5, 225.0]
delta_linf = attack_fgsm(model, X, y, epsilon_list[0], alpha_list[0], norms_list[0], args.init)
delta_l1 = attack_fgsm(model, X, y, epsilon_list[1], alpha_list[1], norms_list[1], args.init)
delta_l2 = attack_fgsm(model, X, y, epsilon_list[2], alpha_list[2], norms_list[2], args.init)
else:
delta = attack_fgsm(model, X, y, args.epsilon, args.alpha, args.norm, args.init)
#if deltas != None:
# deltas = torch.cat((deltas, delta), dim=0)
#else:
# deltas = delta
elif args.attack == 'none':
delta = torch.zeros_like(X)
elif args.attack == 'pgd':
if args.attack_type == 'random':
norms_list = ['linf', 'l2-scaled', 'l1']
epsilon_list = [0.3, 6.5, 180.0]
alpha_list = [0.01, 0.03, 6.0]
curr_norm = np.random.randint(len(norms_list))
selected_attack.append(norms_list[curr_norm])
delta = attack_pgd(model, X, y, epsilon_list[curr_norm], alpha_list[curr_norm], args.attack_iters, args.restarts, norms_list[curr_norm], args.init)
elif args.attack_type == 'max' or args.attack_type == "avg" or args.attack_type == "avg_loss":
norms_list = ['linf', 'l2-scaled', 'l1']
epsilon_list = [0.3, 6.5, 180.0]
alpha_list = [0.01, 0.03, 6.0]
delta_linf = attack_pgd(model, X, y, epsilon_list[0], alpha_list[0], args.attack_iters, args.restarts, norms_list[0], args.init)
delta_l1 = attack_pgd(model, X, y, epsilon_list[1], alpha_list[1], args.attack_iters, args.restarts, norms_list[1], args.init)
delta_l2 = attack_pgd(model, X, y, epsilon_list[2], alpha_list[2], args.attack_iters, args.restarts, norms_list[2], args.init)
else:
delta = attack_pgd(model, X, y, args.epsilon, args.alpha, args.attack_iters, args.restarts, args.norm, args.init)
if args.attack_type == "single" or args.attack_type == "random":
output = model(torch.clamp(X + delta, 0, 1))
loss = criterion(output, y)
opt.zero_grad()
loss.backward()
opt.step()
elif args.attack_type == 'max':
output_linf = model(torch.clamp(X + delta_linf[:X.size(0)], 0, 1))
output_l1 = model(torch.clamp(X + delta_l1[:X.size(0)], 0, 1))
output_l2 = model(torch.clamp(X + delta_l2[:X.size(0)], 0, 1))
loss_linf = criterion(output_linf, y)
loss_l1 = criterion(output_l1, y)
loss_l2 = criterion(output_l2, y)
acc_linf = (output_linf.max(1)[1] == y).sum().item()
acc_l1 = (output_l1.max(1)[1] == y).sum().item()
acc_l2 = (output_l2.max(1)[1] == y).sum().item()
loss_list = [loss_l1, loss_l2, loss_linf]
acc_list = [acc_l1, acc_l2, acc_linf]
delta_list = [delta_l1, delta_l2, delta_linf]
max_loss = max(loss_list)
max_loss_index = loss_list.index(max_loss)
max_delta = delta_list[max_loss_index]
output = model(torch.clamp(X + max_delta[:X.size(0)], 0, 1))
loss = criterion(output, y)
opt.zero_grad()
loss.backward()
nn.utils.clip_grad_norm_(model.parameters(), 0.5)
opt.step()
elif args.attack_type == "avg_loss":
output = model(torch.clamp(X + delta_linf[:X.size(0)], 0, 1))
loss_linf = criterion(output, y)
output_l1 = model(torch.clamp(X + delta_l1[:X.size(0)], 0, 1))
loss_l1 = criterion(output_l1, y)
output_l2 = model(torch.clamp(X + delta_l2[:X.size(0)], 0, 1))
loss_l2 = criterion(output_l2, y)
loss = (loss_linf + loss_l1 + loss_l2)/3
opt.zero_grad()
# with amp.scale_loss(loss, opt) as scaled_loss:
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)
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(model.state_dict(), args.fname)
if deltas != None:
torch.save(deltas.cpu(), args.fname+'_deltas')
if gradients != None:
torch.save(gradients.cpu(), args.fname+'_gradients')
def main():
args = get_args()
logging.basicConfig(
format=
'[%(asctime)s %(filename)s %(name)s %(levelname)s] - %(message)s',
datefmt='%Y/%m/%d %H:%M:%S',
filename=f'{args.ename}.log',
level=logging.DEBUG)
logger.info(args)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
mnist_test = datasets.MNIST("../mnist-data",
train=False,
download=True,
transform=transforms.ToTensor())
test_loader = torch.utils.data.DataLoader(mnist_test,
batch_size=args.batch_size,
shuffle=False)
model = mnist_net().cuda()
checkpoint = torch.load(args.fname)
model.load_state_dict(checkpoint)
model.eval()
total_loss = 0
total_acc = 0
n = 0
deltas = None
if args.attack == 'none':
with torch.no_grad():
for i, (X, y) in enumerate(test_loader):
X, y = X.cuda(), y.cuda()
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 i, (X, y) in enumerate(test_loader):
X, y = X.cuda(), y.cuda()
if args.attack == 'pgd':
delta = attack_pgd(model, X, y, args.epsilon, args.alpha,
args.attack_iters, args.restarts, args.norm,
args.init)
# store deltas for viz
if deltas != None:
deltas = torch.cat((deltas, delta))
else:
deltas = delta
elif args.attack == 'fgsm':
delta = attack_fgsm(model, X, y, args.epsilon, args.alpha,
args.norm, args.init)
with torch.no_grad():
#plot(X, delta, y)
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)
if deltas != None:
torch.save(deltas.cpu(), args.fname + '_eval_deltas')
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)