Keep track of the best model's path saved by ModelCheckpoint

CHANGELOG.md

@@ -10,6 +10,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Added type hints in `Trainer.fit()` and `Trainer.test()` to reflect that also a list of dataloaders can be passed in ([#1723](https://github.com/PyTorchLightning/pytorch-lightning/pull/1723)).
 
+- Attribute `best_model` to `ModelCheckpoint` for storing and later retrieving the path to the best saved model file ([#1799](https://github.com/PyTorchLightning/pytorch-lightning/pull/1799)) 
+
 ### Changed
 
 - Allow user to select individual TPU core to train on ([#1729](https://github.com/PyTorchLightning/pytorch-lightning/pull/1729))

pytorch_lightning/callbacks/model_checkpoint.py

@@ -20,7 +20,10 @@
 
 class ModelCheckpoint(Callback):
     r"""
-    Save the model after every epoch.
+    Save the model after every epoch if it improves.
+
+    After training finishes, use :attr:`best_model` to retrieve the path to the
+    best checkpoint file.
 
     Args:
         filepath: path to save the model file.
@@ -80,6 +83,13 @@ class ModelCheckpoint(Callback):
         ...     filepath='my/path/sample-mnist_{epoch:02d}-{val_loss:.2f}'
         ... )
 
+        # retrieve the best checkpoint after training
+        checkpoint_callback = ModelCheckpoint(filepath='my/path/')
+        trainer = Trainer(checkpoint_callback=checkpoint_callback)
+        model = ...
+        trainer.fit(model)
+        checkpoint_callback.best_model
+
     """
 
     def __init__(self, filepath: Optional[str] = None, monitor: str = 'val_loss', verbose: bool = False,
@@ -112,6 +122,7 @@ def __init__(self, filepath: Optional[str] = None, monitor: str = 'val_loss', ve
         # {filename: monitor}
         self.kth_best_model = ''
         self.best = 0
+        self.best_model = ''
         self.save_function = None
 
         torch_inf = torch.tensor(np.Inf)
@@ -265,7 +276,8 @@ def _do_check_save(self, filepath, current, epoch):
             self.kth_value = self.best_k_models[self.kth_best_model]
 
         _op = min if self.mode == 'min' else max
-        self.best = _op(self.best_k_models.values())
+        self.best_model = _op(self.best_k_models, key=self.best_k_models.get)
+        self.best = self.best_k_models[self.best_model]
 
         if self.verbose > 0:
             log.info(

pytorch_lightning/trainer/training_io.py

@@ -315,6 +315,7 @@ def dump_checkpoint(self, weights_only: bool = False):
         if not weights_only:
             if self.checkpoint_callback:
                 checkpoint['checkpoint_callback_best'] = self.checkpoint_callback.best
+                checkpoint['checkpoint_callback_best_model'] = self.checkpoint_callback.best_model
 
             if self.early_stop_callback:
                 checkpoint['early_stop_callback_wait'] = self.early_stop_callback.wait
@@ -398,10 +399,11 @@ def restore_training_state(self, checkpoint):
                 ' This is probably due to `ModelCheckpoint.save_weights_only` being set to `True`.'
             )
 
-        if self.checkpoint_callback is not None and self.checkpoint_callback is not False:
+        if self.checkpoint_callback:
             self.checkpoint_callback.best = checkpoint['checkpoint_callback_best']
+            self.checkpoint_callback.best_model = checkpoint['checkpoint_callback_best_model']
 
-        if self.early_stop_callback is not None and self.early_stop_callback is not False:
+        if self.early_stop_callback:
             self.early_stop_callback.wait = checkpoint['early_stop_callback_wait']
             self.early_stop_callback.patience = checkpoint['early_stop_callback_patience']