[1 / 3] improvements to saving and loading callback state

pytorch_lightning/callbacks/base.py

@@ -17,7 +17,7 @@
 """
 
 import abc
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Type
 
 from pytorch_lightning.core.lightning import LightningModule
 
@@ -29,6 +29,14 @@ class Callback(abc.ABC):
     Subclass this class and override any of the relevant hooks
     """
 
+    @property
+    def state_identifier(self) -> str:
+        return self.__class__.__name__
+
+    @property
+    def _legacy_state_identifier(self) -> Type:
+        return type(self)
+
     def on_configure_sharded_model(self, trainer, pl_module: LightningModule) -> None:
         """Called before configure sharded model"""
 

pytorch_lightning/trainer/callback_hook.py

@@ -14,6 +14,7 @@
 
 from abc import ABC
 from copy import deepcopy
+from distutils.version import LooseVersion
 from inspect import signature
 from typing import Any, Callable, Dict, List, Optional, Type
 
@@ -243,7 +244,7 @@ def __is_old_signature(fn: Callable) -> bool:
             return True
         return False
 
-    def on_save_checkpoint(self, checkpoint: Dict[str, Any]) -> Dict[Type, dict]:
+    def on_save_checkpoint(self, checkpoint: Dict[str, Any]) -> Dict[str, dict]:
         """Called when saving a model checkpoint."""
         callback_states = {}
         for callback in self.callbacks:
@@ -257,18 +258,22 @@ def on_save_checkpoint(self, checkpoint: Dict[str, Any]) -> Dict[Type, dict]:
             else:
                 state = callback.on_save_checkpoint(self, self.lightning_module, checkpoint)
             if state:
-                callback_states[type(callback)] = state
+                callback_states[callback.state_identifier] = state
         return callback_states
 
-    def on_load_checkpoint(self, checkpoint):
+    def on_load_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
         """Called when loading a model checkpoint."""
         callback_states = checkpoint.get('callbacks')
+        version = checkpoint.get('pytorch-lightning_version')
         # Todo: the `callback_states` are dropped with TPUSpawn as they
         # can't be saved using `xm.save`
         # https://github.com/pytorch/xla/issues/2773
         if callback_states is not None:
             for callback in self.callbacks:
-                state = callback_states.get(type(callback))
+                state = (
+                    callback_states.get(callback.state_identifier)
+                    or callback_states.get(callback._legacy_state_identifier)
+                )
                 if state:
                     state = deepcopy(state)
                     callback.on_load_checkpoint(state)

pytorch_lightning/trainer/connectors/checkpoint_connector.py

@@ -244,7 +244,7 @@ def dump_checkpoint(self, weights_only: bool = False) -> dict:
             structured dictionary: {
                 'epoch':                     training epoch
                 'global_step':               training global step
-                'pytorch-lightning_version': PyTorch Lightning's version
+                'pytorch-lightning_version': The version of PyTorch Lightning that produced this checkpoint
                 'callbacks':                 "callback specific state"[] # if not weights_only
                 'optimizer_states':          "PT optim's state_dict"[]   # if not weights_only
                 'lr_schedulers':             "PT sched's state_dict"[]   # if not weights_only

tests/callbacks/test_early_stopping.py

@@ -76,7 +76,7 @@ def test_resume_early_stopping_from_checkpoint(tmpdir):
     # the checkpoint saves "epoch + 1"
     early_stop_callback_state = early_stop_callback.saved_states[checkpoint["epoch"] - 1]
     assert 4 == len(early_stop_callback.saved_states)
-    assert checkpoint["callbacks"][type(early_stop_callback)] == early_stop_callback_state
+    assert checkpoint["callbacks"]["EarlyStoppingTestRestore"] == early_stop_callback_state
 
     # ensure state is reloaded properly (assertion in the callback)
     early_stop_callback = EarlyStoppingTestRestore(early_stop_callback_state, monitor='train_loss')

tests/checkpointing/test_legacy_checkpoints.py

@@ -60,6 +60,7 @@
         "1.2.5",
         "1.2.6",
         "1.2.7",
+        "1.2.8",
     ]
 )
 def test_resume_legacy_checkpoints(tmpdir, pl_version: str):

