-
Notifications
You must be signed in to change notification settings - Fork 14
/
mart.py
58 lines (45 loc) · 1.84 KB
/
mart.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
def mart_loss(model,
x_natural,
y,
optimizer,
step_size=0.007,
epsilon=0.031,
perturb_steps=10,
beta=6.0,
distance='l_inf'):
kl = nn.KLDivLoss(reduction='none')
model.eval()
batch_size = len(x_natural)
# generate adversarial example
x_adv = x_natural.detach() + 0.001 * torch.randn(x_natural.shape).cuda().detach()
if distance == 'l_inf':
for _ in range(perturb_steps):
x_adv.requires_grad_()
with torch.enable_grad():
loss_ce = F.cross_entropy(model(x_adv), y)
grad = torch.autograd.grad(loss_ce, [x_adv])[0]
x_adv = x_adv.detach() + step_size * torch.sign(grad.detach())
x_adv = torch.min(torch.max(x_adv, x_natural - epsilon), x_natural + epsilon)
x_adv = torch.clamp(x_adv, 0.0, 1.0)
else:
x_adv = torch.clamp(x_adv, 0.0, 1.0)
model.train()
x_adv = Variable(torch.clamp(x_adv, 0.0, 1.0), requires_grad=False)
# zero gradient
optimizer.zero_grad()
logits = model(x_natural)
logits_adv = model(x_adv)
adv_probs = F.softmax(logits_adv, dim=1)
tmp1 = torch.argsort(adv_probs, dim=1)[:, -2:]
new_y = torch.where(tmp1[:, -1] == y, tmp1[:, -2], tmp1[:, -1])
loss_adv = F.cross_entropy(logits_adv, y) + F.nll_loss(torch.log(1.0001 - adv_probs + 1e-12), new_y)
nat_probs = F.softmax(logits, dim=1)
true_probs = torch.gather(nat_probs, 1, (y.unsqueeze(1)).long()).squeeze()
loss_robust = (1.0 / batch_size) * torch.sum(
torch.sum(kl(torch.log(adv_probs + 1e-12), nat_probs), dim=1) * (1.0000001 - true_probs))
loss = loss_adv + float(beta) * loss_robust
return loss