add weights_only arg to checkpoint save. weights_only during test set based on ckpt version

Author: matsumotosan

src/lightning/pytorch/core/saving.py

@@ -56,11 +56,17 @@ def _load_from_checkpoint(
     map_location: _MAP_LOCATION_TYPE = None,
     hparams_file: Optional[_PATH] = None,
     strict: Optional[bool] = None,
+    weights_only: Optional[bool] = None,
     **kwargs: Any,
 ) -> Union["pl.LightningModule", "pl.LightningDataModule"]:
     map_location = map_location or _default_map_location
+
+    if weights_only is None:
+        log.debug("`weights_only` not specified, defaulting to `True`.")
+        weights_only = True
+
     with pl_legacy_patch():
-        checkpoint = pl_load(checkpoint_path, map_location=map_location)
+        checkpoint = pl_load(checkpoint_path, map_location=map_location, weights_only=weights_only)
 
     # convert legacy checkpoints to the new format
     checkpoint = _pl_migrate_checkpoint(

src/lightning/pytorch/trainer/connectors/checkpoint_connector.py

@@ -414,7 +414,9 @@ def dump_checkpoint(self, weights_only: bool = False) -> dict:
         """Creating a model checkpoint dictionary object from various component states.
 
         Args:
-            weights_only: saving model weights only
+            weights_only: If True, only saves model and loops state_dict objects. If False,
+            additionally saves callbacks, optimizers, schedulers, and precision plugin states.
+
         Return:
             structured dictionary: {
                 'epoch':                     training epoch

tests/tests_pytorch/checkpointing/test_legacy_checkpoints.py

@@ -18,6 +18,7 @@
 
 import pytest
 import torch
+from packaging.version import Version
 
 import lightning.pytorch as pl
 from lightning.pytorch import Callback, Trainer
@@ -45,7 +46,12 @@ def test_load_legacy_checkpoints(tmp_path, pl_version: str):
         assert path_ckpts, f'No checkpoints found in folder "{PATH_LEGACY}"'
         path_ckpt = path_ckpts[-1]
 
-        model = ClassificationModel.load_from_checkpoint(path_ckpt, num_features=24)
+        # legacy load utility added in 1.5.0 (see https://github.com/Lightning-AI/pytorch-lightning/pull/9166)
+        if pl_version == "local":
+            pl_version = pl.__version__
+        weights_only = not Version(pl_version) < Version("1.5.0")
+
+        model = ClassificationModel.load_from_checkpoint(path_ckpt, num_features=24, weights_only=weights_only)
         trainer = Trainer(default_root_dir=tmp_path)
         dm = ClassifDataModule(num_features=24, length=6000, batch_size=128, n_clusters_per_class=2, n_informative=8)
         res = trainer.test(model, datamodule=dm)
@@ -73,13 +79,18 @@ def test_legacy_ckpt_threading(pl_version: str):
     assert path_ckpts, f'No checkpoints found in folder "{PATH_LEGACY}"'
     path_ckpt = path_ckpts[-1]
 
+    # legacy load utility added in 1.5.0 (see https://github.com/Lightning-AI/pytorch-lightning/pull/9166)
+    if pl_version == "local":
+        pl_version = pl.__version__
+    weights_only = not Version(pl_version) < Version("1.5.0")
+
     def load_model():
         import torch
 
         from lightning.pytorch.utilities.migration import pl_legacy_patch
 
         with pl_legacy_patch():
-            _ = torch.load(path_ckpt, weights_only=False)
+            _ = torch.load(path_ckpt, weights_only=weights_only)
 
     with patch("sys.path", [PATH_LEGACY] + sys.path):
         t1 = ThreadExceptionHandler(target=load_model)