def train_stage1(self):
"""
In this stage, we will freeze all the convolution blocks and train
only the newly added dense layers. We will add a global spatial average
pooling layer, we will add fully connected dense layers on the output
of the base models. We will freeze the convolution base and train only
the top layers. We will set all the convolution layers to false, the model
should be compiled when all the convolution layers are set to false.
Arguments:
-input_params : This parameter will contain all the information that the user will
input through the terminal
"""
print(
"\nTraining the model by freezing the convolution block and tuning the top layers..."
)
st = dt.now()
utils_obj = Utility(self.input_params, self.path_dict)
#Put if statement here. If model_name != custom then run this block, or else. Do something else.
if (self.input_params['model_name'] != 'custom'):
base_model = utils_obj.load_imagenet_model()
#Adding a global spatial average pooling layer
x = base_model.output
x = GlobalAveragePooling2D()(x)
#Adding a fully-connected dense layer
#x = Dense(self.input_params['dense_neurons'], activation='relu', kernel_initializer='he_normal')(x)
#Adding the custom layers
customlayers = self.input_params['customlayers']
#Adding a final dense output final layer
x = customlayers(x)
n = utils_obj.no_of_classes()
output_layer = Dense(
n,
activation=self.input_params['outputlayer_activation'],
kernel_initializer='glorot_uniform')(x)
model_stg1 = Model(inputs=base_model.input, outputs=output_layer)
#Define the model
model_stg1 = Model(inputs=base_model.input, outputs=output_layer)
#Here we will freeze the convolution base and train only the top layers
#We will set all the convolution layers to false, the model should be
#compiled when all the convolution layers are set to false
for layer in base_model.layers:
layer.trainable = False
else:
model_stg1 = self.input_params['custom_model']
#Compiling the model
model_stg1.compile(
optimizer=optimizers.Adam(lr=self.input_params['stage1_lr']),
loss='categorical_crossentropy',
metrics=[self.input_params['metric']])
#Normalize the images
train_datagen = ImageDataGenerator(
preprocessing_function=utils_obj.init_preprocess_func())
val_datagen = ImageDataGenerator(
preprocessing_function=utils_obj.init_preprocess_func())
df_train = utils_obj.load_data("train")
df_val = utils_obj.load_data("val")
train_generator = train_datagen.flow_from_dataframe(
dataframe=df_train,
directory=self.path_dict['source'],
target_size=utils_obj.init_sizes(),
x_col="filenames",
y_col="class_label",
batch_size=self.input_params['batch_size'],
class_mode='categorical',
color_mode='rgb',
shuffle=True)
val_generator = val_datagen.flow_from_dataframe(
dataframe=df_val,
directory=self.path_dict['source'],
target_size=utils_obj.init_sizes(),
x_col="filenames",
y_col="class_label",
batch_size=self.input_params['batch_size'],
class_mode='categorical',
color_mode='rgb',
shuffle=True)
nb_train_samples = len(train_generator.classes)
nb_val_samples = len(val_generator.classes)
history = model_stg1.fit_generator(
generator=train_generator,
steps_per_epoch=nb_train_samples //
self.input_params['batch_size'],
epochs=self.input_params['epochs1'],
validation_data=val_generator,
validation_steps=nb_val_samples // self.input_params['batch_size'],
callbacks=TrainingUtils.callbacks_list(self, 1),
workers=self.input_params['nworkers'],
use_multiprocessing=False,
max_queue_size=20) #1 for stage 1
hist_df = pd.DataFrame(history.history)
hist_csv_file = self.path_dict['model_path'] + "stage{}/".format(
1) + "{}_history_stage_{}.csv".format(
self.input_params['model_name'], 1)
with open(hist_csv_file, mode='w') as file:
hist_df.to_csv(file, index=None)
#model_stg1.load_weights(self.path_dict['model_path'] + "stage{}/".format(1) + "{}_weights_stage_{}.hdf5".format(self.input_params['model_name'], 1))
model_stg1.save(
self.path_dict['model_path'] + "stage{}/".format(1) +
"{}_model_stage_{}.h5".format(self.input_params['model_name'], 1))
TrainingUtils.save_summary(self, model_stg1, 1)
TrainingUtils.plot_layer_arch(self, model_stg1, 1)
stage1_params = dict()
stage1_params['train_generator'] = train_generator
stage1_params['val_generator'] = val_generator
stage1_params['nb_train_samples'] = nb_train_samples
stage1_params['nb_val_samples'] = nb_val_samples
print("\nTime taken to train the model in stage 1: ", dt.now() - st)
#Start model evaluation for Stage 1
eval_utils = EvalUtils(self.input_params, self.path_dict, 1)
eval_utils.predict_on_test()
return model_stg1, stage1_params
