-
Notifications
You must be signed in to change notification settings - Fork 193
/
tutorial_train_mnist.py
127 lines (109 loc) · 4.7 KB
/
tutorial_train_mnist.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
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
# Copyright (c) 2018-present, Royal Bank of Canada.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
#
from __future__ import print_function
import os
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from advertorch.context import ctx_noparamgrad_and_eval
from advertorch.test_utils import LeNet5
from advertorch_examples.utils import get_mnist_train_loader
from advertorch_examples.utils import get_mnist_test_loader
from advertorch_examples.utils import TRAINED_MODEL_PATH
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Train MNIST')
parser.add_argument('--seed', default=0, type=int)
parser.add_argument('--mode', default="cln", help="cln | adv")
parser.add_argument('--train_batch_size', default=50, type=int)
parser.add_argument('--test_batch_size', default=1000, type=int)
parser.add_argument('--log_interval', default=200, type=int)
args = parser.parse_args()
torch.manual_seed(args.seed)
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
if args.mode == "cln":
flag_advtrain = False
nb_epoch = 10
model_filename = "mnist_lenet5_clntrained.pt"
elif args.mode == "adv":
flag_advtrain = True
nb_epoch = 90
model_filename = "mnist_lenet5_advtrained.pt"
else:
raise
train_loader = get_mnist_train_loader(
batch_size=args.train_batch_size, shuffle=True)
test_loader = get_mnist_test_loader(
batch_size=args.test_batch_size, shuffle=False)
model = LeNet5()
model.to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-4)
if flag_advtrain:
from advertorch.attacks import LinfPGDAttack
adversary = LinfPGDAttack(
model, loss_fn=nn.CrossEntropyLoss(reduction="sum"), eps=0.3,
nb_iter=40, eps_iter=0.01, rand_init=True, clip_min=0.0,
clip_max=1.0, targeted=False)
for epoch in range(nb_epoch):
model.train()
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
ori = data
if flag_advtrain:
# when performing attack, the model needs to be in eval mode
# also the parameters should NOT be accumulating gradients
with ctx_noparamgrad_and_eval(model):
data = adversary.perturb(data, target)
optimizer.zero_grad()
output = model(data)
loss = F.cross_entropy(
output, target, reduction='elementwise_mean')
loss.backward()
optimizer.step()
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()))
model.eval()
test_clnloss = 0
clncorrect = 0
if flag_advtrain:
test_advloss = 0
advcorrect = 0
for clndata, target in test_loader:
clndata, target = clndata.to(device), target.to(device)
with torch.no_grad():
output = model(clndata)
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:
advdata = adversary.perturb(clndata, target)
with torch.no_grad():
output = model(advdata)
test_advloss += F.cross_entropy(
output, target, reduction='sum').item()
pred = output.max(1, keepdim=True)[1]
advcorrect += pred.eq(target.view_as(pred)).sum().item()
test_clnloss /= len(test_loader.dataset)
print('\nTest set: avg cln loss: {:.4f},'
' cln acc: {}/{} ({:.0f}%)\n'.format(
test_clnloss, clncorrect, len(test_loader.dataset),
100. * clncorrect / len(test_loader.dataset)))
if flag_advtrain:
test_advloss /= len(test_loader.dataset)
print('Test set: avg adv loss: {:.4f},'
' adv acc: {}/{} ({:.0f}%)\n'.format(
test_advloss, advcorrect, len(test_loader.dataset),
100. * advcorrect / len(test_loader.dataset)))
torch.save(
model.state_dict(),
os.path.join(TRAINED_MODEL_PATH, model_filename))