use best model from pretrain

train_classifier.py

@@ -6,6 +6,7 @@ from network.Resnet50 import Resnet50
 from helpers.DataLoader import DataLoader
 from tensorflow.keras.callbacks import ModelCheckpoint, TensorBoard, EarlyStopping
 from tensorflow.keras.optimizers import RMSprop, SGD
+from tensorflow.keras.models import load_model
 
 import pathlib
 
@@ -100,6 +101,12 @@ model.compile(optimizer=RMSprop(), loss='categorical_crossentropy', metrics=['ac
     Pre-Train FC Layers 
 '''
 callbacks_prefit = [
+    ModelCheckpoint(
+        filepath= str(checkpoint_path) + '/best_pre_train.hdf5',
+        save_best_only=True,
+        monitor='val_loss',
+        verbose=1),
+    EarlyStopping(monitor='val_loss', patience=10),
     TensorBoard(options.output_path +'/logs/{}_prefit'.format(time)),
 ]
 
@@ -112,6 +119,11 @@ history = model.fit(train_dataset,
 #model.save_weights('last_pre_train_model.h5')
 print('\nHistory dict:', history.history)
 
+'''
+    Run refinement training on best model of pre-train 
+'''
+model = load_model(str(checkpoint_path) + '/best_pre_train.hdf5')
+
 for layer in model.layers[-39:]:
     layer.trainable = True