change Checkpoint callback's save_best_only to save_top_k

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'

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

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'

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'

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'

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'

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'

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'

pytorch_lightning/callbacks/pt_callbacks.py

@@ -156,11 +156,11 @@ 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.
         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 +174,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.bestk = {}
+        # {epoch: monitor}
+        self.best = 0
 
         if mode not in ['auto', 'min', 'max']:
             print('ModelCheckpoint mode %s is unknown, '
@@ -193,17 +196,35 @@ def __init__(self, filepath, monitor='val_loss', verbose=0,
 
         if mode == 'min':
             self.monitor_op = np.less
-            self.best = np.Inf
+            self.kth = np.Inf
+            self.mode = 'min'
         elif mode == 'max':
             self.monitor_op = np.greater
-            self.best = -np.Inf
+            self.kth = -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 = -np.Inf
+                self.mode = 'max'
             else:
                 self.monitor_op = np.less
-                self.best = np.Inf
+                self.kth = np.Inf
+                self.mode = 'min'
+
+    def del_model(self, filepath):
+        dirpath = '/'.join(filepath.split('/')[:-1])
+
+        # make paths
+        os.makedirs(os.path.dirname(filepath), exist_ok=True)
+
+        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)
 
     def save_model(self, filepath, overwrite):
         dirpath = '/'.join(filepath.split('/')[:-1])
@@ -225,29 +246,46 @@ def save_model(self, filepath, overwrite):
 
     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
+        self.epochs_since_last_check += 1
+        if self.epochs_since_last_check >= self.period:
+            self.epochs_since_last_check = 0
             filepath = '{}/{}_ckpt_epoch_{}.ckpt'.format(self.filepath, self.prefix, epoch + 1)
-            if self.save_best_only:
+            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)
                 else:
-                    if self.monitor_op(current, self.best):
+                    if ((len(self.bestk.keys()) < self.save_top_k) or
+                            (self.monitor_op(current, self.bestk[self.kth]))):
+                        if len(self.bestk.keys()) == self.save_top_k:
+                            # need to pop the kth
+                            delpath = '{}/{}_ckpt_epoch_{}.ckpt'.format(
+                                self.filepath, self.prefix, self.kth + 1)
+                            self.bestk.pop(self.kth)
+                            self.del_model(delpath)
+                        self.bestk[epoch] = current
+                        if len(self.bestk.keys()) == self.save_top_k:
+                            # monitor dict has reached k elements
+                            if self.mode == 'min':
+                                self.kth = max(self.bestk, key=self.bestk.get)
+                            else:
+                                self.kth = min(self.bestk, key=self.bestk.get)
+                        if self.mode == 'min':
+                            self.best = min(self.bestk.values())
+                        else:
+                            self.best = max(self.bestk.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
+                            print('\nEpoch %05d: %s reached %s (best %s),'
+                                  ' saving model to %s as top %d'
+                                  % (epoch + 1, self.monitor, current, self.best,
+                                     filepath, self.save_top_k))
                         self.save_model(filepath, overwrite=True)
 
                     else:
                         if self.verbose > 0:
-                            print('\nEpoch %05d: %s did not improve' %
-                                  (epoch + 1, self.monitor))
+                            print('\nEpoch %05d: %s was not in top %d' %
+                                  (epoch + 1, self.monitor, self.save_top_k))
             else:
                 if self.verbose > 0:
                     print('\nEpoch %05d: saving model to %s' % (epoch + 1, filepath))
@@ -262,3 +300,6 @@ def on_epoch_end(self, epoch, logs=None):
         print(loss)
         if should_stop:
             break
+    w = ModelCheckpoint('res', save_top_k=2, verbose=1)
+    for i, loss in enumerate(losses):
+        w.on_epoch_end(i, logs={'val_loss': loss})