tests/checkpointing/test_model_checkpoint.py

@@ -147,7 +147,7 @@ def configure_optimizers(self):
         assert chk['epoch'] == epoch + 1
         assert chk['global_step'] == limit_train_batches * (epoch + 1)
 
-        mc_specific_data = chk['callbacks'][type(checkpoint)]
+        mc_specific_data = chk['callbacks']["ModelCheckpoint"]
         assert mc_specific_data['dirpath'] == checkpoint.dirpath
         assert mc_specific_data['monitor'] == monitor
         assert mc_specific_data['current_score'] == score
@@ -251,7 +251,7 @@ def configure_optimizers(self):
             assert chk['epoch'] == epoch + 1
             assert chk['global_step'] == per_epoch_steps * (global_ix + 1)
 
-            mc_specific_data = chk['callbacks'][type(checkpoint)]
+            mc_specific_data = chk['callbacks']["ModelCheckpoint"]
             assert mc_specific_data['dirpath'] == checkpoint.dirpath
             assert mc_specific_data['monitor'] == monitor
             assert mc_specific_data['current_score'] == score
@@ -840,7 +840,7 @@ def test_model_checkpoint_save_last_checkpoint_contents(tmpdir):
     ckpt_last = torch.load(path_last)
     assert all(ckpt_last_epoch[k] == ckpt_last[k] for k in ("epoch", "global_step"))
 
-    ch_type = type(model_checkpoint)
+    ch_type = "ModelCheckpoint"
     assert ckpt_last["callbacks"][ch_type] == ckpt_last_epoch["callbacks"][ch_type]
 
     # it is easier to load the model objects than to iterate over the raw dict of tensors
@@ -1098,7 +1098,7 @@ def training_step(self, *args):
     trainer.fit(TestModel())
     assert model_checkpoint.current_score == 0.3
     ckpts = [torch.load(str(ckpt)) for ckpt in tmpdir.listdir()]
-    ckpts = [ckpt["callbacks"][type(model_checkpoint)] for ckpt in ckpts]
+    ckpts = [ckpt["callbacks"]["ModelCheckpoint"] for ckpt in ckpts]
     assert sorted(ckpt["current_score"] for ckpt in ckpts) == [0.1, 0.2, 0.3]
 
 

tests/trainer/connectors/test_callback_connector.py

@@ -68,11 +68,11 @@ def test_all_callback_states_saved_before_checkpoint_callback(tmpdir):
     trainer.fit(model)
 
     ckpt = torch.load(str(tmpdir / "all_states.ckpt"))
-    state0 = ckpt["callbacks"][type(callback0)]
-    state1 = ckpt["callbacks"][type(callback1)]
+    state0 = ckpt["callbacks"]["StatefulCallback0"]
+    state1 = ckpt["callbacks"]["StatefulCallback1"]
     assert "content0" in state0 and state0["content0"] == 0
     assert "content1" in state1 and state1["content1"] == 1
-    assert type(checkpoint_callback) in ckpt["callbacks"]
+    assert "ModelCheckpoint" in ckpt["callbacks"]
 
 
 def test_attach_model_callbacks():

tests/trainer/logging_/test_logger_connector.py

@@ -272,7 +272,7 @@ def test_dataloader(self):
 
 def test_call_back_validator(tmpdir):
 
-    funcs_name = sorted([f for f in dir(Callback) if not f.startswith('_')])
+    funcs_name = sorted([f for f in dir(Callback) if not f.startswith('_') and callable(getattr(Callback, f))])
 
     callbacks_func = [
         'on_after_backward',