-
Notifications
You must be signed in to change notification settings - Fork 12
/
train.py
306 lines (241 loc) · 10.2 KB
/
train.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
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
import json
import numpy as np
import os
import glob
import argparse
import matplotlib.pyplot as plt
from nets.inception_v3 import InceptionV3
from nets.inception_resnet_v2 import InceptionResNetV2
from nets.nasnet import *
from nets.densenet import *
from nets.resnet152 import *
from nets.resnet101 import *
from nets.resnet50 import *
from nets.inception_v4 import InceptionV4
from nets.vgg19 import *
#from nets.vgg16 import *
# from nets.NASnet import *
from nets.xception import Xception
from keras.models import Model
from keras.layers import *
from keras import backend as K
from keras.callbacks import ModelCheckpoint
from keras.callbacks import TensorBoard
from keras.models import load_model
from keras.optimizers import *
from nets.densenet import *
from nets.inception_v4 import InceptionV4, preprocess_input
from keras.applications import *
from keras.preprocessing.image import *
train_data_dir = '/home/fenglf/data/dog/stanford/Images/data'
validation_data_dir = '/home/fenglf/data/dog/stanford/Images/test_data'
pretrained_model_root_dir = '/home/fenglf/PycharmProjects/keras-finetuning-master/model/pretrained/'
output_model_root_dir = '/home/fenglf/PycharmProjects/keras-finetuning-master/model/outputk'
#base_model_name = VGG19
base_model_name = InceptionV3
np.random.seed(2018)
top_epochs = 0
fit_epochs = 5
batch_size = 32
nb_fc_hidden_layer = 1024
nb_classes = 120
#NB_LAYERS_TO_FREEZE = 652
NB_LAYERS_TO_FREEZE = 406
# input size
#img_width, img_height = 224, 224
img_width, img_height = 299, 299
# img_width, img_height = 331, 331
learning_rate_finetune = 0.0001
momentum_finetune = 0.9
input_included_model = ['VGG16',
'VGG19',
'ResNet50',
'InceptionV3',
'Xception',
'ResNet152',
'ResNet101',
'NASNetLarge',
'InceptionResNetV2']
# plot switch: whether to visualize the training loss and acc
plot_switch = True
# the path of pretrained model
pretrained_model_dir = os.path.join(pretrained_model_root_dir, base_model_name.__name__, base_model_name.__name__ + '_notop.h5')
# the path of fine-tuned checkpoint path
# top_layers_checkpoint_path refer to the best new added fc layer fine-tuned checkpoint path
top_layers_checkpoint_path = os.path.join(output_model_root_dir, base_model_name.__name__, 'top_layer_weights', base_model_name.__name__ + '.top.best.hdf5')
# fine_tuned_checkpoint_path refer to the best free layers of base_model fine-tuned checkpoint path
fine_tuned_checkpoint_path = os.path.join(output_model_root_dir, base_model_name.__name__, 'fine_tuned_weights', base_model_name.__name__ + '.fine_tuned.best.hdf5')
# final_weights_path refer to the best fine-tuned final weights
final_weights_path = os.path.join(output_model_root_dir, base_model_name.__name__, 'final_weights', base_model_name.__name__ + '.final_weights.hdf5')
final_weights_json_path = os.path.join(output_model_root_dir, base_model_name.__name__, 'final_weights', base_model_name.__name__ + '.final_weights.json')
final_weights_label_path = os.path.join(output_model_root_dir, base_model_name.__name__, 'final_weights', base_model_name.__name__ + '.final_weights-labels.json')
def add_new_last_layer(base_model, nb_classes):
"""Add last layer to the convnet
Args:
base_model: keras model excluding top
nb_classes: # of classes
Returns:
new keras model with last layer
"""
# add a global spatial average pooling layer
x = base_model.output
if base_model_name.__name__ in input_included_model:
#if x.shape[1] < 4:
ax = GlobalAveragePooling2D()(x)
#else:
#ax = AveragePooling2D(pool_size=(2, 2))(x)
#ax = Flatten(name='flatten')(ax)
else:
ax = x
#x = Dropout(0.5)(x)
# let's add a fully-connected layer, random init
x = Dense(nb_fc_hidden_layer)(ax)
#print ax.shape
# add BN layer and Dropout flf
x = BatchNormalization()(x)
x = Activation('relu')(x)
#x = Dropout(0.5)(x)
# and a logistic layer -- we have 120 classes
predictions = Dense(nb_classes, activation='softmax')(x)
# this is the model we will train
model = Model(input=base_model.input, output=predictions)
return model
def setup_to_transfer_learn(model, base_model):
"""Freeze all layers and compile the model"""
# i.e. freeze all convolutional layers
for layer in base_model.layers:
layer.trainable = False
# compile the model (should be done *after* setting layers to non-trainable)
model.compile(optimizer='adadelta', loss='categorical_crossentropy', metrics=['accuracy'], )
def setup_to_finetune(model):
"""Freeze the bottom NB_LAYERS_TO_FREEZE and retrain the remaining top layers.
note: NB_LAYERS_TO_FREEZE corresponds to the top 2 inception blocks in the base_model arch
Args:
model: keras model
"""
for layer in model.layers[:NB_LAYERS_TO_FREEZE]:
layer.trainable = False
for layer in model.layers[NB_LAYERS_TO_FREEZE:]:
layer.trainable = True
# we need to recompile the model for these modifications to take effect
# we use SGD with a low learning rate
model.compile(optimizer=SGD(lr=learning_rate_finetune, momentum=momentum_finetune), loss='categorical_crossentropy', metrics=['accuracy'])
def plot_training(history):
acc = history.history['acc']
val_acc = history.history['val_acc']
loss = history.history['loss']
val_loss = history.history['val_loss']
epochs = range(len(acc))
plt.plot(epochs, acc, 'r.')
plt.plot(epochs, val_acc, 'r')
plt.title('Training and validation accuracy')
plt.figure()
plt.plot(epochs, loss, 'r.')
plt.plot(epochs, val_loss, 'r-')
plt.title('Training and validation loss')
plt.show()
def train(lambda_func):
# prepare data augmentation configuration
train_datagen = ImageDataGenerator(
rescale=1. / 255)
test_datagen = ImageDataGenerator(
rescale=1. / 255)
train_datagen = ImageDataGenerator(
preprocessing_function=lambda_func)
test_datagen = ImageDataGenerator(
preprocessing_function=lambda_func)
train_generator = train_datagen.flow_from_directory(
train_data_dir,
target_size=(img_height, img_width),
batch_size=batch_size,
class_mode='categorical')
class_dict = train_generator.class_indices
print class_dict
validation_generator = test_datagen.flow_from_directory(
validation_data_dir,
target_size=(img_height, img_width),
batch_size=batch_size,
class_mode='categorical')
# define the input_shape
if K.image_data_format() == 'channels_first':
input_shape = (3, img_width, img_height)
else:
input_shape = (img_width, img_height, 3)
# setup model
if base_model_name.__name__ == 'DenseNet201':
# densenet_weights_path = densenet_weights_root + 'densenet169_weights_tf.h5'
base_model = base_model_name(reduction=0.5, input_shape=(img_width, img_height, 3), weights=pretrained_model_dir)
else:
#base_model = base_model_name(weights=pretrained_model_dir, input_shape=input_shape, include_top=False)
base_model = base_model_name(weights=pretrained_model_dir, include_top=False)
model = add_new_last_layer(base_model, nb_classes)
# base_model.layers.pop()
# for i, layer in enumerate(base_model.layers):
# print (i, layer.name)
if os.path.exists(top_layers_checkpoint_path):
model.load_weights(top_layers_checkpoint_path)
print ("Checkpoint " + top_layers_checkpoint_path + " loaded.")
# first: train only the top layers (which were randomly initialized)
setup_to_transfer_learn(model, base_model)
# Save the model after every epoch.
mc_top = ModelCheckpoint(
top_layers_checkpoint_path,
monitor='val_acc',
verbose=0,
save_best_only=True,
save_weights_only=False,
mode='auto',
period=1)
# Save the TensorBoard logs.
tb = TensorBoard(log_dir='./logs', histogram_freq=0, write_graph=True, write_images=True)
# train the model on the new data for a few epochs
model.fit_generator(
train_generator,
steps_per_epoch=len(train_generator.filenames) // batch_size,
epochs=top_epochs,
validation_data=validation_generator,
validation_steps=len(validation_generator.filenames) // batch_size,
callbacks=[mc_top, tb])
# at this point, the top layers are well trained and we can start fine-tuning
# convolutional layers from base model. We will freeze the bottom N layers
# and train the remaining top layers.
# let's visualize layer names and layer indices to see how many layers
# we should freeze:
for i, layer in enumerate(base_model.layers):
print (i, layer.name)
if os.path.exists(fine_tuned_checkpoint_path):
model.load_weights(fine_tuned_checkpoint_path)
print ("Checkpoint '" + fine_tuned_checkpoint_path + "' loaded.")
# we chose to train the some top blocks, i.e. we will freeze
# the first NB_LAYERS_TO_FREEZE layers and unfreeze the rest:
setup_to_finetune(model)
# Save the model after every epoch.
mc_fit = ModelCheckpoint(
fine_tuned_checkpoint_path,
monitor='val_acc',
verbose=0,
save_best_only=True,
save_weights_only=False,
mode='auto',
period=1)
# we train our model again (this time fine-tuning the top 2 inception blocks
# alongside the top Dense layers
history_ft = model.fit_generator(
train_generator,
steps_per_epoch=len(train_generator.filenames) // batch_size,
epochs=fit_epochs,
validation_data=validation_generator,
validation_steps=len(validation_generator.filenames) // batch_size,
callbacks=[mc_fit, tb])
# save final weights
model.save_weights(final_weights_path)
# serialize model to JSON
model_json = model.to_json()
with open(final_weights_json_path, "w") as json_file:
json_file.write(model_json)
with open(final_weights_label_path, "w") as json_file:
json.dump(class_dict, json_file)
if plot_switch:
plot_training(history_ft)
if __name__ == '__main__':
train(inception_v3.preprocess_input)