Repository: Lightning-AI/pytorch-lightning

Issue:
change Checkpoint callback's `save_best_only` to `save_top_k`
**Is your feature request related to a problem? Please describe.**
`save_best_only` is a special case of `save_top_k`. However, `save_tok_k` checkpoints can be used to create ensemble model during the test time.

**Describe the solution you'd like**
keep a dict of `{epoch: monitor}` of length `k`, and save new ckeckpoint that can enter this dict, remove the worst checkpoint.


change Checkpoint callback's `save_best_only` to `save_top_k`
**Is your feature request related to a problem? Please describe.**
`save_best_only` is a special case of `save_top_k`. However, `save_tok_k` checkpoints can be used to create ensemble model during the test time.

**Describe the solution you'd like**
keep a dict of `{epoch: monitor}` of length `k`, and save new ckeckpoint that can enter this dict, remove the worst checkpoint.




Discussion:
@vincentzlt good suggestion. I like the change.

Any takers? @vincentzlt interested in making this change?
@vincentzlt good suggestion. I like the change.

Any takers? @vincentzlt interested in making this change?

Resolving patch:
diff --git a/pytorch_lightning/callbacks/pt_callbacks.py b/pytorch_lightning/callbacks/pt_callbacks.py
--- a/pytorch_lightning/callbacks/pt_callbacks.py
+++ b/pytorch_lightning/callbacks/pt_callbacks.py
@@ -158,11 +158,17 @@ class ModelCheckpoint(Callback):
         filepath: string, path to save the model file.
         monitor: quantity to monitor.
         verbose: verbosity mode, 0 or 1.
-        save_best_only: if `save_best_only=True`,
-            the latest best model according to
-            the quantity monitored will not be overwritten.
+        save_top_k: if `save_top_k == k`,
+            the best k models according to
+            the quantity monitored will be saved.
+            if `save_top_k == 0`, no models are saved.
+            if `save_top_k == -1`, all models are saved.
+            Please note that the monitors are checked every `period` epochs.
+            if `save_top_k >= 2` and the callback is called multiple
+            times inside an epoch, the name of the saved file will be
+            appended with a version count starting with `v0`.
         mode: one of {auto, min, max}.
-            If `save_best_only=True`, the decision
+            If `save_top_k != 0`, the decision
             to overwrite the current save file is made
             based on either the maximization or the
             minimization of the monitored quantity. For `val_acc`,
