update ckpt_manager.py

Author: inzapp

ckpt_manager.py

@@ -34,6 +34,7 @@ class CheckpointManager:
     def __init__(self):
         self.model_name = None
         self.checkpoint_path = None
+        self.best_metric = None
 
     def set_model_name(self, model_name):
         self.model_name = model_name
@@ -69,16 +70,33 @@ def init_checkpoint_dir(self):
         print(f'checkpoint path : {self.checkpoint_path}')
 
     def remove_last_model(self):
-        for last_model_path in glob(f'{self.checkpoint_path}/model_*_iter.h5'):
+        for last_model_path in glob(f'{self.checkpoint_path}/last_*.h5'):
             os.remove(last_model_path)
 
-    def save_last_model(self, model, iteration_count):
+    def save_last_model(self, model, iteration_count, content=''):
         self.make_checkpoint_dir()
-        save_path = f'{self.checkpoint_path}/model_{iteration_count}_iter.h5'
+        save_path = f'{self.checkpoint_path}/last_{iteration_count}_iter{content}.h5'
         model.save(save_path, include_optimizer=False)
         backup_path = f'{save_path}.bak'
         sh.move(save_path, backup_path)
         self.remove_last_model()
         sh.move(backup_path, save_path)
         return save_path
 
+    def remove_best_model(self):
+        for best_model_path in glob(f'{self.checkpoint_path}/best_*.h5'):
+            os.remove(best_model_path)
+
+    def save_best_model(self, model, iteration_count, metric, content=''):
+        save_path = None
+        if self.best_metric is None or metric > self.best_metric:
+            self.best_metric = metric
+            self.make_checkpoint_dir()
+            save_path = f'{self.checkpoint_path}/best_{iteration_count}_iter{content}.h5'
+            model.save(save_path, include_optimizer=False)
+            backup_path = f'{save_path}.bak'
+            sh.move(save_path, backup_path)
+            self.remove_best_model()
+            sh.move(backup_path, save_path)
+        return save_path
+