migrate checkpoint on load

wrap

Author: awaelchli

pytorch_lightning/plugins/training_type/training_type_plugin.py

@@ -26,6 +26,8 @@
 from pytorch_lightning.utilities import rank_zero_warn
 from pytorch_lightning.utilities.cloud_io import atomic_save
 from pytorch_lightning.utilities.cloud_io import load as pl_load
+from pytorch_lightning.utilities.migration.base import pl_legacy_patch
+from pytorch_lightning.utilities.migration.migrations import migrate_checkpoint
 
 TBroadcast = TypeVar("T")
 
@@ -216,7 +218,10 @@ def restore_model_state_from_ckpt_path(
             bool: Wether to load optimizer / lr_schedulers states from checkpoint
 
         """
-        ckpt = pl_load(ckpt_path, map_location=map_location)
+        with pl_legacy_patch():
+            ckpt = pl_load(ckpt_path, map_location=map_location)
+        ckpt = migrate_checkpoint(ckpt)
+
         # restore datamodule states
         if self.lightning_module.trainer.datamodule is not None:
             self.lightning_module.trainer.datamodule.on_load_checkpoint(ckpt)

pytorch_lightning/trainer/connectors/checkpoint_connector.py

@@ -32,6 +32,8 @@
 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.migration.base import pl_legacy_patch
+from pytorch_lightning.utilities.migration.migrations import migrate_checkpoint
 
 if _APEX_AVAILABLE:
     from apex import amp
@@ -89,9 +91,11 @@ def restore(self, checkpoint_path: str, on_gpu: bool) -> bool:
             rank_zero_warn("No checkpoint file exists at `resume_from_checkpoint`. Start from scratch")
             return False
 
-        checkpoint, load_optimizer_states = self.trainer.training_type_plugin.restore_model_state_from_ckpt_path(
-            checkpoint_path, map_location=lambda storage, loc: storage
-        )
+        with pl_legacy_patch():
+            checkpoint, load_optimizer_states = self.trainer.training_type_plugin.restore_model_state_from_ckpt_path(
+                checkpoint_path, map_location=lambda storage, loc: storage
+            )
+        migrate_checkpoint(checkpoint)
 
         model = self.trainer.lightning_module
 

pytorch_lightning/utilities/migration/base.py

@@ -26,4 +26,5 @@ def __enter__(self):
         return self
 
     def __exit__(self, exc_type, exc_value, exc_traceback):
-        delattr(pytorch_lightning.utilities.argparse, "_gpus_arg_default")
+        if hasattr(pytorch_lightning.utilities.argparse, "_gpus_arg_default"):
+            delattr(pytorch_lightning.utilities.argparse, "_gpus_arg_default")