-
Notifications
You must be signed in to change notification settings - Fork 0
/
fgsm.py
59 lines (42 loc) · 1.49 KB
/
fgsm.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
59
import argparse
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms
import numpy as np
from PIL import Image
import random
import model
def save_image(x, filename):
img = transforms.ToPILImage()(x.view(-1,28,28))
img = img.convert('RGB')
img.save(filename)
def generate_adversarial_examples(args):
eps = args.eps
# data
data, target = np.load('./data/image.npy'), np.load('./data/labels.npy')
i = random.randint(0, len(data)-1)
x, y = torch.from_numpy(data[i]).view(-1,1,28,28), torch.from_numpy(np.array([target[i]]))
x.requires_grad=True
# model
net = model.MnistClassifer()
param = torch.load('mnist_model_params.pth', map_location='cpu')
net.load_state_dict(param)
y_pred = net(x)
creterion = nn.CrossEntropyLoss()
loss = creterion(y_pred, y)
loss.backward()
# fast gradient sign method
eta = eps * np.sign(x.grad.view(-1))
adv = x.data + eta.view(-1,1,28,28)
adv_y_pred = net(adv)
save_image(x, 'result/x.png')
save_image(x.grad, 'result/noize.png')
save_image(adv, 'result/adv_sample.png')
print('y: ', int(y), ' y_pred: ', int(y_pred.max(1)[1]), ' adv_y_pred: ', int(adv_y_pred.max(1)[1]))
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='MNIST TRAINING')
parser.add_argument('--eps', type=int, default=0.001,
help='Epsilon')
args = parser.parse_args()
generate_adversarial_examples(args)