suboptimal

Author: awaelchli

pytorch_lightning/core/saving.py

@@ -30,6 +30,7 @@
 from pytorch_lightning.utilities.apply_func import apply_to_collection
 from pytorch_lightning.utilities.cloud_io import get_filesystem
 from pytorch_lightning.utilities.cloud_io import load as pl_load
+from pytorch_lightning.utilities.migration.migrations import migrate_checkpoint
 from pytorch_lightning.utilities.parsing import parse_class_init_keys
 
 log = logging.getLogger(__name__)
@@ -134,6 +135,9 @@ def load_from_checkpoint(
         else:
             checkpoint = pl_load(checkpoint_path, map_location=lambda storage, loc: storage)
 
+        # convert legacy checkpoints to the new format
+        checkpoint = migrate_checkpoint(checkpoint)
+
         if hparams_file is not None:
             extension = hparams_file.split('.')[-1]
             if extension.lower() == 'csv':
@@ -148,6 +152,7 @@ def load_from_checkpoint(
             # overwrite hparams by the given file
             checkpoint[cls.CHECKPOINT_HYPER_PARAMS_KEY] = hparams
 
+        # TODO: make this a migration:
         # for past checkpoint need to add the new key
         if cls.CHECKPOINT_HYPER_PARAMS_KEY not in checkpoint:
             checkpoint[cls.CHECKPOINT_HYPER_PARAMS_KEY] = {}
@@ -171,6 +176,7 @@ def _load_model_state(cls, checkpoint: Dict[str, Any], strict: bool = True, **cl
         if cls.CHECKPOINT_HYPER_PARAMS_KEY in checkpoint:
 
             # 1. (backward compatibility) Try to restore model hparams from checkpoint using old/past keys
+            # TODO: make this a migration:
             for _old_hparam_key in CHECKPOINT_PAST_HPARAMS_KEYS:
                 cls_kwargs_loaded.update(checkpoint.get(_old_hparam_key, {}))
 

pytorch_lightning/trainer/connectors/checkpoint_connector.py

@@ -32,7 +32,6 @@
 from pytorch_lightning.utilities.cloud_io import atomic_save, get_filesystem
 from pytorch_lightning.utilities.cloud_io import load as pl_load
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
-from pytorch_lightning.utilities.upgrade_checkpoint import KEYS_MAPPING as DEPRECATED_CHECKPOINT_KEYS
 
 if _APEX_AVAILABLE:
     from apex import amp
@@ -134,14 +133,6 @@ def restore_training_state(self, checkpoint, load_optimizer_states: bool = True)
                 ' This is probably due to `ModelCheckpoint.save_weights_only` being set to `True`.'
             )
 
-        if any([key in checkpoint for key in DEPRECATED_CHECKPOINT_KEYS]):
-            raise ValueError(
-                "The checkpoint you're attempting to load follows an"
-                " outdated schema. You can upgrade to the current schema by running"
-                " `python -m pytorch_lightning.utilities.upgrade_checkpoint --file model.ckpt`"
-                " where `model.ckpt` is your checkpoint file."
-            )
-
         # restore amp scaling
         if self.trainer.amp_backend == AMPType.NATIVE and 'native_amp_scaling_state' in checkpoint:
             self.trainer.scaler.load_state_dict(checkpoint['native_amp_scaling_state'])

pytorch_lightning/utilities/argparse.py

@@ -286,11 +286,11 @@ def _gpus_allowed_type(x) -> Union[int, str]:
         return int(x)
 
 
-def _gpus_arg_default(x) -> Union[int, str]:  # pragma: no-cover
-    # unused, but here for backward compatibility with old checkpoints that need to be able to
-    # unpickle the function from the checkpoint, as it was not filtered out in versions < 1.2.8
-    # see: https://github.com/PyTorchLightning/pytorch-lightning/pull/6898
-    pass
+# def _gpus_arg_default(x) -> Union[int, str]:  # pragma: no-cover
+#     # unused, but here for backward compatibility with old checkpoints that need to be able to
+#     # unpickle the function from the checkpoint, as it was not filtered out in versions < 1.2.8
+#     # see: https://github.com/PyTorchLightning/pytorch-lightning/pull/6898
+#     pass
 
 
 def _int_or_float_type(x) -> Union[int, float]:

pytorch_lightning/utilities/migration/__init__.py

@@ -0,0 +1,29 @@
+from distutils.version import LooseVersion
+
+import pytorch_lightning.utilities.argparse
+
+
+def get_version(checkpoint: dict) -> str:
+    return checkpoint["pytorch-lightning_version"]
+
+
+def set_version(checkpoint: dict, version: str):
+    checkpoint["pytorch-lightning_version"] = version
+
+
+def should_upgrade(checkpoint: dict, target: str) -> bool:
+    return LooseVersion(get_version(checkpoint)) < LooseVersion(target)
+
+
+class pl_legacy_patch:
+    """
+    Registers legacy artifacts (classes, methods, etc.) that were removed but still need to be
+    included for unpickling old checkpoints.
+    """
+
+    def __enter__(self):
+        setattr(pytorch_lightning.utilities.argparse, "_gpus_arg_default", lambda x: x)
+        return self
+
+    def __exit__(self, exc_type, exc_value, exc_traceback):
+        delattr(pytorch_lightning.utilities.argparse, "_gpus_arg_default")

pytorch_lightning/utilities/migration/base.py

@@ -0,0 +1,45 @@
+import pytorch_lightning as pl
+from pytorch_lightning.utilities.migration.base import set_version, should_upgrade
+
+
+# v0.10.0
+def migrate_model_checkpoint_early_stopping(checkpoint: dict) -> dict:
+    from pytorch_lightning.callbacks.early_stopping import EarlyStopping
+    from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint
+    keys_mapping = {
+        "checkpoint_callback_best_model_score": (ModelCheckpoint, "best_model_score"),
+        "checkpoint_callback_best_model_path": (ModelCheckpoint, "best_model_path"),
+        "checkpoint_callback_best": (ModelCheckpoint, "best_model_score"),
+        "early_stop_callback_wait": (EarlyStopping, "wait_count"),
+        "early_stop_callback_patience": (EarlyStopping, "patience"),
+    }
+    checkpoint["callbacks"] = checkpoint.get("callbacks") or {}
+
+    for key, new_path in keys_mapping.items():
+        if key in checkpoint:
+            value = checkpoint[key]
+            callback_type, callback_key = new_path
+            checkpoint["callbacks"][callback_type] = checkpoint["callbacks"].get(callback_type) or {}
+            checkpoint["callbacks"][callback_type][callback_key] = value
+            del checkpoint[key]
+    return checkpoint
+
+
+# v1.3.1
+def migrate_callback_state_identifiers(checkpoint):
+    if "callbacks" not in checkpoint:
+        return
+    callbacks = checkpoint["callbacks"]
+    checkpoint["callbacks"] = dict((callback_type.__name__, state) for callback_type, state in callbacks.items())
+    return checkpoint
+
+
+def migrate_checkpoint(checkpoint: dict):
+    """ Applies all the above migrations in order. """
+    if should_upgrade(checkpoint, "0.10.0"):
+        migrate_model_checkpoint_early_stopping(checkpoint)
+    if should_upgrade(checkpoint, "1.3.0"):
+        migrate_callback_state_identifiers(checkpoint)
+        set_version(checkpoint, "1.3.0")
+    set_version(checkpoint, pl.__version__)
+    return checkpoint

pytorch_lightning/utilities/migration/migrations.py

@@ -2,6 +2,7 @@
 
 import torch
 from torchmetrics import Metric
+
 """
 Convention:
  - Do not include any `_TYPE` suffix

pytorch_lightning/utilities/types.py

@@ -17,34 +17,11 @@
 
 import torch
 
-from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
-
-KEYS_MAPPING = {
-    "checkpoint_callback_best_model_score": (ModelCheckpoint, "best_model_score"),
-    "checkpoint_callback_best_model_path": (ModelCheckpoint, "best_model_path"),
-    "checkpoint_callback_best": (ModelCheckpoint, "best_model_score"),
-    "early_stop_callback_wait": (EarlyStopping, "wait_count"),
-    "early_stop_callback_patience": (EarlyStopping, "patience"),
-}
+from pytorch_lightning.utilities.migration.base import pl_legacy_patch
+from pytorch_lightning.utilities.migration.migrations import migrate_checkpoint
 
 log = logging.getLogger(__name__)
 
-
-def upgrade_checkpoint(filepath):
-    checkpoint = torch.load(filepath)
-    checkpoint["callbacks"] = checkpoint.get("callbacks") or {}
-
-    for key, new_path in KEYS_MAPPING.items():
-        if key in checkpoint:
-            value = checkpoint[key]
-            callback_type, callback_key = new_path
-            checkpoint["callbacks"][callback_type] = checkpoint["callbacks"].get(callback_type) or {}
-            checkpoint["callbacks"][callback_type][callback_key] = value
-            del checkpoint[key]
-
-    torch.save(checkpoint, filepath)
-
-
 if __name__ == "__main__":
 
     parser = argparse.ArgumentParser(
@@ -57,4 +34,7 @@ def upgrade_checkpoint(filepath):
 
     log.info("Creating a backup of the existing checkpoint file before overwriting in the upgrade process.")
     copyfile(args.file, args.file + ".bak")
-    upgrade_checkpoint(args.file)
+    with pl_legacy_patch():
+        checkpoint = torch.load(args.file)
+    migrate_checkpoint(checkpoint)
+    torch.save(checkpoint, args.file)

pytorch_lightning/utilities/upgrade_checkpoint.py

@@ -11,13 +11,11 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import os
-
 import pytest
-import torch
+import pytorch_lightning as pl
 
-from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
-from pytorch_lightning.utilities.upgrade_checkpoint import upgrade_checkpoint
+from pytorch_lightning.utilities.migration.base import set_version, get_version
+from pytorch_lightning.utilities.migration.migrations import migrate_checkpoint
 
 
 @pytest.mark.parametrize(
@@ -33,7 +31,7 @@
                 "epoch": 1,
                 "global_step": 23,
                 "callbacks": {
-                    ModelCheckpoint: {
+                    "ModelCheckpoint": {
                         "best_model_score": 0.34
                     }
                 }
@@ -49,7 +47,7 @@
                 "epoch": 1,
                 "global_step": 23,
                 "callbacks": {
-                    ModelCheckpoint: {
+                    "ModelCheckpoint": {
                         "best_model_score": 0.99
                     }
                 }
@@ -65,7 +63,7 @@
                 "epoch": 1,
                 "global_step": 23,
                 "callbacks": {
-                    ModelCheckpoint: {
+                    "ModelCheckpoint": {
                         "best_model_path": 'path'
                     }
                 }
@@ -82,7 +80,7 @@
                 "epoch": 1,
                 "global_step": 23,
                 "callbacks": {
-                    EarlyStopping: {
+                    "EarlyStopping": {
                         "wait_count": 2,
                         "patience": 4
                     }
@@ -92,8 +90,8 @@
     ],
 )
 def test_upgrade_checkpoint(tmpdir, old_checkpoint, new_checkpoint):
-    filepath = os.path.join(tmpdir, "model.ckpt")
-    torch.save(old_checkpoint, filepath)
-    upgrade_checkpoint(filepath)
-    updated_checkpoint = torch.load(filepath)
+    set_version(old_checkpoint, "0.9.0")
+    set_version(new_checkpoint, pl.__version__)
+    updated_checkpoint = migrate_checkpoint(old_checkpoint)
     assert updated_checkpoint == new_checkpoint
+    assert get_version(updated_checkpoint) == pl.__version__