save model at every training batch

JinZhuXing · Mar 5, 2021 · e6e6b5d · e6e6b5d
1 parent 0469dfb
commit e6e6b5d
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/source/04_generalized_model/train.py b/source/04_generalized_model/train.py
@@ -293,6 +293,10 @@ def main(args):
         # training model
         model.fit(train_gen, epochs = train_epoch, validation_data = val_gen, callbacks = [cp_callback])
 
+        # save model
+        model_path = save_model_path + 'fp' + str(image_width) + '_' + str(image_height) + '.h5'
+        model.save(model_path)
+
 
 # argument parser
 def parse_arguments(argv):