diff --git a/pytorch_lightning/callbacks/model_checkpoint.py b/pytorch_lightning/callbacks/model_checkpoint.py index 4dc73e4b..298a7bb0 100644 --- a/pytorch_lightning/callbacks/model_checkpoint.py +++ b/pytorch_lightning/callbacks/model_checkpoint.py @@ -10,6 +10,7 @@ import os import re import numpy as np +from typing import Optional from pytorch_lightning import _logger as log from pytorch_lightning.callbacks.base import Callback @@ -37,6 +38,9 @@ class ModelCheckpoint(Callback): ... filepath='my/path/{epoch}-{val_loss:.2f}-{other_metric:.2f}' ... ) + Can also be set to `None`, then it will be set to default location + during trainer construction. + monitor: quantity to monitor. verbose: verbosity mode. Default: ``False``. save_top_k: if `save_top_k == k`, @@ -78,7 +82,7 @@ class ModelCheckpoint(Callback): """ - def __init__(self, filepath: str, monitor: str = 'val_loss', verbose: bool = False, + def __init__(self, filepath: Optional[str] = None, monitor: str = 'val_loss', verbose: bool = False, save_top_k: int = 1, save_weights_only: bool = False, mode: str = 'auto', period: int = 1, prefix: str = ''): super().__init__() @@ -90,12 +94,14 @@ class ModelCheckpoint(Callback): self.monitor = monitor self.verbose = verbose - if os.path.isdir(filepath): - self.dirpath, self.filename = filepath, '{epoch}' + if filepath is None: # will be determined by trainer at runtime + self.dirpath, self.filename = None, None else: - self.dirpath, self.filename = os.path.split(filepath) - - os.makedirs(self.dirpath, exist_ok=True) + if os.path.isdir(filepath): + self.dirpath, self.filename = filepath, '{epoch}' + else: + self.dirpath, self.filename = os.path.split(filepath) + os.makedirs(self.dirpath, exist_ok=True) self.save_top_k = save_top_k self.save_weights_only = save_weights_only self.period = period diff --git a/pytorch_lightning/trainer/callback_config.py b/pytorch_lightning/trainer/callback_config.py index c288cdd6..7ace0fb2 100644 --- a/pytorch_lightning/trainer/callback_config.py +++ b/pytorch_lightning/trainer/callback_config.py @@ -33,7 +33,7 @@ class TrainerCallbackConfigMixin(ABC): Otherwise use os.getcwd() """ ckpt_path = self.default_root_dir - if self.checkpoint_callback is True: + if self.checkpoint_callback: # init a default one if self.logger is not None: save_dir = (getattr(self.logger, 'save_dir', None) or @@ -57,12 +57,18 @@ class TrainerCallbackConfigMixin(ABC): train_step_only = not self.is_overriden('validation_step') monitor_key = 'loss' if train_step_only else 'val_loss' - self.ckpt_path = ckpt_path - os.makedirs(ckpt_path, exist_ok=True) - self.checkpoint_callback = ModelCheckpoint( - filepath=ckpt_path, - monitor=monitor_key - ) + if self.checkpoint_callback is True: + os.makedirs(ckpt_path, exist_ok=True) + self.checkpoint_callback = ModelCheckpoint( + filepath=ckpt_path, + monitor=monitor_key + ) + # If user specified None in filepath, override with runtime default + elif isinstance(self.checkpoint_callback, ModelCheckpoint) \ + and self.checkpoint_callback.dirpath is None: + self.checkpoint_callback.dirpath = ckpt_path + self.checkpoint_callback.filename = '{epoch}' + os.makedirs(self.checkpoint_callback.dirpath, exist_ok=True) elif self.checkpoint_callback is False: self.checkpoint_callback = None diff --git a/tests/trainer/test_callbacks.py b/tests/trainer/test_callbacks.py index ee208059..4f77dabb 100644 --- a/tests/trainer/test_callbacks.py +++ b/tests/trainer/test_callbacks.py @@ -1,7 +1,7 @@ import tests.base.utils as tutils from pytorch_lightning import Callback from pytorch_lightning import Trainer, LightningModule -from pytorch_lightning.callbacks import EarlyStopping +from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint from tests.base import ( LightTrainDataloader, LightTestMixin, @@ -181,3 +181,27 @@ def test_early_stopping_no_val_step(tmpdir): assert result == 1, 'training failed to complete' assert trainer.current_epoch < trainer.max_epochs + + +def test_model_checkpoint_with_non_string_input(tmpdir): + """ Test that None in checkpoint callback is valid and that chkp_path is + set correctly """ + tutils.reset_seed() + + class CurrentTestModel(LightTrainDataloader, TestModelBase): + pass + + hparams = tutils.get_default_hparams() + model = CurrentTestModel(hparams) + + checkpoint = ModelCheckpoint(filepath=None, save_top_k=-1) + + trainer = Trainer(default_root_dir=tmpdir, + checkpoint_callback=checkpoint, + overfit_pct=0.20, + max_epochs=5 + ) + result = trainer.fit(model) + + # These should be different if the dirpath has be overridden + assert trainer.ckpt_path != trainer.default_root_dir