From 8d02fa043cc3cd69e89704ebc6e53aebc8e472bf Mon Sep 17 00:00:00 2001
From: Maciej Sobkowiak <masobkowiak@gmail.com>
Date: Wed, 16 Feb 2022 22:47:40 +0100
Subject: [PATCH] save model after fitting

---
 main.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/main.py b/main.py
index d9f348a..d5f61c4 100644
--- a/main.py
+++ b/main.py
@@ -4,6 +4,7 @@ from src.loss import jaccard_loss
 from src.metrics import IOU
 from src.consts import EPOCHS, STEPS, SEED
 from src.generators import create_generators
+from src.helpers import create_folder
 from tensorflow.keras.callbacks import ModelCheckpoint
 import tensorflow as tf
 
@@ -16,10 +17,8 @@ if __name__ == "__main__":
         'optimizer':'rmsprop',
         'metrics':[IOU]
         }
-        
                     
     model.compile(**compile_params)
-    # tf.keras.utils.plot_model(model, show_shapes=True)
 
     model_name = "models/unet.h5"
     modelcheckpoint = ModelCheckpoint(model_name,
@@ -38,4 +37,6 @@ if __name__ == "__main__":
                         steps_per_epoch=STEPS,
                         validation_steps = STEPS,
                         shuffle=True,
-    )
\ No newline at end of file
+    )
+    create_folder('models', '.')
+    model.save(filepath=model_name)
\ No newline at end of file