-
Notifications
You must be signed in to change notification settings - Fork 0
/
train_models.py
executable file
·121 lines (94 loc) · 4.12 KB
/
train_models.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
## train_models.py -- train the neural network models for attacking
##
## Copyright (C) 2016, Nicholas Carlini <nicholas@carlini.com>.
##
## This program is licenced under the BSD 2-Clause licence,
## contained in the LICENCE file in this directory.
# modified by Juan C. Perdomo 2017
from keras.models import Sequential
from keras.layers import Dense, Dropout, Activation, Flatten
from keras.layers import Conv2D, MaxPooling2D
from keras.optimizers import SGD
import tensorflow as tf
from setup_mnist import MNIST
import os
def train(data, file_name, params, num_epochs=50, batch_size=128, train_temp=1, init=None, conv=True):
"""
Standard neural network training procedure.
"""
model = Sequential()
print(data.train_data.shape)
if conv:
model.add(Conv2D(params[0], (3, 3),
input_shape=data.train_data.shape[1:]))
model.add(Activation('relu'))
model.add(Conv2D(params[1], (3, 3)))
model.add(Activation('relu'))
model.add(MaxPooling2D(pool_size=(2, 2)))
model.add(Conv2D(params[2], (3, 3)))
model.add(Activation('relu'))
model.add(Conv2D(params[3], (3, 3)))
model.add(Activation('relu'))
model.add(MaxPooling2D(pool_size=(2, 2)))
model.add(Flatten())
model.add(Dense(params[4]))
model.add(Activation('relu'))
model.add(Dropout(0.5))
model.add(Dense(params[5]))
model.add(Activation('relu'))
model.add(Dense(10))
else:
model.add(Flatten(input_shape=data.train_data.shape[1:]))
for p in params:
model.add(Dense(p))
model.add(Activation('relu'))
model.add(Dropout(0.1))
model.add(Dense(10))
if init != None:
model.load_weights(init)
def fn(correct, predicted):
return tf.nn.softmax_cross_entropy_with_logits(labels=correct,
logits=predicted/train_temp)
sgd = SGD(lr=0.01, decay=1e-6, momentum=0.9, nesterov=True)
model.compile(loss=fn,
optimizer=sgd,
metrics=['accuracy'])
model.fit(data.train_data, data.train_labels,
batch_size=batch_size,
validation_data=(data.validation_data, data.validation_labels),
nb_epoch=num_epochs,
shuffle=True)
if file_name != None:
model.save(file_name)
return model
def train_distillation(data, file_name, params, num_epochs=50, batch_size=128, train_temp=1, conv=True):
"""
Train a network using defensive distillation.
Distillation as a Defense to Adversarial Perturbations against Deep Neural Networks
Nicolas Papernot, Patrick McDaniel, Xi Wu, Somesh Jha, Ananthram Swami
IEEE S&P, 2016.
"""
if not os.path.exists(file_name+"_init"):
# Train for one epoch to get a good starting point.
train(data, file_name+"_init", params, 1, batch_size, conv=conv)
# now train the teacher at the given temperature
teacher = train(data, file_name+"_teacher", params, num_epochs, batch_size, train_temp,
init=file_name+"_init", conv=conv)
# evaluate the labels at temperature t
predicted = teacher.predict(data.train_data)
with tf.Session() as sess:
y = sess.run(tf.nn.softmax(predicted/train_temp))
print(y)
data.train_labels = y
# train the student model at temperature t
student = train(data, file_name, params, num_epochs, batch_size, train_temp,
init=file_name+"_init", conv=conv)
# and finally we predict at temperature 1
predicted = student.predict(data.train_data)
print(predicted)
if not os.path.isdir('models'):
os.makedirs('models')
train(MNIST(), "models/mnist", [32, 32, 64, 64, 200, 200], num_epochs=30)
train(MNIST(), "models/multilayer", [128, 128, 128, 128], conv=False, num_epochs=50)
train_distillation(MNIST(), "models/multilayer-distilled-100", [128, 128, 128, 128], conv=False, num_epochs=50, train_temp=100)
train_distillation(MNIST(), "models/mnist-distilled-100", [32, 32, 64, 64, 200, 200], num_epochs=30, train_temp=100)