Keep track of the best model's path saved by ModelCheckpoint (#1799)

* Add an additional attribute to ModelCheckpoint to keep track of the best model's path

Currently, only the best metric value is directly tracked. This new attribute will help in uses cases where the trained model needs to be used or tracked right after training.

* Add small description and usage example to docs

* Fix PEP8 issues

* Fix doctest example

* Fix expected output in doctest

* Apply suggestions from code review

* Show example as code block instead of doctest

* Apply suggestions from code review

* Update CHANGELOG.md

* Rename `ModelCheckpoint.best` to `ModelCheckpoint.best_model_score`

Also rename `ModelCheckpoint.best_model` (added in this PR) to `ModelCheckpoint.best_model_path`, for consistency, and `kth_best_model` to `kth_best_model_path`.

* Update pytorch_lightning/trainer/training_io.py

Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com>

* Apply suggestions from code review

Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com>

* Add warning when loading checkpoint from an old version

Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com>
This commit is contained in:
Fabio Natanael Kepler
2020-05-31 08:47:13 -04:00
committed by GitHub
co-authored by Jirka Borovec
parent 55fdfe3845
commit 8b9b923ca8
3 changed files with 55 additions and 15 deletions
+6
View File
@@ -18,6 +18,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
- Early stopping checks `on_validation_end` ([#1458](https://github.com/PyTorchLightning/pytorch-lightning/pull/1458))
- Attribute `best_model_path` 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))
@@ -26,10 +28,14 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
- Allow passing model hyperparameters as complete kwarg list ([#1896](https://github.com/PyTorchLightning/pytorch-lightning/pull/1896))
- Renamed `ModelCheckpoint`'s attributes `best` to `best_model_score` and `kth_best_model` to `kth_best_model_path` ([#1799](https://github.com/PyTorchLightning/pytorch-lightning/pull/1799))
- Re-Enable Logger's `ImportError`s ([#1938](https://github.com/PyTorchLightning/pytorch-lightning/pull/1938))
### Deprecated
- Deprecated `ModelCheckpoint`'s attributes `best` and `kth_best_model` ([#1799](https://github.com/PyTorchLightning/pytorch-lightning/pull/1799))
- Dropped official support/testing for older PyTorch versions <1.3 ([#1917](https://github.com/PyTorchLightning/pytorch-lightning/pull/1917))
### Removed
+35 -11
View File
@@ -20,7 +20,10 @@ from pytorch_lightning.utilities import rank_zero_warn, rank_zero_only
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_path` to retrieve the path to the
best checkpoint file and :attr:`best_model_score` to retrieve its score.
Args:
filepath: path to save the model file.
@@ -81,6 +84,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_path
"""
def __init__(self, filepath: Optional[str] = None, monitor: str = 'val_loss', verbose: bool = False,
@@ -112,8 +122,9 @@ class ModelCheckpoint(Callback):
self.prefix = prefix
self.best_k_models = {}
# {filename: monitor}
self.kth_best_model = ''
self.best = 0
self.kth_best_model_path = ''
self.best_model_score = 0
self.best_model_path = ''
self.save_function = None
torch_inf = torch.tensor(np.Inf)
@@ -131,6 +142,18 @@ class ModelCheckpoint(Callback):
self.kth_value, self.mode = mode_dict[mode]
@property
def best(self):
rank_zero_warn("Attribute `best` has been renamed to `best_model_score` since v0.8.0"
" and will be removed in v0.10.0", DeprecationWarning)
return self.best_model_score
@property
def kth_best_model(self):
rank_zero_warn("Attribute `kth_best_model` has been renamed to `kth_best_model_path` since v0.8.0"
" and will be removed in v0.10.0", DeprecationWarning)
return self.kth_best_model_path
def _del_model(self, filepath):
if os.path.isfile(filepath):
os.remove(filepath)
@@ -162,7 +185,7 @@ class ModelCheckpoint(Callback):
"max": torch.gt,
}[self.mode]
return monitor_op(current, self.best_k_models[self.kth_best_model])
return monitor_op(current, self.best_k_models[self.kth_best_model_path])
def format_checkpoint_name(self, epoch, metrics, ver=None):
"""Generate a filename according to the defined template.
@@ -258,25 +281,26 @@ class ModelCheckpoint(Callback):
del_list = []
if len(self.best_k_models) == self.save_top_k and self.save_top_k > 0:
delpath = self.kth_best_model
self.best_k_models.pop(self.kth_best_model)
delpath = self.kth_best_model_path
self.best_k_models.pop(self.kth_best_model_path)
del_list.append(delpath)
self.best_k_models[filepath] = current
if len(self.best_k_models) == self.save_top_k:
# monitor dict has reached k elements
_op = max if self.mode == 'min' else min
self.kth_best_model = _op(self.best_k_models,
key=self.best_k_models.get)
self.kth_value = self.best_k_models[self.kth_best_model]
self.kth_best_model_path = _op(self.best_k_models,
key=self.best_k_models.get)
self.kth_value = self.best_k_models[self.kth_best_model_path]
_op = min if self.mode == 'min' else max
self.best = _op(self.best_k_models.values())
self.best_model_path = _op(self.best_k_models, key=self.best_k_models.get)
self.best_model_score = self.best_k_models[self.best_model_path]
if self.verbose > 0:
log.info(
f'\nEpoch {epoch:05d}: {self.monitor} reached'
f' {current:0.5f} (best {self.best:0.5f}), saving model to'
f' {current:0.5f} (best {self.best_model_score:0.5f}), saving model to'
f' {filepath} as top {self.save_top_k}')
self._save_model(filepath)
+14 -4
View File
@@ -330,7 +330,8 @@ class TrainerIOMixin(ABC):
if not weights_only:
if self.checkpoint_callback:
checkpoint['checkpoint_callback_best'] = self.checkpoint_callback.best
checkpoint['checkpoint_callback_best_model_score'] = self.checkpoint_callback.best_model_score
checkpoint['checkpoint_callback_best_model_path'] = self.checkpoint_callback.best_model_path
if self.early_stop_callback:
checkpoint['early_stop_callback_wait'] = self.early_stop_callback.wait
@@ -401,10 +402,19 @@ class TrainerIOMixin(ABC):
' 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:
self.checkpoint_callback.best = checkpoint['checkpoint_callback_best']
if self.checkpoint_callback:
if 'checkpoint_callback_best_model_score' in checkpoint:
self.checkpoint_callback.best_model_score = checkpoint['checkpoint_callback_best_model_score']
else:
# Old naming until version 0.7.6
rank_zero_warn(
'Loading a checkpoint created with an old version of Lightning; '
'this will not be supported in the future.'
)
self.checkpoint_callback.best_model_score = checkpoint['checkpoint_callback_best']
self.checkpoint_callback.best_model_path = checkpoint['checkpoint_callback_best_model_path']
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']