@@ -176,27 +182,33 @@ class ModelCheckpoint(Callback):
     """
 
     def __init__(self, filepath, monitor='val_loss', verbose=0,
-                 save_best_only=True, save_weights_only=False,
+                 save_top_k=1, save_weights_only=False,
                  mode='auto', period=1, prefix=''):
         super(ModelCheckpoint, self).__init__()
         if (
-            save_best_only and
+            save_top_k and
             os.path.isdir(filepath) and
             len(os.listdir(filepath)) > 0
         ):
             warnings.warn(
-                f"Checkpoint directory {filepath} exists and is not empty with save_best_only=True."
+                f"Checkpoint directory {filepath} exists and is not empty with save_top_k != 0."
                 "All files in this directory will be deleted when a checkpoint is saved!"
             )
 
         self.monitor = monitor
         self.verbose = verbose
         self.filepath = filepath
-        self.save_best_only = save_best_only
+        if not os.path.exists(filepath):
+            os.makedirs(filepath)
+        self.save_top_k = save_top_k
         self.save_weights_only = save_weights_only
         self.period = period
-        self.epochs_since_last_save = 0
+        self.epochs_since_last_check = 0
         self.prefix = prefix
+        self.best_k_models = {}
+        # {filename: monitor}
+        self.kth_best_model = ''
+        self.best = 0
 
         if mode not in ['auto', 'min', 'max']:
             warnings.warn(
@@ -206,66 +218,112 @@ def __init__(self, filepath, monitor='val_loss', verbose=0,
 
         if mode == 'min':
             self.monitor_op = np.less
-            self.best = np.Inf
+            self.kth_value = np.Inf
+            self.mode = 'min'
         elif mode == 'max':
             self.monitor_op = np.greater
-            self.best = -np.Inf
+            self.kth_value = -np.Inf
+            self.mode = 'max'
         else:
             if 'acc' in self.monitor or self.monitor.startswith('fmeasure'):
                 self.monitor_op = np.greater
-                self.best = -np.Inf
+                self.kth_value = -np.Inf
+                self.mode = 'max'
             else:
                 self.monitor_op = np.less
-                self.best = np.Inf
+                self.kth_value = np.Inf
+                self.mode = 'min'
 
-    def save_model(self, filepath, overwrite):
-        dirpath = '/'.join(filepath.split('/')[:-1])
+    def _del_model(self, filepath):
+        dirpath = os.path.dirname(filepath)
 
         # make paths
-        os.makedirs(os.path.dirname(filepath), exist_ok=True)
+        os.makedirs(dirpath, exist_ok=True)
 
-        if overwrite:
-            for filename in os.listdir(dirpath):
-                if self.prefix in filename:
-                    path_to_delete = os.path.join(dirpath, filename)
-                    try:
-                        shutil.rmtree(path_to_delete)
-                    except OSError:
-                        os.remove(path_to_delete)
+        try:
+            shutil.rmtree(filepath)
+        except OSError:
+            os.remove(filepath)
+
+    def _save_model(self, filepath):
+        dirpath = os.path.dirname(filepath)
+
+        # make paths
+        os.makedirs(dirpath, exist_ok=True)
 
         # delegate the saving to the model
         self.save_function(filepath)
 
+    def check_monitor_top_k(self, current):
+        less_than_k_models = len(self.best_k_models.keys()) < self.save_top_k
+        if less_than_k_models:
+            return True
+        return self.monitor_op(current, self.best_k_models[self.kth_best_model])
+
     def on_epoch_end(self, epoch, logs=None):
         logs = logs or {}
-        self.epochs_since_last_save += 1
-        if self.epochs_since_last_save >= self.period:
-            self.epochs_since_last_save = 0
-            filepath = '{}/{}_ckpt_epoch_{}.ckpt'.format(self.filepath, self.prefix, epoch + 1)
-            if self.save_best_only:
+        self.epochs_since_last_check += 1
+
+        if self.save_top_k == 0:
+            # no models are saved
+            return
+        if self.epochs_since_last_check >= self.period:
+            self.epochs_since_last_check = 0
+            filepath = f'{self.filepath}/{self.prefix}_ckpt_epoch_{epoch}.ckpt'
+            version_cnt = 0
+            while os.path.isfile(filepath):
+                # this epoch called before
+                filepath = f'{self.filepath}/{self.prefix}_ckpt_epoch_{epoch}_v{version_cnt}.ckpt'
+                version_cnt += 1
+
+            print(filepath)
+
+            if self.save_top_k != -1:
                 current = logs.get(self.monitor)
+
                 if current is None:
                     warnings.warn(
                         f'Can save best model only with {self.monitor} available,'
                         ' skipping.', RuntimeWarning)
                 else:
-                    if self.monitor_op(current, self.best):
+                    if self.check_monitor_top_k(current):
+
+                        # remove kth
+                        if len(self.best_k_models.keys()) == self.save_top_k:
+                            delpath = self.kth_best_model
+                            self.best_k_models.pop(self.kth_best_model)
+                            self._del_model(delpath)
+
+                        self.best_k_models[filepath] = current
+                        if len(self.best_k_models.keys()) == self.save_top_k:
+                            # monitor dict has reached k elements
+                            if self.mode == 'min':
+                                self.kth_best_model = max(self.best_k_models, key=self.best_k_models.get)
+                            else:
+                                self.kth_best_model = min(self.best_k_models, key=self.best_k_models.get)
+                            self.kth_value = self.best_k_models[self.kth_best_model]
+
+                        if self.mode == 'min':
+                            self.best = min(self.best_k_models.values())
+                        else:
+                            self.best = max(self.best_k_models.values())
                         if self.verbose > 0:
                             logging.info(
-                                f'\nEpoch {epoch + 1:05d}: {self.monitor} improved'
-                                f' from {self.best:0.5f} to {current:0.5f},'
-                                f' saving model to {filepath}')
-                        self.best = current
-                        self.save_model(filepath, overwrite=True)
+                                f'\nEpoch {epoch:05d}: {self.monitor} reached',
+                                f'{current:0.5f} (best {self.best:0.5f}), saving model to',
+                                f'{filepath} as top {self.save_top_k}')
+                        self._save_model(filepath)
 
                     else:
                         if self.verbose > 0:
                             logging.info(
-                                f'\nEpoch {epoch + 1:05d}: {self.monitor} did not improve')
+                                f'\nEpoch {epoch:05d}: {self.monitor}',
+                                f'was not in top {self.save_top_k}')
+
             else:
                 if self.verbose > 0:
-                    logging.info(f'\nEpoch {epoch + 1:05d}: saving model to {filepath}')
-                self.save_model(filepath, overwrite=False)
+                    logging.info(f'\nEpoch {epoch:05d}: saving model to {filepath}')
+                self._save_model(filepath)
 
 
 class GradientAccumulationScheduler(Callback):


Test patch:
diff --git a/tests/test_a_restore_models.py b/tests/test_a_restore_models.py
--- a/tests/test_a_restore_models.py
+++ b/tests/test_a_restore_models.py
@@ -131,7 +131,7 @@ def test_load_model_from_checkpoint():
     # correct result and ok accuracy
     assert result == 1, 'training failed to complete'
     pretrained_model = LightningTestModel.load_from_checkpoint(
-        os.path.join(trainer.checkpoint_callback.filepath, "_ckpt_epoch_1.ckpt")
+        os.path.join(trainer.checkpoint_callback.filepath, "_ckpt_epoch_0.ckpt")
     )
 
     # test that hparams loaded correctly
diff --git a/tests/test_trainer.py b/tests/test_trainer.py
--- a/tests/test_trainer.py
+++ b/tests/test_trainer.py
@@ -232,6 +232,118 @@ def test_dp_output_reduce():
     assert reduced['b']['c'] == out['b']['c']
 
 
+def test_model_checkpoint_options():
+    """
+    Test ModelCheckpoint options
+    :return:
+    """
+    def mock_save_function(filepath):
+        open(filepath, 'a').close()
+
+    hparams = testing_utils.get_hparams()
+    model = LightningTestModel(hparams)
+
+    # simulated losses
+    save_dir = testing_utils.init_save_dir()
+    losses = [10, 9, 2.8, 5, 2.5]
+
+    # -----------------
+    # CASE K=-1  (all)
+    w = ModelCheckpoint(save_dir, save_top_k=-1, verbose=1)
+    w.save_function = mock_save_function
+    for i, loss in enumerate(losses):
+        w.on_epoch_end(i, logs={'val_loss': loss})
+
+    file_lists = set(os.listdir(save_dir))
+
+    assert len(file_lists) == len(losses), "Should save all models when save_top_k=-1"
+
+    # verify correct naming
+    for i in range(0, len(losses)):
+        assert f'_ckpt_epoch_{i}.ckpt' in file_lists
+
+    testing_utils.clear_save_dir()
+
+    # -----------------
+    # CASE K=0 (none)
+    w = ModelCheckpoint(save_dir, save_top_k=0, verbose=1)
+    w.save_function = mock_save_function
+    for i, loss in enumerate(losses):
+        w.on_epoch_end(i, logs={'val_loss': loss})
+
+    file_lists = os.listdir(save_dir)
+
+    assert len(file_lists) == 0, "Should save 0 models when save_top_k=0"
+
+    testing_utils.clear_save_dir()
+
+    # -----------------
+    # CASE K=1 (2.5, epoch 4)
+    w = ModelCheckpoint(save_dir, save_top_k=1, verbose=1, prefix='test_prefix')
+    w.save_function = mock_save_function
+    for i, loss in enumerate(losses):
+        w.on_epoch_end(i, logs={'val_loss': loss})
+
+    file_lists = set(os.listdir(save_dir))
+
+    assert len(file_lists) == 1, "Should save 1 model when save_top_k=1"
+    assert 'test_prefix_ckpt_epoch_4.ckpt' in file_lists
+
+    testing_utils.clear_save_dir()
+
+    # -----------------
+    # CASE K=2 (2.5 epoch 4, 2.8 epoch 2)
+    # make sure other files don't get deleted
+
+    w = ModelCheckpoint(save_dir, save_top_k=2, verbose=1)
+    open(f'{save_dir}/other_file.ckpt', 'a').close()
+    w.save_function = mock_save_function
+    for i, loss in enumerate(losses):
+        w.on_epoch_end(i, logs={'val_loss': loss})
+
+    file_lists = set(os.listdir(save_dir))
+
+    assert len(file_lists) == 3, 'Should save 2 model when save_top_k=2'
+    assert '_ckpt_epoch_4.ckpt' in file_lists
+    assert '_ckpt_epoch_2.ckpt' in file_lists
+    assert 'other_file.ckpt' in file_lists
+
+    testing_utils.clear_save_dir()
+
+    # -----------------
+    # CASE K=4 (save all 4 models)
+    # multiple checkpoints within same epoch
+
+    w = ModelCheckpoint(save_dir, save_top_k=4, verbose=1)
+    w.save_function = mock_save_function
+    for loss in losses:
+        w.on_epoch_end(0, logs={'val_loss': loss})
+
+    file_lists = set(os.listdir(save_dir))
+
+    assert len(file_lists) == 4, 'Should save all 4 models when save_top_k=4 within same epoch'
+
+    testing_utils.clear_save_dir()
+
+    # -----------------
+    # CASE K=3 (save the 2nd, 3rd, 4th model)
+    # multiple checkpoints within same epoch
+
+    w = ModelCheckpoint(save_dir, save_top_k=3, verbose=1)
+    w.save_function = mock_save_function
+    for loss in losses:
+        w.on_epoch_end(0, logs={'val_loss': loss})
+
+    file_lists = set(os.listdir(save_dir))
+
+    assert len(file_lists) == 3, 'Should save 3 models when save_top_k=3'
+    assert '_ckpt_epoch_0_v2.ckpt' in file_lists
+    assert '_ckpt_epoch_0_v1.ckpt' in file_lists
+    assert '_ckpt_epoch_0.ckpt' in file_lists
+
+    testing_utils.clear_save_dir()
+
+
 def test_model_freeze_unfreeze():
     testing_utils.reset_seed()
 
diff --git a/tests/testing_utils.py b/tests/testing_utils.py
--- a/tests/testing_utils.py
+++ b/tests/testing_utils.py
@@ -149,7 +149,7 @@ def init_save_dir():
 
 def clear_save_dir():
     root_dir = os.path.dirname(os.path.realpath(__file__))
-    save_dir = os.path.join(root_dir, 'save_dir')
+    save_dir = os.path.join(root_dir, 'tests', 'save_dir')
     if os.path.exists(save_dir):
         n = RANDOM_FILE_PATHS.pop()
         shutil.move(save_dir, save_dir + f'_{n}')