Repository: Lightning-AI/pytorch-lightning

Issue:
Call load_from_checkpoint when trainer load state
## 🚀 Feature
When use `Trainer` class to load the model state instead of using `load_from_checkpoint`, the hook `on_load_checkpoint` won't be called. It's better to let trainer call `load_from_checkpoint` instead of directly call `model.load_state_dict`, thus `on_load_checkpoint` and other hooks can be called.

### Motivation

I need to do some sanity check of saved `hparams` and I found the suitable location for this is `on_load_checkpoint`. But this method is not called when use Trainer to load states.



Discussion:
to me it makes sense call `on_load_checkpoint` inside `load_from_checkpoint` :]
@jeffling @hadim ano thoughts?
This is a good change, but we may want to think about the versioning here since it may break some codebases. 
@PyTorchLightning/core-contributors any other thoughts?
I just hit this issue, and would love it in the next release!

Resolving patch:
diff --git a/pytorch_lightning/trainer/training_io.py b/pytorch_lightning/trainer/training_io.py
--- a/pytorch_lightning/trainer/training_io.py
+++ b/pytorch_lightning/trainer/training_io.py
@@ -278,6 +278,10 @@ def restore(self, checkpoint_path: str, on_gpu: bool):
 
         # load the state_dict on the model automatically
         model.load_state_dict(checkpoint['state_dict'])
+
+        # give model a chance to load something
+        model.on_load_checkpoint(checkpoint)
+
         if on_gpu:
             model.cuda(self.root_gpu)
 


Test patch:
diff --git a/tests/trainer/test_trainer.py b/tests/trainer/test_trainer.py
--- a/tests/trainer/test_trainer.py
+++ b/tests/trainer/test_trainer.py
@@ -309,8 +309,8 @@ def test_model_freeze_unfreeze():
     model.unfreeze()
 
 
-def test_resume_from_checkpoint_epoch_restored(tmpdir):
-    """Verify resuming from checkpoint runs the right number of epochs"""
+def test_resume_from_checkpoint(tmpdir):
+    """Verify resuming from checkpoint (epoch, batch numbers and on_load_checkpoint())"""
     import types
 
     tutils.reset_seed()
@@ -322,6 +322,7 @@ def _new_model():
         model = LightningTestModel(hparams)
         model.num_epochs_seen = 0
         model.num_batches_seen = 0
+        model.num_on_load_checkpoint_called = 0
 
         def increment_epoch(self):
             self.num_epochs_seen += 1
@@ -329,10 +330,14 @@ def increment_epoch(self):
         def increment_batch(self, _):
             self.num_batches_seen += 1
 
-        # Bind the increment_epoch function on_epoch_end so that the
-        # model keeps track of the number of epochs it has seen.
+        def increment_on_load_checkpoint(self, _):
+            self.num_on_load_checkpoint_called += 1
+
+        # Bind methods to keep track of epoch numbers, batch numbers it has seen
+        # as well as number of times it has called on_load_checkpoint()
         model.on_epoch_end = types.MethodType(increment_epoch, model)
         model.on_batch_start = types.MethodType(increment_batch, model)
+        model.on_load_checkpoint = types.MethodType(increment_on_load_checkpoint, model)
         return model
 
     model = _new_model()
@@ -356,6 +361,7 @@ def increment_batch(self, _):
 
     assert model.num_epochs_seen == 2
     assert model.num_batches_seen == training_batches * 2
+    assert model.num_on_load_checkpoint_called == 0
 
     # Other checkpoints can be uncommented if/when resuming mid-epoch is supported
     checkpoints = sorted(glob.glob(os.path.join(trainer.checkpoint_callback.dirpath, '*.ckpt')))
@@ -369,6 +375,7 @@ def increment_batch(self, _):
         new_trainer = Trainer(**trainer_options, resume_from_checkpoint=check)
         new_trainer.fit(next_model)
         assert state['global_step'] + next_model.num_batches_seen == training_batches * trainer_options['max_epochs']
+        assert next_model.num_on_load_checkpoint_called == 1
 
 
 def _init_steps_model():