Repository: Lightning-AI/pytorch-lightning
Pull request: change Checkpoint callback's `save_best_only` to `save_top_k`
this pr should close #70 
bringing `save_function` outside of `pytorch_lightning/callbacks/pt_callbacks.py` is really confusing, since in that case `pt_callbacks.py` cannot be run on its own :new_moon_with_face: 
not sure where to add the tests, so I appended them at the end of the file.

Problem statement:
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.

Hints from issue discussion:
@vincentzlt good suggestion. I like the change.

Any takers? @vincentzlt interested in making this change?
i'm interested in this feature, i'll give it a try :smile: 
awesome!
Hi, I have some questions.

I noticed `save_function` option was deleted in this commit: https://github.com/williamFalcon/pytorch-lightning/commit/63d84283a4cd91578c413b2ad6aaaec80e3fa3bc

But how does this line work then? https://github.com/williamFalcon/pytorch-lightning/blob/master/pytorch_lightning/callbacks/pt_callbacks.py#L223-L224

Also, there were some `save_function=None` in docs, which make it somehow confusing.
https://github.com/williamFalcon/pytorch-lightning/blob/master/pytorch_lightning/models/trainer.py#L137


https://github.com/williamFalcon/pytorch-lightning/blob/master/pytorch_lightning/root_module/model_saving.py#L57
@vincentzlt https://gitter.im/PyTorch-Lightning/community?utm_source=badge&utm_medium=badge&utm_campaign=pr-badge

Commit message: Update test_models.py

