def test_pgd(args, model, device, test_loader, epsilon=0.063):
model.eval()
model.reset()
adversary = LinfPGDAttack(model.forward_adv,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
eps=epsilon,
nb_iter=args.nb_iter,
eps_iter=args.eps_iter,
rand_init=True,
clip_min=-1.0,
clip_max=1.0,
targeted=False)
correct = 0
for batch_idx, (data, target) in enumerate(test_loader):
data, target = data.to(device), target.to(device)
model.reset()
with ctx_noparamgrad_and_eval(model):
adv_images = adversary.perturb(data, target)
output = model.run_cycles(adv_images)
pred = output.argmax(dim=1, keepdim=True)
correct += pred.eq(target.view_as(pred)).sum().item()
acc = correct / len(test_loader.dataset)
print('PGD attack Acc {:.3f}'.format(100. * acc))
return acc
def eval_test(model, device, test_loader):
model.eval()
test_loss = 0
correct = 0
adv_correct = 0
for data, target in test_loader:
data, target = data.to(device), target.to(device)
with torch.no_grad():
output = model(data)
test_loss += F.cross_entropy(output, target, size_average=False).item()
pred = output.max(1, keepdim=True)[1]
correct += pred.eq(target.view_as(pred)).sum().item()
with ctx_noparamgrad_and_eval(model):
adv_data = PGD_adversary.perturb(data, target)
with torch.no_grad():
output = model(adv_data)
pred = output.max(1, keepdim=True)[1]
adv_correct += pred.eq(target.view_as(pred)).sum().item()
test_loss /= len(test_loader.dataset)
print('Test: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)'.format(
test_loss, correct, len(test_loader.dataset),
100. * correct / len(test_loader.dataset)))
print('Test Adv Accuracy: {}/{} ({:.0f}%)'.format(
adv_correct, len(test_loader.dataset),
100. * adv_correct / len(test_loader.dataset)))
test_accuracy = correct / len(test_loader.dataset)
return test_loss, test_accuracy
def train_one_epoch(model, dataloader, epoch):
model.train()
n_total = 0
n_correct = 0
for i, (img, label) in enumerate(dataloader):
# training model using source data
img = img.expand(img.data.shape[0], 3, 28, 28)
img = img.cuda()
label = label.cuda()
with ctx_noparamgrad_and_eval(model):
img_adv = adversary_train.perturb(img, label)
output = model(input_data=img_adv)
loss = loss_func(output, label)
pred = output.data.max(1, keepdim=True)[1]
n_correct += pred.eq(label.data.view_as(pred)).cpu().sum()
n_total += batch_size
optimizer.zero_grad()
loss.backward()
optimizer.step()
if i % 100 == 0:
print('Epoch: {}, Iter: {}, Loss: {}'.format(epoch, i + 1, loss))
accu = n_correct.data.numpy() * 1.0 / n_total
print('Epoch: {}, Train Acc: {}'.format(epoch, accu))
return accu
def test(model, dataloader, epoch):
""" training """
model = model.eval()
model = model.cuda()
# i = 0
n_total = 0
n_adv_correct = 0
n_correct = 0
for img, label in dataloader:
batch_size = img.shape[0]
img = img.expand(img.data.shape[0], 3, 28, 28)
img = img.cuda()
label = label.cuda()
with ctx_noparamgrad_and_eval(model):
img_adv = adversary_test.perturb(img, label)
adv_output = model(input_data=img_adv)
pred = adv_output.data.max(1, keepdim=True)[1]
n_adv_correct += pred.eq(label.data.view_as(pred)).cpu().sum()
output = model(input_data=img)
pred = output.data.max(1, keepdim=True)[1]
n_correct += pred.eq(label.data.view_as(pred)).cpu().sum()
n_total += batch_size
accu = n_correct.data.numpy() * 1.0 / n_total
print('Epoch: {}, Test Acc: {}'.format(epoch, accu))
accu = n_adv_correct.data.numpy() * 1.0 / n_total
print('Epoch: {}, Adv Test Acc: {}'.format(epoch, accu))
def train(epoch):
print('\nEpoch: %d' % epoch)
net.train()
total_train_loss = 0
total_grad_loss = 0
correct = 0
total = 0
for batch_idx, (inputs, targets) in enumerate(trainloader):
inputs, targets = inputs.to(device), targets.to(device)
# inputs += torch.zeros_like(inputs).uniform_(-8/255, 8/255)
if args.adv:
with ctx_noparamgrad_and_eval(net):
train_im = adversary.perturb(inputs, targets).detach()
else:
train_im = inputs
loss, grad_loss, outputs = calc_grad(train_im, targets, net)
optimizer.zero_grad()
(loss+grad_loss).backward()
optimizer.step()
# writer.add_scalar("loss", loss.item(), len(trainloader)*epoch+batch_idx)
# writer.add_scalar("grad loss", grad_loss.item(), len(trainloader)*epoch+batch_idx)
total_train_loss += loss.item()
total_grad_loss += grad_loss.item()
_, predicted = outputs.max(1)
total += targets.size(0)
correct += predicted.eq(targets).sum().item()
progress_bar(batch_idx, len(trainloader), 'Loss: %.3f | Grad Loss: %.6f | Acc: %.3f%% (%d/%d)'
% (total_train_loss/(batch_idx+1), total_grad_loss/(batch_idx+1), 100.*correct/total, correct, total))
def test_one_epoch(self):
print("Evaluating on {}, epoch {}".format(self.dataname, self.epochs))
self.model.eval()
self.model.to(self.device)
self.reset_epoch_meters()
self.reset_disp_meters()
for data, idx in self.loader:
data, idx = data.to(self.device), idx.to(self.device)
target = self.loader.targets[idx]
with ctx_noparamgrad_and_eval(self.model):
# this advdata is a fixed eps adv, not scaled
advdata, curr_eps = self.adversary.perturb(data, target)
update_eps_meter(self.eps_meter, curr_eps.mean().item(), len(data))
self.update_eps(curr_eps, idx)
clnloss, clnacc = self.predict_then_update_loss_acc_meter(
self.cln_meter, data, target)
advloss, advacc = self.predict_then_update_loss_acc_meter(
self.adv_meter, advdata, target)
self.epochs += 1
self.print_disp_meters()
self.disp_eps_hist()
return (self.meters['cln']['epoch_acc'].avg,
self.meters['adv']['epoch_acc'].avg,
np.array(list(self.dct_eps.values())).mean())
def eval_fn(model):
advcorrect = 0
model.to(args.dev)
model.eval()
with ctx_noparamgrad_and_eval(model):
if args.source_arch == 'googlenet':
adv_complete_list = []
for batch_idx, (x_batch, y_batch) in enumerate(test_loader):
if (batch_idx + 1) * args.test_batch_size > args.batch_size:
break
x_batch, y_batch = x_batch.to(args.dev), y_batch.to(args.dev)
adv_complete_list.append(attacker.perturb(x_batch, target=y_batch))
adv_complete = torch.cat(adv_complete_list)
else:
adv_complete = attacker.perturb(x_test[:args.batch_size],
target=y_test[:args.batch_size])
adv_complete = torch.clamp(adv_complete, min=0., max=1.0)
output = model(adv_complete)
pred = output.max(1, keepdim=True)[1]
advcorrect += pred.eq(y_test[:args.batch_size].view_as(pred)).sum().item()
fool_rate = 1 - advcorrect / float(args.batch_size)
print('Test set base model fool rate: %f' %(fool_rate))
model.cpu()
if args.transfer:
adv_img_list = []
y_orig = y_test[:args.batch_size]
for i in range(0, len(adv_complete)):
adv_img_list.append([adv_complete[i].unsqueeze(0), y_orig[i]])
# Free memory
del model
torch.cuda.empty_cache()
baseline_transfer(args, attacker, "AEG", model_type,
adv_img_list, l_test_classif_paths, adv_models)
def test(epoch):
progress_bar = tqdm(testloader)
net.eval()
acc_clean = AverageMeter()
acc_adv = AverageMeter()
for batch_idx, (images, labels) in enumerate(progress_bar):
images, labels = images.cuda(), labels.cuda()
with ctx_noparamgrad_and_eval(net):
images_adv = test_attacker.perturb(images, labels)
pred_clean = net(images).argmax(dim=1)
pred_adv = net(images_adv).argmax(dim=1)
acc_clean.update((pred_clean == labels).float().mean().item() * 100.0,
images.size(0))
acc_adv.update((pred_adv == labels).float().mean().item() * 100.0,
images.size(0))
progress_bar.set_description(
'Test Epoch: [{0}] '
'Clean Acc: {acc_clean.val:.3f} ({acc_clean.avg:.3f}) '
'Adv Acc: {acc_adv.val:.3f} ({acc_adv.avg:.3f}) '.format(
epoch, acc_clean=acc_clean, acc_adv=acc_adv))
logging.info(
f'Epoch: {epoch} | Clean: {acc_clean.avg:.2f} % | Adv: {acc_adv.avg:.2f} %'
)
def attack_pgd_transfer(self, model_attacker, clean_loader, epsilon=0.1, eps_iter=0.02, test='average', nb_iter=7):
""" Use adversarial samples generated against model_attacker to attack the current model. """
self.model.eval()
self.model.reset()
model_attacker.eval()
model_attacker.reset()
adversary = LinfPGDAttack(
model_attacker.forward_adv, loss_fn=nn.CrossEntropyLoss(reduction="sum"), eps=epsilon,
nb_iter=nb_iter, eps_iter=eps_iter, rand_init=True, clip_min=-1.0, clip_max=1.0, targeted=False)
correct = 0
for batch_idx, (data, target) in enumerate(clean_loader):
data, target = data.to(self.device), target.to(self.device)
self.model.reset()
model_attacker.reset()
with ctx_noparamgrad_and_eval(model_attacker):
adv_images = adversary.perturb(data, target)
if(test=='last'):
output = self.model.run_cycles(adv_images)
elif(test=='average'):
output = self.model.run_average(adv_images)
else:
self.model.reset()
output = self.model(adv_images)
pred = output.argmax(dim=1, keepdim=True)
correct += pred.eq(target.view_as(pred)).sum().item()
acc = correct / len(clean_loader.dataset)
print('PGD attack Acc {:.3f}'.format(100. * acc))
return acc
def perturb(self, x, y):
with ctx_noparamgrad_and_eval(self.predict):
xadv = self.pgdadv.perturb(x, y)
unitptb, curr_eps = self._get_unitptb_and_eps(xadv, x, y,
self.pgdadv.eps)
return xadv, curr_eps
def train(train_loader, model, criterion, optimizer, epoch, adversary, args):
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
# switch to train mode
model.train()
end = time.time()
for i, (input, target) in enumerate(train_loader):
# measure data loading time
data_time.update(time.time() - end)
if args.gpu is not None:
input = input.cuda(args.gpu, non_blocking=True)
target = target.cuda(args.gpu, non_blocking=True)
if args.advtrain:
with ctx_noparamgrad_and_eval(model):
input = adversary.perturb(input, target)
# compute output
output = model(input)
loss = criterion(output, target)
# measure accuracy and record loss
acc1, acc5 = accuracy(output, target, topk=(1, 5))
losses.update(loss.item(), input.size(0))
top1.update(acc1[0], input.size(0))
top5.update(acc5[0], input.size(0))
# compute gradient and do SGD step
optimizer.zero_grad()
loss.backward()
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % args.print_freq == 0:
print('Epoch: [{0}][{1}/{2}]\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
'Acc@1 {top1.val:.3f} ({top1.avg:.3f})\t'
'Acc@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
epoch,
i,
len(train_loader),
batch_time=batch_time,
data_time=data_time,
loss=losses,
top1=top1,
top5=top5))
def test(model,
device,
test_loader,
epoch=None,
attack=None,
summary=None,
name=None):
model.eval()
test_loss = 0
correct = 0
with ctx_noparamgrad_and_eval(model):
for batch_idx, (data, target) in enumerate(test_loader):
data, target = data.to(device), target.to(device)
first_img = data[0]
if attack:
data = attack.perturb(data, target)
# plt.imshow(example_img.squeeze(0).cpu().numpy(), cmap='gray')
# plt.show()
output = model(data)
test_loss += F.nll_loss(
output, target, reduction="sum").item() # sum up batch loss
pred = output.argmax(
dim=1,
keepdim=True) # get the index of the max log-probability
if attack and summary:
example_img = data[0]
example_label = target[0].cpu().item()
example_pred = pred[0].cpu().item()
summary.add_image(
f"{name}_adversarial_image_eps{epoch:.3f}_l{example_label}_p{example_pred}",
example_img, 0)
summary.add_image(
f"{name}_adversarial_image_eps{epoch:.3f}_l{example_label}_p{example_pred}",
first_img, 1)
correct += pred.eq(target.view_as(pred)).sum().item()
# Average loss
test_loss /= len(test_loader.dataset)
acc = correct / len(test_loader.dataset)
print("\nTest set: Average loss: {:.4f}, {}/{} ({:.4f})\n".format(
test_loss,
correct,
len(test_loader.dataset),
acc,
))
if summary and (attack is None):
summary.add_scalar("test_loss", test_loss, epoch)
summary.add_scalar("test_acc", acc, epoch)
# Returns loss and accuracy
return test_loss, acc
def forward(self, X, mask, y):
if self.attack == 'fgsm':
with ctx_noparamgrad_and_eval(self.model):
delta = torch.zeros_like(X).uniform_(
-self.eps, self.eps)
delta.requires_grad = True
output = self.model(X + delta)
loss = self.loss_fn(output, y)
loss.backward()
grad = delta.grad.detach()
delta.data = torch.clamp(delta + self.alpha * torch.sign(grad),
-self.eps,
self.eps)
delta.data = torch.max(torch.min(
self.clip_max - X, delta.data), self.clip_min - X)
delta.data[(mask == 1).expand_as(delta)] = 0.
delta = delta.detach()
elif self.attack == 'pgd':
delta = torch.zeros_like(X).uniform_(
-self.eps, self.eps)
delta.data = torch.max(torch.min(
self.clip_max - X, delta.data), self.clip_min - X)
delta.requires_grad = True
with ctx_noparamgrad_and_eval(self.model):
for _ in range(self.nb_iter):
delta.data[(mask == 1).expand_as(delta)] = 0.
output = self.model(X + delta)
loss = self.loss_fn(output, y)
loss.backward()
grad = delta.grad.detach()
I = output.max(1)[1] == y
delta.data[I] = torch.clamp(
delta + self.alpha * torch.sign(grad),
-self.eps,
self.eps)[I]
delta.data[I] = torch.max(torch.min(
self.clip_max - X, delta.data), self.clip_min - X)[I]
delta.data[(mask == 1).expand_as(delta)] = 0.
delta = delta.detach()
return X + delta
def perturb(self, x, y, prev_eps):
self.pgdadv.eps = prev_eps
self.pgdadv.eps_iter = self.scale_eps_iter(self.pgdadv.eps,
self.pgdadv.nb_iter)
with ctx_noparamgrad_and_eval(self.predict):
xadv = self.pgdadv.perturb(x, y)
unitptb, curr_eps = self._get_unitptb_and_eps(xadv, x, y, prev_eps)
xadv = x + batch_multiply(curr_eps, unitptb)
return xadv, curr_eps
def attack_spsa(self,
clean_loader,
epsilon=0.1,
test='average',
ete=False,
nb_iter=7):
""" Use SPSA to attack the model. """
self.model.eval()
self.model.reset()
if (ete == False):
adv_func = self.model.forward_adv
else:
adv_func = self.model.run_cycles_adv
adversary = LinfSPSAAttack(
adv_func,
loss_fn=nn.CrossEntropyLoss(reduction="none"),
eps=epsilon,
nb_iter=nb_iter,
delta=0.01,
nb_sample=128,
max_batch_size=64,
clip_min=-1.0,
clip_max=1.0,
targeted=False)
correct = 0
numofdata = 0
for batch_idx, (data, target) in enumerate(clean_loader):
# To speed up the evaluation of attack, evaluate the first 10 batches
if (batch_idx < 10):
data, target = data.to(self.device), target.to(self.device)
self.model.reset()
with ctx_noparamgrad_and_eval(self.model):
adv_images = adversary.perturb(data, target)
if (test == 'last'):
output = self.model.run_cycles(adv_images)
elif (test == 'average'):
output = self.model.run_average(adv_images)
else:
self.model.reset()
output = self.model(adv_images)
pred = output.argmax(dim=1, keepdim=True)
correct += pred.eq(target.view_as(pred)).sum().item()
numofdata += data.shape[0]
acc = correct / numofdata
print('SPSA attack Acc {:.3f}'.format(100. * acc))
return acc
def train(epoch):
print('\nEpoch: %d' % epoch)
global Train_acc
net.train()
train_loss = 0
correct = 0
total = 0
if epoch > learning_rate_decay_start and learning_rate_decay_start >= 0:
frac = (epoch - learning_rate_decay_start) // learning_rate_decay_every
decay_factor = learning_rate_decay_rate**frac
current_lr = opt.lr * decay_factor
utils.set_lr(optimizer, current_lr) # set the decayed rate
else:
current_lr = opt.lr
print('learning_rate: %s' % str(current_lr))
for batch_idx, (inputs, targets) in enumerate(trainloader):
if use_cuda:
inputs, targets = inputs.cuda(), targets.cuda()
ori = inputs
# generate adversarial samples
with ctx_noparamgrad_and_eval(net):
inputs = adversary.perturb(inputs, targets)
# concatenate with clean samples
inputs = torch.cat((ori, inputs), 0)
targets = torch.cat((targets, targets), 0)
optimizer.zero_grad()
inputs, targets = Variable(inputs), Variable(targets)
outputs = net(inputs)
loss = criterion(outputs, targets)
loss.backward()
utils.clip_gradient(optimizer, 0.1)
optimizer.step()
train_loss += loss.data[0]
_, predicted = torch.max(outputs.data, 1)
total += targets.size(0)
correct += predicted.eq(targets.data).cpu().sum()
utils.progress_bar(
batch_idx, len(trainloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)' %
(train_loss /
(batch_idx + 1), 100. * correct / total, correct, total))
Train_acc = 100. * correct / total
Train_loss = train_loss / len(trainloader.dataset)
return Train_acc, Train_loss
def perturb(self, x, y, target_y=None):
with ctx_noparamgrad_and_eval(self.predict):
if self.pgdadv.targeted:
self.target_y = target_y
xadv = self.pgdadv.perturb(x, target_y)
adv_pred = self.pgdadv.predict(xadv).argmax(1)
# print((adv_pred == target_y).float().mean())
else:
xadv = self.pgdadv.perturb(x, y)
# print(self.pgdadv.eps, x.shape, xadv.shape, torch.norm((x-xadv).view(x.shape[0],-1), p=float('inf'), dim=1).mean())
unitptb, curr_eps = self._get_unitptb_and_eps(
xadv, x, y, self.pgdadv.eps)
xadv = clamp(x + batch_multiply(curr_eps, unitptb),
min=self.pgdadv.clip_min, max=self.pgdadv.clip_max)
# print('')
return xadv
def train_step(engine, batch):
model.train()
x, y = batch
x = x.to(device)
y = y.to(device) - min_y_train
with ctx_noparamgrad_and_eval(model):
x_adv = attack.perturb(x, y)
optimizer.zero_grad()
x = torch.cat((x, x_adv))
y = torch.cat((y, y))
ans = model.forward(x)
l = loss(ans, y)
optimizer.zero_grad()
l.backward()
optimizer.step()
# return ans, y
return l.item()
def test(epoch, is_adv=False):
global is_training, best_acc
is_training = False
net.eval()
test_loss = 0
correct = 0
total = 0
# with torch.no_grad():
for batch_idx, (inputs, targets) in enumerate(test_loader):
if use_cuda:
inputs, targets = inputs.requires_grad_().cuda(), targets.cuda()
if is_adv:
adversary = PGDAttack(net,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
eps=args.eps,
nb_iter=args.iter,
eps_iter=args.eps / args.iter,
rand_init=True,
clip_min=0.0,
clip_max=1.0,
targeted=False)
with ctx_noparamgrad_and_eval(net):
inputs = adversary.perturb(inputs, targets)
optimizer.zero_grad()
outputs = net(inputs)
loss = criterion(outputs, targets)
test_loss += loss.item()
_, predicted = torch.max(outputs.data, 1)
total += targets.size(0)
correct += predicted.eq(targets.data).cpu().sum()
correct = correct.item()
progress_bar(
batch_idx, len(test_loader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)' %
(test_loss /
(batch_idx + 1), 100. * correct / total, correct, total))
# Save checkpoint.
acc = 100. * correct / total
if acc > best_acc:
best_acc = acc
checkpoint(acc, epoch)
return (test_loss / batch_idx, 100. * correct / total)
def train(epoch):
global is_training
is_training = True
print('\nEpoch: %d' % epoch)
net.train()
train_loss = 0
correct = 0
total = 0
for batch_idx, (inputs, targets) in enumerate(train_loader):
if use_cuda:
inputs, targets = inputs.cuda().requires_grad_(), targets.cuda()
# generate adv img
if args.adv_train:
adversary = PGDAttack(net,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
eps=args.eps,
nb_iter=args.iter,
eps_iter=args.eps / args.iter,
rand_init=True,
clip_min=0.0,
clip_max=1.0,
targeted=False)
with ctx_noparamgrad_and_eval(net):
inputs = adversary.perturb(inputs, targets)
optimizer.zero_grad()
outputs = net(inputs)
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()
train_loss += loss.item()
_, predicted = torch.max(outputs.data, 1)
total += targets.size(0)
correct += predicted.eq(targets.data).cpu().sum()
correct = correct.item()
progress_bar(
batch_idx, len(train_loader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)' %
(train_loss /
(batch_idx + 1), 100. * correct / total, correct, total))
return (train_loss / batch_idx, 100. * correct / total)
def train(args, model, device, train_loader, optimizer, epoch):
model.train()
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
# attack
with ctx_noparamgrad_and_eval(model):
adv_data = adversary.perturb(data, target)
loss = F.cross_entropy(model(adv_data),
target,
reduction='elementwise_mean')
loss.backward()
optimizer.step()
# print progress
if batch_idx % args.log_interval == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item()))
def train(self, lr=1e-4, betas=(0.5, 0.9)):
model = self.model
optimizer = torch.optim.Adam(model.parameters(), lr=lr, betas=betas)
for epoch in range(self.iters):
for batch_idx, (data, target) in enumerate(train_loader):
data = data.to(device)
target = target.to(device)
optimizer.zero_grad()
real_data = data
real_data_v = autograd.Variable(real_data).to(device)
model.zero_grad()
# train with real
D_real = model(real_data_v)
D_real = D_real.mean()
with ctx_noparamgrad_and_eval(model):
fake_data = adversary.perturb(data, target)
# fake_data = adversary.perturb(data.cpu(), target.cpu())
fake_data_v = autograd.Variable(fake_data).to(device)
# train with fake
D_fake = model(fake_data_v)
D_fake = D_fake.mean()
gradient_penalty = calc_gradient_penalty(
model, real_data_v.data, fake_data_v.data)
loss = D_fake - D_real + gradient_penalty
loss.backward()
Wasserstein_D = D_real - D_fake
optimizer.step()
self.losses['loss'].append(loss.item())
self.losses['w_dist'].append(Wasserstein_D.item())
if batch_idx % args.log_interval == 0:
print(
'Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}\tW_dist: {:.6f}'
.format(epoch, batch_idx * len(data),
len(train_loader.dataset),
100. * batch_idx / len(train_loader),
loss.item(), Wasserstein_D.item()))
# 每 10 batch 保存一下
if epoch % 20 == 0:
self.save_model(epoch)
self.save_loss(epoch)
self.save_model('final')
self.save_loss('final')
def eval(args, attacker, attack_name, eval_helpers, num_eval_samples=None):
model, model_type, adv_models, l_test_classif_paths, test_loader = eval_helpers[:]
advcorrect = 0
with ctx_noparamgrad_and_eval(model):
adv_complete_list, output_list, y_list = [], [], []
for batch_idx, (x_batch, y_batch) in enumerate(test_loader):
x_batch, y_batch = x_batch.to(args.dev), y_batch.to(args.dev)
adv_imgs = attacker.perturb(x_batch, y_batch).detach()
adv_complete_list.append(adv_imgs.cpu())
output_list.append(model(adv_imgs).cpu())
y_list.append(y_batch.cpu())
output = torch.cat(output_list)
adv_complete = torch.cat(adv_complete_list)
y_test = torch.cat(y_list)
pred = output.max(1, keepdim=True)[1]
advcorrect += pred.eq(y_test.view_as(pred)).sum().item()
fool_rate = 1 - advcorrect / float(len(test_loader.dataset))
print('Test set base model fool rate: %f' % (fool_rate))
if args.transfer:
if num_eval_samples is not None:
adv_complete = adv_complete[:num_eval_samples]
y_test = y_test[:num_eval_samples]
adv_img_list = torch.utils.data.TensorDataset(adv_complete, y_test)
adv_img_list = torch.utils.data.DataLoader(
adv_img_list, batch_size=args.test_batch_size)
# Free memory
del model
torch.cuda.empty_cache()
fool_rate = []
if args.target_arch is not None:
model_type = args.target_arch
if not isinstance(model_type, list):
model_type = [model_type]
for name in model_type:
fool_rate.append(
baseline_transfer(args, attacker, attack_name, name,
adv_img_list, l_test_classif_paths,
adv_models))
return np.mean(fool_rate)
return fool_rate
def train(args, model, device, cifar_nat_x,cifar_x,cifar_y, optimizer, epoch):
model.train()
num_of_example = 50000
batch_size = args.batch_size
cur_order = np.random.permutation(num_of_example)
iter_num = num_of_example // batch_size + (0 if num_of_example % batch_size == 0 else 1)
batch_idx = -batch_size
for i in range(iter_num):
# get shuffled data
batch_idx = (batch_idx + batch_size) if batch_idx + batch_size < num_of_example else 0
x_batch = cifar_x[cur_order[batch_idx:min(num_of_example, batch_idx + batch_size)]].to(device)
x_nat_batch = cifar_nat_x[cur_order[batch_idx:min(num_of_example, batch_idx + batch_size)]].to(device)
y_batch = cifar_y[cur_order[batch_idx:min(num_of_example, batch_idx + batch_size)]].to(device)
batch_size = y_batch.shape[0]
# attack
with ctx_noparamgrad_and_eval(model):
adv_data = adversary.perturb(x_nat_batch,y_batch,epoch,x_batch)
# update cifar_x
cifar_x[cur_order[batch_idx:min(num_of_example, batch_idx + batch_size)]] = adv_data.clone().detach().cpu()
# loss backward
optimizer.zero_grad()
loss = F.cross_entropy(model(adv_data),y_batch,reduction='elementwise_mean')
loss.backward()
optimizer.step()
# print max perturbation & min
if i == 0:
eta = adv_data - x_nat_batch
print("max:{}".format(torch.max(eta)))
print("min:{}".format(torch.min(eta)))
# print progress
if i % args.log_interval == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, i * len(x_batch), len(train_loader.dataset),
100. * i / len(train_loader), loss.item()))
def train(epoch):
progress_bar = tqdm(trainloader)
net.train()
gan.eval()
image_loss_meter = AverageMeter()
latent_loss_meter = AverageMeter()
total_loss_meter = AverageMeter()
for batch_idx, (images, latents, labels) in enumerate(progress_bar):
images, latents, labels = images.cuda(), latents.cuda(), labels.cuda()
lr = cfg.optimizer.args.lr
if lr_schedule is not None:
lr = lr_schedule(epoch + (batch_idx + 1) / len(trainloader))
for param_group in optimizer.param_groups:
param_group['lr'] = lr
with ctx_noparamgrad_and_eval(model):
images_adv = image_attacker.perturb(images, labels)
# latents_adv = latent_attacker.perturb(latents, labels)
# images_ladv = gan(latents_adv).detach()
optimizer.zero_grad()
total_loss = criterion(net(images_adv), labels)
# latent_loss = criterion(net(images_ladv), labels)
# total_loss = 0.5 * image_loss + 0.5 * latent_loss
total_loss.backward()
optimizer.step()
# image_loss_meter.update(image_loss.item(), images.size(0))
# latent_loss_meter.update(latent_loss.item(), images.size(0))
total_loss_meter.update(total_loss.item(), images.size(0))
progress_bar.set_description('Train Epoch: [{0}] | '
'Total Loss: {t_loss.val:.3f} | '
'lr: {1:.6f}'.format(
epoch, lr, t_loss=total_loss_meter))
def _run_batch(self, data, target, eval_metrics=False, train=True):
"""Runs batch optimization and returns loss
:param data: input batch
:param target: labels batch
:return: loss value
"""
# generate adversarial data
original_data = data
with ctx_noparamgrad_and_eval(self.model):
adversarial_data = self.adversary.perturb(data, target)
self.optimizer.zero_grad()
# get adversarial output and normal+adversarial loss
if train is True:
output = self.model(original_data)
adv_output = self.model(adversarial_data)
else:
with torch.no_grad():
output = self.model(original_data)
adv_output = self.model(adversarial_data)
adv_loss = self.loss(adv_output, target)
loss = self.loss(output, target)
if train is True:
if self.worst_case_training is False:
loss.backward()
self.optimizer.step()
adv_loss.backward()
self.optimizer.step()
if eval_metrics is True:
metrics = self.eval_metrics(output, adv_output, target)
return loss.item(), adv_loss.item(), metrics, self.get_metrics_dic(
metrics)
else:
return loss.item(), adv_loss.item()
def train_epoch(self, model: nn.Module, train_loader: DataLoader,
val_clean_loader: DataLoader, val_triggered_loader: DataLoader,
epoch_num: int, use_amp: bool = False):
"""
Runs one epoch of training on the specified model
:param model: the model to train for one epoch
:param train_loader: a DataLoader object pointing to the training dataset
:param val_clean_loader: a DataLoader object pointing to the validation dataset that is clean
:param val_triggered_loader: a DataLoader object pointing to the validation dataset that is triggered
:param epoch_num: the epoch number that is being trained
:param use_amp: if True, uses automated mixed precision for FP16 training.
:return: a list of statistics for batches where statistics were computed
"""
# Probability of Adversarial attack to occur in each iteration
attack_prob = self.optimizer_cfg.training_cfg.adv_training_ratio
pid = os.getpid()
train_dataset_len = len(train_loader.dataset)
loop = tqdm(train_loader, disable=self.optimizer_cfg.reporting_cfg.disable_progress_bar)
scaler = None
if use_amp:
scaler = torch.cuda.amp.GradScaler()
train_n_correct, train_n_total = None, None
# Define parameters of the adversarial attack
attack_eps = float(self.optimizer_cfg.training_cfg.adv_training_eps)
attack_iterations = int(self.optimizer_cfg.training_cfg.adv_training_iterations)
eps_iter = (2.0 * attack_eps) / float(attack_iterations)
attack = LinfPGDAttack(
predict=model,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
eps=attack_eps,
nb_iter=attack_iterations,
eps_iter=eps_iter)
sum_batchmean_train_loss = 0
running_train_acc = 0
num_batches = len(train_loader)
model.train()
for batch_idx, (x, y_truth) in enumerate(loop):
x = x.to(self.device)
y_truth = y_truth.to(self.device)
# put network into training mode & zero out previous gradient computations
self.optimizer.zero_grad()
# get predictions based on input & weights learned so far
if use_amp:
with torch.cuda.amp.autocast():
# add adversarial noise via l_inf PGD attack
# only apply attack to attack_prob of the batches
if attack_prob and np.random.rand() <= attack_prob:
with ctx_noparamgrad_and_eval(model):
x = attack.perturb(x, y_truth)
y_hat = model(x)
# compute metrics
batch_train_loss = self._eval_loss_function(y_hat, y_truth)
else:
# add adversarial noise vis lin PGD attack
if attack_prob and np.random.rand() <= attack_prob:
with ctx_noparamgrad_and_eval(model):
x = attack.perturb(x, y_truth)
y_hat = model(x)
batch_train_loss = self._eval_loss_function(y_hat, y_truth)
sum_batchmean_train_loss += batch_train_loss.item()
running_train_acc, train_n_total, train_n_correct = default_optimizer._running_eval_acc(y_hat, y_truth,
n_total=train_n_total,
n_correct=train_n_correct,
soft_to_hard_fn=self.soft_to_hard_fn,
soft_to_hard_fn_kwargs=self.soft_to_hard_fn_kwargs)
# compute gradient
if use_amp:
# Scales loss. Calls backward() on scaled loss to create scaled gradients.
# Backward passes under autocast are not recommended.
# Backward ops run in the same dtype autocast chose for corresponding forward ops.
scaler.scale(batch_train_loss).backward()
else:
if np.isnan(sum_batchmean_train_loss) or np.isnan(running_train_acc):
default_optimizer._save_nandata(x, y_hat, y_truth, batch_train_loss, sum_batchmean_train_loss, running_train_acc,
train_n_total, train_n_correct, model)
batch_train_loss.backward()
# perform gradient clipping if configured
if self.optimizer_cfg.training_cfg.clip_grad:
if use_amp:
# Unscales the gradients of optimizer's assigned params in-place
scaler.unscale_(self.optimizer)
if self.optimizer_cfg.training_cfg.clip_type == 'norm':
# clip_grad_norm_ modifies gradients in place
# see: https://pytorch.org/docs/stable/_modules/torch/nn/utils/clip_grad.html
torch_clip_grad.clip_grad_norm_(model.parameters(), self.optimizer_cfg.training_cfg.clip_val,
**self.optimizer_cfg.training_cfg.clip_kwargs)
elif self.optimizer_cfg.training_cfg.clip_type == 'val':
# clip_grad_val_ modifies gradients in place
# see: https://pytorch.org/docs/stable/_modules/torch/nn/utils/clip_grad.html
torch_clip_grad.clip_grad_value_(
model.parameters(), self.optimizer_cfg.training_cfg.clip_val)
else:
msg = "Unknown clipping type for gradient clipping!"
logger.error(msg)
raise ValueError(msg)
if use_amp:
# scaler.step() first unscales the gradients of the optimizer's assigned params.
# If these gradients do not contain infs or NaNs, optimizer.step() is then called,
# otherwise, optimizer.step() is skipped.
scaler.step(self.optimizer)
# Updates the scale for next iteration.
scaler.update()
else:
self.optimizer.step()
# report batch statistics to tensorflow
if self.tb_writer:
try:
batch_num = int(epoch_num * num_batches + batch_idx)
self.tb_writer.add_scalar(self.optimizer_cfg.reporting_cfg.experiment_name + '-train_loss',
batch_train_loss.item(), global_step=batch_num)
self.tb_writer.add_scalar(self.optimizer_cfg.reporting_cfg.experiment_name + '-running_train_acc',
running_train_acc, global_step=batch_num)
except:
# TODO: catch specific expcetions
pass
loop.set_description('Epoch {}/{}'.format(epoch_num + 1, self.num_epochs))
loop.set_postfix(avg_train_loss=batch_train_loss.item())
if batch_idx % self.num_batches_per_logmsg == 0:
logger.info('{}\tTrain Epoch: {} [{}/{} ({:.0f}%)]\tTrainLoss: {:.6f}\tTrainAcc: {:.6f}'.format(
pid, epoch_num, batch_idx * len(x), train_dataset_len,
100. * batch_idx / num_batches, batch_train_loss.item(), running_train_acc))
train_stats = EpochTrainStatistics(running_train_acc, sum_batchmean_train_loss / float(num_batches))
# if we have validation data, we compute on the validation dataset
num_val_batches_clean = len(val_clean_loader)
if num_val_batches_clean > 0:
logger.info('Running Validation on Clean Data')
running_val_clean_acc, _, _, val_clean_loss = \
default_optimizer._eval_acc(val_clean_loader, model, self.device,
self.soft_to_hard_fn, self.soft_to_hard_fn_kwargs, self._eval_loss_function)
else:
logger.info("No dataset computed for validation on clean dataset!")
running_val_clean_acc = None
val_clean_loss = None
num_val_batches_triggered = len(val_triggered_loader)
if num_val_batches_triggered > 0:
logger.info('Running Validation on Triggered Data')
running_val_triggered_acc, _, _, val_triggered_loss = \
default_optimizer._eval_acc(val_triggered_loader, model, self.device,
self.soft_to_hard_fn, self.soft_to_hard_fn_kwargs, self._eval_loss_function)
else:
logger.info(
"No dataset computed for validation on triggered dataset!")
running_val_triggered_acc = None
val_triggered_loss = None
validation_stats = EpochValidationStatistics(running_val_clean_acc, val_clean_loss,
running_val_triggered_acc, val_triggered_loss)
if num_val_batches_clean > 0:
logger.info('{}\tTrain Epoch: {} \tCleanValLoss: {:.6f}\tCleanValAcc: {:.6f}'.format(
pid, epoch_num, val_clean_loss, running_val_clean_acc))
if num_val_batches_triggered > 0:
logger.info('{}\tTrain Epoch: {} \tTriggeredValLoss: {:.6f}\tTriggeredValAcc: {:.6f}'.format(
pid, epoch_num, val_triggered_loss, running_val_triggered_acc))
if self.tb_writer:
try:
batch_num = int((epoch_num + 1) * num_batches)
if num_val_batches_clean > 0:
self.tb_writer.add_scalar(self.optimizer_cfg.reporting_cfg.experiment_name +
'-clean-val-loss', val_clean_loss, global_step=batch_num)
self.tb_writer.add_scalar(self.optimizer_cfg.reporting_cfg.experiment_name +
'-clean-val_acc', running_val_clean_acc, global_step=batch_num)
if num_val_batches_triggered > 0:
self.tb_writer.add_scalar(self.optimizer_cfg.reporting_cfg.experiment_name +
'-triggered-val-loss', val_triggered_loss, global_step=batch_num)
self.tb_writer.add_scalar(self.optimizer_cfg.reporting_cfg.experiment_name +
'-triggered-val_acc', running_val_triggered_acc, global_step=batch_num)
except:
pass
# update the lr-scheduler if necessary
if self.lr_scheduler is not None:
if self.optimizer_cfg.training_cfg.lr_scheduler_call_arg is None:
self.lr_scheduler.step()
elif self.optimizer_cfg.training_cfg.lr_scheduler_call_arg.lower() == 'val_acc':
val_acc = validation_stats.get_val_acc()
if val_acc is not None:
self.lr_scheduler.step(val_acc)
else:
msg = "val_clean_acc not defined b/c validation dataset is not defined! Ignoring LR step!"
logger.warning(msg)
elif self.optimizer_cfg.training_cfg.lr_scheduler_call_arg.lower() == 'val_loss':
val_loss = validation_stats.get_val_loss()
if val_loss is not None:
self.lr_scheduler.step(val_loss)
else:
msg = "val_clean_loss not defined b/c validation dataset is not defined! Ignoring LR step!"
logger.warning(msg)
else:
msg = "Unknown mode for calling lr_scheduler!"
logger.error(msg)
raise ValueError(msg)
return train_stats, validation_stats
def train_adv(args,
model,
device,
train_loader,
optimizer,
scheduler,
epoch,
cycles,
mse_parameter=1.0,
clean_parameter=1.0,
clean='supclean'):
model.train()
correct = 0
train_loss = 0.0
model.reset()
adversary = LinfPGDAttack(model,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
eps=args.eps,
nb_iter=args.nb_iter,
eps_iter=args.eps_iter,
rand_init=True,
clip_min=-1.0,
clip_max=1.0,
targeted=False)
print(len(train_loader))
for batch_idx, (images, targets) in enumerate(train_loader):
optimizer.zero_grad()
images = images.cuda()
targets = targets.cuda()
model.reset()
with ctx_noparamgrad_and_eval(model):
adv_images = adversary.perturb(images, targets)
images_all = torch.cat((images, adv_images), 0)
# Reset the model latent variables
model.reset()
if (args.dataset == 'cifar10'):
logits, orig_feature_all, block1_all, block2_all, block3_all = model(
images_all, first=True, inter=True)
elif (args.dataset == 'fashion'):
logits, orig_feature_all, block1_all, block2_all = model(
images_all, first=True, inter=True)
ff_prev = orig_feature_all
# f1 the original feature of clean images
orig_feature, _ = torch.split(orig_feature_all, images.size(0))
block1_clean, _ = torch.split(block1_all, images.size(0))
block2_clean, _ = torch.split(block2_all, images.size(0))
if (args.dataset == 'cifar10'):
block3_clean, _ = torch.split(block3_all, images.size(0))
logits_clean, logits_adv = torch.split(logits, images.size(0))
if not ('no' in clean):
loss = (clean_parameter * F.cross_entropy(logits_clean, targets) +
F.cross_entropy(logits_adv, targets)) / (2 * (cycles + 1))
else:
loss = F.cross_entropy(logits_adv, targets) / (cycles + 1)
for i_cycle in range(cycles):
if (args.dataset == 'cifar10'):
recon, block1_recon, block2_recon, block3_recon = model(
logits, step='backward', inter_recon=True)
elif (args.dataset == 'fashion'):
recon, block1_recon, block2_recon = model(logits,
step='backward',
inter_recon=True)
recon_clean, recon_adv = torch.split(recon, images.size(0))
recon_block1_clean, recon_block1_adv = torch.split(
block1_recon, images.size(0))
recon_block2_clean, recon_block2_adv = torch.split(
block2_recon, images.size(0))
if (args.dataset == 'cifar10'):
recon_block3_clean, recon_block3_adv = torch.split(
block3_recon, images.size(0))
loss += (F.mse_loss(recon_adv, orig_feature) +
F.mse_loss(recon_block1_adv, block1_clean) +
F.mse_loss(recon_block2_adv, block2_clean) +
F.mse_loss(recon_block3_adv, block3_clean)
) * mse_parameter / (4 * cycles)
elif (args.dataset == 'fashion'):
loss += (F.mse_loss(recon_adv, orig_feature) +
F.mse_loss(recon_block1_adv, block1_clean) +
F.mse_loss(recon_block2_adv, block2_clean)
) * mse_parameter / (3 * cycles)
# feedforward
ff_current = ff_prev + args.res_parameter * (recon - ff_prev)
logits = model(ff_current, first=False)
ff_prev = ff_current
logits_clean, logits_adv = torch.split(logits, images.size(0))
if not ('no' in clean):
loss += (
clean_parameter * F.cross_entropy(logits_clean, targets) +
F.cross_entropy(logits_adv, targets)) / (2 * (cycles + 1))
else:
loss += F.cross_entropy(logits_adv, targets) / (cycles + 1)
pred = logits_clean.argmax(
dim=1, keepdim=True) # get the index of the max log-probability
correct += pred.eq(targets.view_as(pred)).sum().item()
loss.backward()
if (args.grad_clip):
nn.utils.clip_grad_norm_(model.parameters(), 0.5)
optimizer.step()
scheduler.step()
train_loss += loss
if batch_idx % args.log_interval == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(images[0]), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item()))
train_loss /= len(train_loader)
acc = correct / len(train_loader.dataset)
return train_loss, acc
adversary = LinfPGDAttack(
net, loss_fn=nn.CrossEntropyLoss().cuda(), eps=16/255,
nb_iter=7, eps_iter=4/255, rand_init=True, clip_min=-1.0, clip_max=1.0,
targeted=False)
if args.alp:
criterion_alp = nn.MSELoss().cuda()
for epoch in range(1, args.nepoch+1):
net.train()
for batch_idx, (inputs, labels) in enumerate(trloader):
inputs_cls, labels_cls = inputs.cuda(), labels.cuda()
optimizer.zero_grad()
with ctx_noparamgrad_and_eval(net):
inputs_adv = adversary.perturb(inputs_cls, labels_cls)
if args.weight == 0:
outputs_adv = net(inputs_adv)
loss = criterion(outputs_adv, labels_cls)
else:
inputs_all = torch.cat([inputs_cls, inputs_adv], dim=0)
labels_all = torch.cat([labels_cls, labels_cls], dim=0)
outputs_all = net(inputs_all)
outputs_cls, outputs_adv = torch.split(outputs_all, inputs_cls.size(0), dim=0)
loss = criterion(outputs_cls, labels_cls)
if args.alp:
loss += args.weight * criterion_alp(outputs_cls, outputs_adv)
else:
pred = output.max(1, keepdim=True)[1]
clncorrect_nodefence += pred.eq(target.view_as(pred)).sum().item()
# clean data with defence
clndata_test_one = clndata
with torch.no_grad():
output = model(clndata_test_one.float())
test_clnloss += F.cross_entropy(output, target,
reduction='sum').item()
pred = output.max(1, keepdim=True)[1]
clncorrect += pred.eq(target.view_as(pred)).sum().item()
if flag_advtrain:
with ctx_noparamgrad_and_eval(model):
advdata = adversary.perturb(clndata, target)
# no defence
with torch.no_grad():
output = model(advdata.float())
test_advloss_nodefence += F.cross_entropy(
output, target, reduction='sum').item()
pred = output.max(1, keepdim=True)[1]
advcorrect_nodefence += pred.eq(
target.view_as(pred)).sum().item()
# with defence
# # gaussian_block
if args.gaussian_block:
