-
Notifications
You must be signed in to change notification settings - Fork 1
/
CNN_train_UCF101.py
140 lines (117 loc) · 4.57 KB
/
CNN_train_UCF101.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
"""
Train on images split into directories. This assumes we've split
our videos into frames and moved them to their respective folders.
Use keras 2+ and tensorflow 1+
Based on:
https://keras.io/preprocessing/image/
and
https://keras.io/applications/
"""
from tensorflow.keras.applications.inception_v3 import InceptionV3
from tensorflow.keras.applications.efficientnet import EfficientNetB3
from tensorflow.keras.optimizers import SGD
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Dense, GlobalAveragePooling2D, BatchNormalization
from tensorflow.keras.callbacks import ModelCheckpoint, TensorBoard, EarlyStopping
from UCFdata import DataSet
data = DataSet()
# Helper: Save the min val_loss model in each epoch.
checkpointer = ModelCheckpoint(
filepath='./data/checkpoints/inception.{epoch:03d}-{val_loss:.2f}.hdf5',
verbose=1,
save_best_only=True)
# Helper: Stop when we stop learning.
# patience: number of epochs with no improvement after which training will be stopped.
early_stopper = EarlyStopping(patience=10)
# Helper: TensorBoard
tensorboard = TensorBoard(log_dir='./data/logs/')
def get_generators():
train_datagen = ImageDataGenerator(
# rescale=1./255,
shear_range=0.2,
horizontal_flip=True,
rotation_range=10.,
width_shift_range=0.2,
height_shift_range=0.2)
test_datagen = ImageDataGenerator(
# rescale=1./255
)
train_generator = train_datagen.flow_from_directory(
'./data/UCF101/train1/',
target_size=(299, 299),
batch_size=16,
classes=data.classes,
class_mode='categorical')
validation_generator = test_datagen.flow_from_directory(
'./data/UCF101/test1/',
target_size=(299, 299),
batch_size=16,
classes=data.classes,
class_mode='categorical')
return train_generator, validation_generator
def get_model(weights='imagenet'):
# create the base pre-trained model
# base_model = InceptionV3(weights=weights, include_top=False)
base_model = EfficientNetB3(weights='imagenet', include_top=False)
for layer in base_model.layers:
if isinstance(layer, BatchNormalization ):
layer.trainable = True
else:
layer.trainable = False
# add a global spatial average pooling layer
x = base_model.output
x = GlobalAveragePooling2D()(x)
# let's add a fully-connected layer
x = Dense(1024, activation='relu')(x)
# and a logistic layer -- let's say we have 2 classes
predictions = Dense(len(data.classes), activation='softmax')(x)
# this is the model we will train
model = Model(inputs=base_model.input, outputs=predictions)
# for layer in base_model.layers:
# layer.trainable = False
# compile the model (should be done *after* setting layers to non-trainable)
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
return model
def fine_tune_inception_layer(model):
"""After we fine-tune the dense layers, train deeper."""
# we chose to train the top 2 inception blocks, i.e. we will freeze
# the first 172 layers and unfreeze the rest:
for layer in model.layers[:172]:
layer.trainable = False
for layer in model.layers[172:]:
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='adam'(lr=0.0005, momentum=0.9),
loss='categorical_crossentropy',
metrics=['accuracy', 'top_k_categorical_accuracy'])
return model
def train_model(model, nb_epoch, generators, callbacks=[]):
train_generator, validation_generator = generators
model.summary()
model.fit_generator(
train_generator,
steps_per_epoch=100,
validation_data=validation_generator,
validation_steps=10,
epochs=nb_epoch,
callbacks=callbacks)
return model
def main(weights_file):
model = get_model()
generators = get_generators()
if weights_file is None:
print("Training Top layers.")
model = train_model(model, 10, generators)
else:
print("Loading saved model: %s." % weights_file)
model.load_weights(weights_file)
# Get and train the mid layers.
# model = fine_tune_inception_layer(model)
model = train_model(model, 1000, generators,
[checkpointer, early_stopper, tensorboard])
if __name__ == '__main__':
weights_file = None
main(weights_file)