Patch to review:
diff --git a/docs/Trainer/Checkpointing.md b/docs/Trainer/Checkpointing.md
index de99f9589b141..7abb25e90e724 100644
--- a/docs/Trainer/Checkpointing.md
+++ b/docs/Trainer/Checkpointing.md
@@ -9,7 +9,7 @@ from pytorch_lightning.callbacks import ModelCheckpoint
 
 checkpoint_callback = ModelCheckpoint(
     filepath='/path/to/store/weights.ckpt',
-    save_best_only=True,
+    save_top_k=1,
     verbose=True,
     monitor='val_loss',
     mode='min'
diff --git a/docs/examples/Examples.md b/docs/examples/Examples.md
index 46481bce6e1d4..63d86d215f9cb 100644
--- a/docs/examples/Examples.md
+++ b/docs/examples/Examples.md
@@ -77,7 +77,7 @@ def main(hparams, cluster, results_dict):
     checkpoint = ModelCheckpoint(
         filepath=model_save_path,
         save_function=None,
-        save_best_only=True,
+        save_top_k=1,
         verbose=True,
         monitor=hparams.model_save_monitor_value,
         mode=hparams.model_save_monitor_mode
diff --git a/examples/new_project_templates/multi_node_cluster_template.py b/examples/new_project_templates/multi_node_cluster_template.py
index c4af6d416b20a..577c8dde09952 100644
--- a/examples/new_project_templates/multi_node_cluster_template.py
+++ b/examples/new_project_templates/multi_node_cluster_template.py
@@ -66,7 +66,7 @@ def main(hparams, cluster, results_dict):
 
     checkpoint = ModelCheckpoint(
         filepath=model_save_path,
-        save_best_only=True,
+        save_top_k=1,
         verbose=True,
         monitor='val_loss',
         mode='min'
diff --git a/examples/new_project_templates/single_cpu_template.py b/examples/new_project_templates/single_cpu_template.py
index c0f4826f99634..c2dedfa654f11 100644
--- a/examples/new_project_templates/single_cpu_template.py
+++ b/examples/new_project_templates/single_cpu_template.py
@@ -57,7 +57,7 @@ def main(hparams):
 
     checkpoint = ModelCheckpoint(
         filepath=model_save_path,
-        save_best_only=True,
+        save_top_k=1,
         verbose=True,
         monitor='val_loss',
         mode='min'
diff --git a/examples/new_project_templates/single_gpu_node_16bit_template.py b/examples/new_project_templates/single_gpu_node_16bit_template.py
index c491f9542ee39..d3a79cb469b8d 100644
--- a/examples/new_project_templates/single_gpu_node_16bit_template.py
+++ b/examples/new_project_templates/single_gpu_node_16bit_template.py
@@ -57,7 +57,7 @@ def main(hparams):
 
     checkpoint = ModelCheckpoint(
         filepath=model_save_path,
-        save_best_only=True,
+        save_top_k=1,
         verbose=True,
         monitor='val_loss',
         mode='min'
diff --git a/examples/new_project_templates/single_gpu_node_ddp_template.py b/examples/new_project_templates/single_gpu_node_ddp_template.py
index dac7b6d19dd23..e05dfa161769d 100644
--- a/examples/new_project_templates/single_gpu_node_ddp_template.py
+++ b/examples/new_project_templates/single_gpu_node_ddp_template.py
@@ -57,7 +57,7 @@ def main(hparams):
 
     checkpoint = ModelCheckpoint(
         filepath=model_save_path,
-        save_best_only=True,
+        save_top_k=1,
         verbose=True,
         monitor='val_loss',
         mode='min'
diff --git a/examples/new_project_templates/single_gpu_node_dp_template.py b/examples/new_project_templates/single_gpu_node_dp_template.py
index dc7dbd5327fe8..657fb48ee8697 100644
--- a/examples/new_project_templates/single_gpu_node_dp_template.py
+++ b/examples/new_project_templates/single_gpu_node_dp_template.py
@@ -57,7 +57,7 @@ def main(hparams):
 
     checkpoint = ModelCheckpoint(
         filepath=model_save_path,
-        save_best_only=True,
+        save_top_k=1,
         verbose=True,
         monitor='val_loss',
         mode='min'
diff --git a/examples/new_project_templates/trainer_cpu_template.py b/examples/new_project_templates/trainer_cpu_template.py
index 84a29a9bf3627..47a695a9566a8 100644
--- a/examples/new_project_templates/trainer_cpu_template.py
+++ b/examples/new_project_templates/trainer_cpu_template.py
@@ -42,7 +42,7 @@ def main(hparams):
     model_save_path = '{}/{}/{}'.format(hparams.model_save_path, exp.name, exp.version)
     checkpoint = ModelCheckpoint(
         filepath=model_save_path,
-        save_best_only=True,
+        save_top_k=1,
         verbose=True,
         monitor='val_acc',
         mode='min'
diff --git a/pytorch_lightning/callbacks/pt_callbacks.py b/pytorch_lightning/callbacks/pt_callbacks.py
index 24d46e756db62..9232e2674f65e 100644
--- a/pytorch_lightning/callbacks/pt_callbacks.py
+++ b/pytorch_lightning/callbacks/pt_callbacks.py
@@ -156,11 +156,12 @@ 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`, the models are saved every `period` epochs.
         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`,
@@ -174,17 +175,20 @@ class ModelCheckpoint(Callback):
     """
 
     def __init__(self, filepath, monitor='val_loss', verbose=0,
-                 save_best_only=False, save_weights_only=False,
+                 save_top_k=0, save_weights_only=False,
                  mode='auto', period=1, prefix=''):
         super(ModelCheckpoint, self).__init__()
         self.monitor = monitor
         self.verbose = verbose
         self.filepath = filepath
-        self.save_best_only = save_best_only
+        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 = {}
+        # {epoch: monitor}
+        self.best = 0
 
         if mode not in ['auto', 'min', 'max']:
             print('ModelCheckpoint mode %s is unknown, '
@@ -193,72 +197,95 @@ 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, overwrite):
+        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_value])
+
     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.epochs_since_last_check >= self.period:
+            self.epochs_since_last_check = 0
+            filepath = f'{self.filepath}/{self.prefix}_ckpt_epoch_{epoch + 1}.ckpt'
+
+            if self.save_top_k:
                 current = logs.get(self.monitor)
+
                 if current is None:
-                    print('Can save best model only with %s available,'
-                          ' skipping.' % (self.monitor), RuntimeWarning)
+                    print(f'Can save best model only with {self.monitor} available,'
+                          f' skipping.', RuntimeWarning)
                 else:
-                    if self.monitor_op(current, self.best):
+                    if self.check_monitor_top_k(current):
+
+                        # nremove kth
+                        if len(self.best_k_models.keys()) == self.save_top_k:
+                            delpath = f'{self.filepath}/{self.prefix}_ckpt_epoch_{self.kth_value + 1}.ckpt'
+                            self.best_k_models.pop(self.kth_value)
+                            self._del_model(delpath)
+                            
+                        self.best_k_models[epoch] = current
+                        if len(self.best_k_models.keys()) == self.save_top_k:
+                            # monitor dict has reached k elements
+                            if self.mode == 'min':
+                                self.kth_value = max(self.best_k_models, key=self.best_k_models.get)
+                            else:
+                                self.kth_value = min(self.best_k_models, key=self.best_k_models.get)
+
+                        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:
-                            print('\nEpoch %05d: %s improved from %0.5f to %0.5f,'
-                                  ' saving model to %s'
-                                  % (epoch + 1, self.monitor, self.best,
-                                     current, filepath))
-                        self.best = current
-                        self.save_model(filepath, overwrite=True)
+                            print(f'\nEpoch {epoch + 1:05d}: {self.monitor} reached',
+                                  f' {current} (best {self.best}), saving model to',
+                                  f' {filepath} as top {self.save_top_k}')
+                        self._save_model(filepath, overwrite=False)
 
                     else:
                         if self.verbose > 0:
-                            print('\nEpoch %05d: %s did not improve' %
-                                  (epoch + 1, self.monitor))
+                            print(f'\nEpoch {epoch + 1:05d}: {self.monitor}',
+                                  f' was not in top {self.save_top_k}')
+
             else:
                 if self.verbose > 0:
-                    print('\nEpoch %05d: saving model to %s' % (epoch + 1, filepath))
-                self.save_model(filepath, overwrite=False)
-
-
-if __name__ == '__main__':
-    c = EarlyStopping(min_delta=0.9, patience=2, verbose=True)
-    losses = [10, 9, 8, 8, 6, 4.3, 5, 4.4, 2.8, 2.5]
-    for i, loss in enumerate(losses):
-        should_stop = c.on_epoch_end(i, logs={'val_loss': loss})
-        print(loss)
-        if should_stop:
-            break
+                    print(f'\nEpoch {epoch + 1:05d}: saving model to {filepath}')
+                self._save_model(filepath, overwrite=False)
diff --git a/tests/test_models.py b/tests/test_models.py
index 4813c542e8557..5ec5ab765816a 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -27,6 +27,85 @@
 # TESTS
 # ------------------------------------------------------------------------
 
+def test_model_checkpoint_options():
+    """
+    Test ModelCheckpoint options
+    :return:
+    """
+    def my_own_save_function(filepath):
+        open(filepath, 'a').close()
+
+    hparams = get_hparams()
+    model = LightningTestModel(hparams)
+
+    # simulated losses
+    save_dir = 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 = my_own_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
+
+    clear_save_dir()
+
+    # -----------------
+    # CASE K=0 (none)
+    w = ModelCheckpoint(save_dir, save_top_k=0, verbose=1)
+    w.save_function = my_own_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"
+
+    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 = my_own_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
+
+    clear_save_dir()
+
+    # -----------------
+    # CASE K=2 (2.5 epoch 4, 2.8 epoch 2)
+    # make sure other files don't get deleted
+    open(f'{save_dir}/other_file.ckpt', 'a').close()
+
+    w = ModelCheckpoint(save_dir, save_top_k=2, verbose=1)
+    w.save_function = my_own_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
+
+    clear_save_dir()
+
+
 def test_optimizer_return_options():
 
     trainer = Trainer()
@@ -98,7 +177,6 @@ def test_early_stopping_cpu_model():
     Test each of the trainer options
     :return:
     """
-
     stopping = EarlyStopping(monitor='val_loss')
     trainer_options = dict(
         early_stop_callback=stopping,


Code review comments:
diff_hunk: @@ -27,6 +27,85 @@
 # TESTS
 # ------------------------------------------------------------------------
 
+def test_model_checkpoint_options():
path: tests/test_models.py
text: @user1:
@author I updated the tests to reflect the desired behavior. Probably easier to modify the checkpoint accordingly until these tests pass.

@user1:
@author your fixes were off by one. the tests are correct. the implementation is off by 1 so needs to be fixed