-
Notifications
You must be signed in to change notification settings - Fork 4
/
train_cifar.py
191 lines (164 loc) · 7.43 KB
/
train_cifar.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
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
from __future__ import print_function
import argparse
import copy
import time
import chainer
import chainer.links as L
import chainer.functions as F
from chainer import training
from chainer.training import extensions
from chainer import reporter
from chainer.datasets import get_cifar10
from chainer.datasets import get_cifar100
import nets
# VGG16VD
def main():
parser = argparse.ArgumentParser(description='Chainer CIFAR example:')
parser.add_argument('--dataset', '-d', default='cifar10',
help='The dataset to use: cifar10 or cifar100')
parser.add_argument('--batchsize', '-b', type=int, default=128,
help='Number of images in each mini-batch')
parser.add_argument('--epoch', '-e', type=int, default=100,
help='Number of sweeps over the dataset to train')
parser.add_argument('--gpu', '-g', type=int, default=0,
help='GPU ID (negative value indicates CPU)')
parser.add_argument('--out', '-o', default='result',
help='Directory to output the result')
parser.add_argument('--pretrain', default=0,
help='Pretrain (w/o VD) or not (w/ VD).' +
' default is not (0).')
parser.add_argument('--resume', '-r', default='',
help='Resume the training from snapshot')
parser.add_argument('--resume-opt', '-ro', default='',
help='Resume optimizer the training from snapshot')
args = parser.parse_args()
print('GPU: {}'.format(args.gpu))
print('# Minibatch-size: {}'.format(args.batchsize))
print('# epoch: {}'.format(args.epoch))
print('')
# Set up a neural network to train.
# Classifier reports softmax cross entropy loss and accuracy at every
# iteration, which will be used by the PrintReport extension below.
if args.dataset == 'cifar10':
print('Using CIFAR10 dataset.')
class_labels = 10
train, test = get_cifar10()
elif args.dataset == 'cifar100':
print('Using CIFAR100 dataset.')
class_labels = 100
train, test = get_cifar100()
else:
raise RuntimeError('Invalid dataset choice.')
print('# train:', len(train))
print('# test :', len(test))
if args.pretrain:
model = nets.VGG16(class_labels)
def calc_loss(x, t):
model.y = model(x)
model.loss = F.softmax_cross_entropy(model.y, t)
reporter.report({'loss': model.loss}, model)
model.accuracy = F.accuracy(model.y, t)
reporter.report({'accuracy': model.accuracy}, model)
return model.loss
model.calc_loss = calc_loss
model.use_raw_dropout = True
elif args.resume:
model = nets.VGG16VD(class_labels, warm_up=1.)
model(train[0][0][None, ]) # for setting in_channels automatically
model.to_variational_dropout()
chainer.serializers.load_npz(args.resume, model)
else:
model = nets.VGG16VD(class_labels, warm_up=0.0001)
model(train[0][0][None, ]) # for setting in_channels automatically
model.to_variational_dropout()
if args.gpu >= 0:
chainer.cuda.get_device(args.gpu).use() # Make a specified GPU current
model.to_gpu() # Copy the model to the GPU
if args.pretrain:
# Original Torch code (http://torch.ch/blog/2015/07/30/cifar.html)
# uses lr=1. However, it doesn't work well as people say in the post.
# This follows a version of Chainer example using lr=0.1.
optimizer = chainer.optimizers.MomentumSGD(0.1)
optimizer.setup(model)
optimizer.add_hook(chainer.optimizer.WeightDecay(5e-4))
elif args.resume:
optimizer = chainer.optimizers.Adam(1e-5)
optimizer.setup(model)
else:
optimizer = chainer.optimizers.Adam(1e-4)
optimizer.setup(model)
optimizer.add_hook(chainer.optimizer.GradientClipping(10.))
train_iter = chainer.iterators.SerialIterator(train, args.batchsize)
test_iter = chainer.iterators.SerialIterator(test, args.batchsize,
repeat=False, shuffle=False)
if args.resume:
classifier = L.Classifier(model.copy())
accuracy = extensions.Evaluator(
test_iter, classifier, device=args.gpu)()['main/accuracy']
print('test accuracy VD:', accuracy)
# Set up a trainer
updater = training.StandardUpdater(
train_iter, optimizer, device=args.gpu, loss_func=model.calc_loss)
trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out)
# Evaluate the model with the test dataset for each epoch
trainer.extend(extensions.Evaluator(test_iter, L.Classifier(model),
device=args.gpu))
if args.pretrain:
trainer.extend(extensions.ExponentialShift('lr', 0.5),
trigger=(25, 'epoch'))
elif not args.resume:
trainer.extend(extensions.LinearShift(
'alpha', (1e-4, 0.),
(0, args.epoch * len(train) // args.batchsize)))
# Take a snapshot at each epoch
# trainer.extend(extensions.snapshot(), trigger=(args.epoch, 'epoch'))
if args.pretrain:
trainer.extend(extensions.snapshot_object(
model, 'model_snapshot_{.updater.epoch}'),
trigger=(10, 'epoch'))
# Write a log of evaluation statistics for each epoch
# trainer.extend(extensions.LogReport())
per = min(len(train) // args.batchsize // 2, 1000)
trainer.extend(extensions.LogReport(trigger=(per, 'iteration')))
# Print selected entries of the log to stdout
# Here "main" refers to the target link of the "main" optimizer again, and
# "validation" refers to the default name of the Evaluator extension.
# Entries other than 'epoch' are reported by the Classifier link, called by
# either the updater or the evaluator.
trainer.extend(extensions.PrintReport(
['epoch', 'iteration', 'main/loss', 'validation/main/loss',
'main/accuracy', 'validation/main/accuracy',
'main/class', 'main/kl', 'main/mean_p', 'main/sparsity',
'main/W/Wnz', 'main/kl_coef',
'elapsed_time']))
# Print a progress bar to stdout
trainer.extend(extensions.ProgressBar())
# Run the training
trainer.run()
print('Measure inference speeds for 1 sample inference...')
test_iter = chainer.iterators.SerialIterator(
test, 1, repeat=False, shuffle=False)
if not args.pretrain:
if args.gpu >= 0:
classifier = L.Classifier(model.copy())
start = time.time()
accuracy = extensions.Evaluator(
test_iter, classifier, device=args.gpu)()['main/accuracy']
print('dense Gpu:', time.time() - start,
's/{} imgs'.format(len(test)))
model.to_cpu()
classifier = L.Classifier(model.copy())
start = time.time()
accuracy = extensions.Evaluator(
test_iter, classifier, device=-1)()['main/accuracy']
print('dense Cpu:', time.time() - start, 's/{} imgs'.format(len(test)))
model.to_cpu_sparse()
model.name = None
classifier = L.Classifier(copy.deepcopy(model))
start = time.time()
accuracy = extensions.Evaluator(
test_iter, classifier, device=-1)()['main/accuracy']
print('sparse Cpu:', time.time() - start,
's/{} imgs'.format(len(test)))
if __name__ == '__main__':
main()