Default value for ModelCheckpoint filepath (Lightning-AI#1548)

* allow determine of filepath at runtime

* typing

Co-authored-by: Nicki Skafte <nugginea@gmail.com>

Author: SkafteNicki

pytorch_lightning/callbacks/model_checkpoint.py

@@ -10,6 +10,7 @@
 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 @@ def __init__(self, filepath: str, monitor: str = 'val_loss', verbose: bool = Fal
 
         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

pytorch_lightning/trainer/callback_config.py

@@ -33,7 +33,7 @@ def configure_checkpoint_callback(self):
         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 @@ def configure_checkpoint_callback(self):
             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
 

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 training_step(self, *args, **kwargs):
 
     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