Repository: Lightning-AI/pytorch-lightning
Pull request: callback method for on_save_checkpoint
## What does this PR do?

Provides a generic way for callbacks to persist state in checkpoints.

This is still in a draft state for review. If the core team agrees with this approach, we still need to:
- [x] Provide a utility for upgrading old checkpoints to this new standard
- [x] Test utility for upgrading old checkpoints to this new standard
- [x] Provide more documentation

Fixes #2401 #2631
cc @awaelchli 

# Before submitting

- [x] Was this discussed/approved via a Github issue? (no need for typos and docs improvements)
- [x] Did you read the [contributor guideline](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/.github/CONTRIBUTING.md), Pull Request section?
- [x] Did you make sure your PR does only one thing, instead of bundling different changes together? Otherwise, we ask you to create a separate PR for every change.
- [ ] Did you make sure to update the documentation with your changes?
- [x] Did you write any new necessary tests? 
- [x] Did you verify new and existing tests pass locally with your changes?
- [ ] If you made a notable change (that affects users), did you update the [CHANGELOG](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/CHANGELOG.md)?

<!-- For CHANGELOG separate each item in unreleased section by a blank line to reduce collisions -->

## PR review    
Anyone in the community is free to review the PR once the tests have passed.     
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

## Did you have fun?
Make sure you had fun coding 🙃


Problem statement:
Add callback method for on_save_checkpoint
## 🚀 Feature

We should allow Callback objects to optionally persist state that can be reloaded from checkpoints.

### Motivation

We already manually save the state for early stopping and model checkpoint callbacks. This refactor would eliminate callback-specific code in the Trainer and extend the ability to user-written callbacks.

### Pitch

This callback would just return a state_dict which the Trainer could store. The only thing that I am unclear how we should handle is for other callbacks how we want to reinitialize the state. If we can expect that the same exact callbacks will be passed to the Trainer then it should be trivial. Or we could expect that you only pass in a single instance of each callback class (eg. callbacks=[CustomerLogger(), EarlyStopping(), ModelCheckpoint()] and not callbacks=[CustomerLogger(params_a), CustomerLogger(params_b), EarlyStopping(), ModelCheckpoint()] and just keep a mapping of callback class to state dicts. However, if the user passed multiple callback instances of the same class I'm not sure how we would want to handle that.

I would recommend that we document the following constraints:
- All objects in the dictionary must be pickle-able.
- You cannot persist multiple instances of the same callback class.

Hints from issue discussion:
I think it is reasonable to assume there is only one instance of these special callbacks (and should raise error otherwise, e.g. see progress bar callback). Note that currently the logger is not a callback. 
Also, the documentation for callbacks should probably let the user know that the the order of the list input to the Trainer is not preserved (e.g. the Trainer should reorder it so that the earlystopping callback comes after checkpoint, right?)

Cool, I'll put together a draft PR once we merge #2391 :)
also https://github.com/PyTorchLightning/pytorch-lightning/issues/2631

Commit message: clean up and add tests

Patch to review:
diff --git a/docs/source/callbacks.rst b/docs/source/callbacks.rst
index cd816081496cb..55b9653b22534 100644
--- a/docs/source/callbacks.rst
+++ b/docs/source/callbacks.rst
@@ -104,6 +104,18 @@ We successfully extended functionality without polluting our super clean
 
 ----------------
 
+Persiting State
+---------------
+
+If you need to persist your callback's state to be reloaded from a checkpoint,
+we provide hooks to allow for this. However, you must follow two constraints:
+
+1. Your state must be able to be pickled.
+2. If you implement the hooks to persist state for a callback, you can only use one 
+instance of that class in the Trainer callbacks list. We don't support persisting state
+for multiple callbacks of the same class.
+
+
 Best Practices
 --------------
 The following are best practices when using/designing callbacks.
diff --git a/pytorch_lightning/callbacks/base.py b/pytorch_lightning/callbacks/base.py
index f241318fbfe7e..7ea7bc2b48af3 100644
--- a/pytorch_lightning/callbacks/base.py
+++ b/pytorch_lightning/callbacks/base.py
@@ -144,3 +144,12 @@ def on_test_end(self, trainer, pl_module):
 
     def on_keyboard_interrupt(self, trainer, pl_module):
         """Called when the training is interrupted by KeyboardInterrupt."""
+        pass
+
+    def on_save_checkpoint(self, trainer, pl_module):
+        """Called when saving a model checkpoint, use to persist state."""
+        pass
+
+    def on_load_checkpoint(self, checkpointed_state):
+        """Called when loading a model checkpoint, use to reload state."""
+        pass
diff --git a/pytorch_lightning/callbacks/early_stopping.py b/pytorch_lightning/callbacks/early_stopping.py
index 3b4aed8566bb6..ca6fe1437e603 100644
--- a/pytorch_lightning/callbacks/early_stopping.py
+++ b/pytorch_lightning/callbacks/early_stopping.py
@@ -5,8 +5,6 @@
 Monitor a validation metric and stop training when it stops improving.
 
 """
-from copy import deepcopy
-
 import numpy as np
 import torch
 import torch.distributed as dist
@@ -118,7 +116,7 @@ def _validate_condition_metric(self, logs):
     def monitor_op(self):
         return self.mode_dict[self.mode]
 
-    def state_dict(self):
+    def on_save_checkpoint(self, trainer, pl_module):
         return {
             'wait_count': self.wait_count,
             'stopped_epoch': self.stopped_epoch,
@@ -126,12 +124,11 @@ def state_dict(self):
             'patience': self.patience
         }
 
-    def load_state_dict(self, state_dict):
-        state_dict = deepcopy(state_dict)
-        self.wait_count = state_dict['wait_count']
-        self.stopped_epoch = state_dict['stopped_epoch']
-        self.best_score = state_dict['best_score']
-        self.patience = state_dict['patience']
+    def on_load_checkpoint(self, checkpointed_state):
+        self.wait_count = checkpointed_state['wait_count']
+        self.stopped_epoch = checkpointed_state['stopped_epoch']
+        self.best_score = checkpointed_state['best_score']
+        self.patience = checkpointed_state['patience']
 
     def on_validation_end(self, trainer, pl_module):
         self._run_early_stopping_check(trainer, pl_module)
diff --git a/pytorch_lightning/callbacks/model_checkpoint.py b/pytorch_lightning/callbacks/model_checkpoint.py
index 2458007db218d..a4839a25f9ddd 100644
--- a/pytorch_lightning/callbacks/model_checkpoint.py
+++ b/pytorch_lightning/callbacks/model_checkpoint.py
@@ -375,3 +375,13 @@ def _do_check_save(self, filepath, current, epoch, trainer, pl_module):
         for cur_path in del_list:
             if cur_path != filepath:
                 self._del_model(cur_path)
+
+    def on_save_checkpoint(self, trainer, pl_module):
+        return {
+            'best_model_score': self.best_model_score,
+            'best_model_path': self.best_model_path,
+        }
+
+    def on_load_checkpoint(self, checkpointed_state):
+        self.best_model_score = checkpointed_state['best_model_score']
+        self.best_model_path = checkpointed_state['best_model_path']
\ No newline at end of file
diff --git a/pytorch_lightning/trainer/callback_hook.py b/pytorch_lightning/trainer/callback_hook.py
index b10c86af451b0..b74d56b30e2f3 100644
--- a/pytorch_lightning/trainer/callback_hook.py
+++ b/pytorch_lightning/trainer/callback_hook.py
@@ -1,4 +1,5 @@
 from abc import ABC
+from copy import deepcopy
 from typing import Callable, List
 
 from pytorch_lightning.callbacks import Callback
@@ -175,3 +176,22 @@ def on_keyboard_interrupt(self):
         """Called when the training is interrupted by KeyboardInterrupt."""
         for callback in self.callbacks:
             callback.on_keyboard_interrupt(self, self.get_model())
+
+    def on_save_checkpoint(self):
+        """Called when saving a model checkpoint."""
+        callback_states = {}
+        for callback in self.callbacks:
+            callback_class = type(callback)
+            state = callback.on_save_checkpoint(self, self.get_model())
+            if state:
+                callback_states[callback_class] = state
+        return callback_states
+
+    def on_load_checkpoint(self, checkpoint):
+        """Called when loading a model checkpoint."""
+        callback_states = checkpoint.get('callbacks')
+        for callback in self.callbacks:
+            state = callback_states.get(type(callback))
+            if state:
+                state = deepcopy(state)
+                callback.on_load_checkpoint(state)
diff --git a/pytorch_lightning/trainer/training_io.py b/pytorch_lightning/trainer/training_io.py
index a4137dc3879c4..a89b8fbecff26 100644
--- a/pytorch_lightning/trainer/training_io.py
+++ b/pytorch_lightning/trainer/training_io.py
@@ -347,20 +347,9 @@ def dump_checkpoint(self, weights_only: bool = False) -> dict:
 
         if not weights_only:
 
-            # TODO support more generic way for callbacks to persist a state_dict in a checkpoint
-            checkpoint_callbacks = [c for c in self.callbacks if isinstance(c, ModelCheckpoint)]
-            early_stopping_callbacks = [c for c in self.callbacks if isinstance(c, EarlyStopping)]
-
-            if checkpoint_callbacks:
-                # we add the official checkpoint callback to the end of the list
-                # extra user provided callbacks will not be persisted yet
-                checkpoint[ModelCheckpoint.CHECKPOINT_STATE_BEST_SCORE] = self.checkpoint_callback.best_model_score
-                checkpoint[ModelCheckpoint.CHECKPOINT_STATE_BEST_PATH] = self.checkpoint_callback.best_model_path
-
-            if early_stopping_callbacks and checkpoint_callbacks:
-                # we add the official early stopping callback to the end of the list
-                # extra user provided callbacks will not be persisted yet
-                checkpoint['early_stop_callback_state_dict'] = early_stopping_callbacks[-1].state_dict()
+            # save callbacks
+            callback_states = self.on_save_checkpoint()
+            checkpoint['callbacks'] = callback_states
 
             # save optimizers
             optimizer_states = []
@@ -431,25 +420,23 @@ def restore_training_state(self, checkpoint):
                 ' This is probably due to `ModelCheckpoint.save_weights_only` being set to `True`.'
             )
 
-        # TODO support more generic way for callbacks to load callback state_dicts
-        checkpoint_callbacks = [c for c in self.callbacks if isinstance(c, ModelCheckpoint)]
-        early_stopping_callbacks = [c for c in self.callbacks if isinstance(c, EarlyStopping)]
-
-        if checkpoint_callbacks:
-            if ModelCheckpoint.CHECKPOINT_STATE_BEST_SCORE in checkpoint:
-                checkpoint_callbacks[-1].best_model_score = checkpoint[ModelCheckpoint.CHECKPOINT_STATE_BEST_SCORE]
-            else:
-                # Old naming until version 0.7.6
-                rank_zero_warn(
-                    'Loading a checkpoint created with an old version of Lightning; '
-                    'this will not be supported in the future.'
-                )
-                checkpoint_callbacks[-1].best_model_score = checkpoint['checkpoint_callback_best']
-            checkpoint_callbacks[-1].best_model_path = checkpoint[ModelCheckpoint.CHECKPOINT_STATE_BEST_PATH]
+        deprecated_keys = (
+            'checkpoint_callback_best_model_score',
+            'checkpoint_callback_best_model_path',
+            'checkpoint_callback_best',
+            'early_stop_callback_wait',
+            'early_stop_callback_patience'
+        )
+        if any([key in checkpoint for key in deprecated_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."
+            )
 
-        if early_stopping_callbacks:
-            state = checkpoint['early_stop_callback_state_dict']
-            early_stopping_callbacks[-1].load_state_dict(state)
+        # load callback states
+        self.on_load_checkpoint(checkpoint)
 
         self.global_step = checkpoint['global_step']
         self.current_epoch = checkpoint['epoch']
diff --git a/pytorch_lightning/utilities/upgrade_checkpoint.py b/pytorch_lightning/utilities/upgrade_checkpoint.py
new file mode 100644
index 0000000000000..a1b224d51b36c
--- /dev/null
+++ b/pytorch_lightning/utilities/upgrade_checkpoint.py
@@ -0,0 +1,39 @@
+import argparse
+from shutil import copyfile
+
+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"),
+}
+
+
+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(description="Upgrade an old checkpoint to the current schema.")
+    parser.add_argument("--file", help="filepath for a checkpoint to upgrade")
+    args = parser.parse_args()
+
+    copyfile(args.file, args.file + ".bak")
+    upgrade_checkpoint(args.file)
diff --git a/tests/callbacks/test_callbacks.py b/tests/callbacks/test_callbacks.py
index 90a3d9117c376..9d5b638ee75cc 100644
--- a/tests/callbacks/test_callbacks.py
+++ b/tests/callbacks/test_callbacks.py
@@ -1,3 +1,10 @@
+from pathlib import Path
+import os
+
+import pytest
+import torch
+
+import tests.base.develop_utils as tutils
 from pytorch_lightning import Callback
 from pytorch_lightning import Trainer, LightningModule
 from tests.base import EvalModelTemplate
diff --git a/tests/callbacks/test_early_stopping.py b/tests/callbacks/test_early_stopping.py
index c9b6850a1ad28..2615865e2541e 100644
--- a/tests/callbacks/test_early_stopping.py
+++ b/tests/callbacks/test_early_stopping.py
@@ -10,6 +10,23 @@
 from tests.base import EvalModelTemplate
 
 
+class EarlyStoppingTestRestore(EarlyStopping):
+    # this class has to be defined outside the test function, otherwise we get pickle error
+    def __init__(self, expected_state=None):
+        super().__init__()
+        self.expected_state = expected_state
+        # cache the state for each epoch
+        self.saved_states = []
+
+    def on_train_start(self, trainer, pl_module):
+        if self.expected_state:
+            assert self.on_save_checkpoint(trainer, pl_module) == self.expected_state
+
+    def on_validation_end(self, trainer, pl_module):
+        super().on_validation_end(trainer, pl_module)
+        self.saved_states.append(self.on_save_checkpoint(trainer, pl_module).copy())
+
+
 def test_resume_early_stopping_from_checkpoint(tmpdir):
     """
     Prevent regressions to bugs:
@@ -17,27 +34,9 @@ def test_resume_early_stopping_from_checkpoint(tmpdir):
     https://github.com/PyTorchLightning/pytorch-lightning/issues/1463
     """
 
-    class EarlyStoppingTestStore(EarlyStopping):
-        def __init__(self, *args, **kwargs):
-            super().__init__(*args, **kwargs)
-            # cache the state for each epoch
-            self.saved_states = []
-
-        def on_validation_end(self, trainer, pl_module):
-            super().on_validation_end(trainer, pl_module)
-            self.saved_states.append(deepcopy(self.state_dict()))
-
-    class EarlyStoppingTestRestore(EarlyStopping):
-        def __init__(self, expected_state):
-            super().__init__()
-            self.expected_state = expected_state
-
-        def on_train_start(self, trainer, pl_module):
-            assert self.state_dict() == self.expected_state
-
     model = EvalModelTemplate()
     checkpoint_callback = ModelCheckpoint(save_top_k=1)
-    early_stop_callback = EarlyStoppingTestStore()
+    early_stop_callback = EarlyStoppingTestRestore()
     trainer = Trainer(
         default_root_dir=tmpdir,
         checkpoint_callback=checkpoint_callback,
@@ -50,9 +49,9 @@ def on_train_start(self, trainer, pl_module):
     # ensure state is persisted properly
     checkpoint = torch.load(checkpoint_filepath)
     # the checkpoint saves "epoch + 1"
-    early_stop_callback_state = early_stop_callback.saved_states[checkpoint['epoch'] - 1]
+    early_stop_callback_state = early_stop_callback.saved_states[checkpoint["epoch"] - 1]
     assert 4 == len(early_stop_callback.saved_states)
-    assert checkpoint['early_stop_callback_state_dict'] == early_stop_callback_state
+    assert checkpoint["callbacks"][type(early_stop_callback)] == early_stop_callback_state
 
     # ensure state is reloaded properly (assertion in the callback)
     early_stop_callback = EarlyStoppingTestRestore(early_stop_callback_state)
@@ -65,19 +64,22 @@ def on_train_start(self, trainer, pl_module):
     new_trainer.fit(model)
 
 
-def test_early_stopping_no_extraneous_invocations(tmpdir):
-    """Test to ensure that callback methods aren't being invoked outside of the callback handler."""
-    class EarlyStoppingTestInvocations(EarlyStopping):
-        def __init__(self, expected_count):
-            super().__init__()
-            self.count = 0
-            self.expected_count = expected_count
+class EarlyStoppingTestInvocations(EarlyStopping):
+    # this class has to be defined outside the test function, otherwise we get pickle error
+    def __init__(self, expected_count):
+        super().__init__()
+        self.count = 0
+        self.expected_count = expected_count
 
-        def on_validation_end(self, trainer, pl_module):
-            self.count += 1
+    def on_validation_end(self, trainer, pl_module):
+        self.count += 1
+
+    def on_train_end(self, trainer, pl_module):
+        assert self.count == self.expected_count
 
-        def on_train_end(self, trainer, pl_module):
-            assert self.count == self.expected_count
+
+def test_early_stopping_no_extraneous_invocations(tmpdir):
+    """Test to ensure that callback methods aren't being invoked outside of the callback handler."""
 
     model = EvalModelTemplate()
     expected_count = 4
@@ -91,11 +93,10 @@ def on_train_end(self, trainer, pl_module):
     trainer.fit(model)
 
 
-@pytest.mark.parametrize('loss_values, patience, expected_stop_epoch', [
-    ([6, 5, 5, 5, 5, 5], 3, 4),
-    ([6, 5, 4, 4, 3, 3], 1, 3),
-    ([6, 5, 6, 5, 5, 5], 3, 4),
-])
+@pytest.mark.parametrize(
+    "loss_values, patience, expected_stop_epoch",
+    [([6, 5, 5, 5, 5, 5], 3, 4), ([6, 5, 4, 4, 3, 3], 1, 3), ([6, 5, 6, 5, 5, 5], 3, 4),],
+)
 def test_early_stopping_patience(tmpdir, loss_values, patience, expected_stop_epoch):
     """Test to ensure that early stopping is not triggered before patience is exhausted."""
 
diff --git a/tests/callbacks/test_model_checkpoint.py b/tests/callbacks/test_model_checkpoint.py
index dfcc1e036854d..b97b6bb6c947c 100644
--- a/tests/callbacks/test_model_checkpoint.py
+++ b/tests/callbacks/test_model_checkpoint.py
@@ -14,7 +14,7 @@
 from tests.base import EvalModelTemplate
 
 
-@pytest.mark.parametrize('save_top_k', [-1, 0, 1, 2])
+@pytest.mark.parametrize("save_top_k", [-1, 0, 1, 2])
 def test_model_checkpoint_with_non_string_input(tmpdir, save_top_k):
     """ Test that None in checkpoint callback is valid and that chkp_path is set correctly """
     tutils.reset_seed()
@@ -24,11 +24,11 @@ def test_model_checkpoint_with_non_string_input(tmpdir, save_top_k):
 
     trainer = Trainer(default_root_dir=tmpdir, checkpoint_callback=checkpoint, overfit_batches=0.20, max_epochs=2)
     trainer.fit(model)
-    assert checkpoint.dirpath == tmpdir / trainer.logger.name / 'version_0' / 'checkpoints'
+    assert checkpoint.dirpath == tmpdir / trainer.logger.name / "version_0" / "checkpoints"
 
 
 @pytest.mark.parametrize(
-    'logger_version,expected', [(None, 'version_0'), (1, 'version_1'), ('awesome', 'awesome')],
+    "logger_version,expected", [(None, "version_0"), (1, "version_1"), ("awesome", "awesome")],
 )
 def test_model_checkpoint_path(tmpdir, logger_version, expected):
     """Test that "version_" prefix is only added when logger's version is an integer"""
@@ -85,7 +85,7 @@ def test_model_checkpoint_no_extraneous_invocations(tmpdir):
     num_epochs = 4
     model_checkpoint = ModelCheckpointTestInvocations(expected_count=num_epochs, save_top_k=-1)
     trainer = Trainer(
-        distributed_backend='ddp_cpu',
+        distributed_backend="ddp_cpu",
         num_processes=2,
         default_root_dir=tmpdir,
         early_stop_callback=False,
@@ -103,10 +103,7 @@ def test_model_checkpoint_save_last_checkpoint_contents(tmpdir):
     num_epochs = 3
     model_checkpoint = ModelCheckpoint(filepath=tmpdir, save_top_k=num_epochs, save_last=True)
     trainer = Trainer(
-        default_root_dir=tmpdir,
-        early_stop_callback=False,
-        checkpoint_callback=model_checkpoint,
-        max_epochs=num_epochs,
+        default_root_dir=tmpdir, early_stop_callback=False, checkpoint_callback=model_checkpoint, max_epochs=num_epochs,
     )
     trainer.fit(model)
     path_last_epoch = model_checkpoint.format_checkpoint_name(num_epochs - 1, {})  # epoch=3.ckpt
@@ -114,15 +111,24 @@ def test_model_checkpoint_save_last_checkpoint_contents(tmpdir):
     assert path_last_epoch != path_last
     ckpt_last_epoch = torch.load(path_last_epoch)
     ckpt_last = torch.load(path_last)
-    matching_keys = (
+
+    trainer_keys = (
         "epoch",
         "global_step",
-        ModelCheckpoint.CHECKPOINT_STATE_BEST_SCORE,
-        ModelCheckpoint.CHECKPOINT_STATE_BEST_PATH,
     )
-    for key in matching_keys:
+    for key in trainer_keys:
         assert ckpt_last_epoch[key] == ckpt_last[key]
 
+    checkpoint_callback_keys = (
+        "best_model_score",
+        "best_model_path",
+    )
+    for key in checkpoint_callback_keys:
+        assert (
+            ckpt_last_epoch["callbacks"][type(model_checkpoint)][key]
+            == ckpt_last_epoch["callbacks"][type(model_checkpoint)][key]
+        )
+
     # it is easier to load the model objects than to iterate over the raw dict of tensors
     model_last_epoch = EvalModelTemplate.load_from_checkpoint(path_last_epoch)
     model_last = EvalModelTemplate.load_from_checkpoint(path_last)
diff --git a/tests/utilities/test_upgrade_checkpoint.py b/tests/utilities/test_upgrade_checkpoint.py
new file mode 100644
index 0000000000000..5f5ecd6b5f2df
--- /dev/null
+++ b/tests/utilities/test_upgrade_checkpoint.py
@@ -0,0 +1,36 @@
+import pytest
+import os
+
+import torch
+
+from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
+from pytorch_lightning.utilities.upgrade_checkpoint import upgrade_checkpoint
+
+
+@pytest.mark.parametrize(
+    "old_checkpoint, new_checkpoint",
+    [
+        (
+            {"epoch": 1, "global_step": 23, "checkpoint_callback_best": 0.34},
+            {"epoch": 1, "global_step": 23, "callbacks": {ModelCheckpoint: {"best_model_score": 0.34}}},
+        ),
+        (
+            {"epoch": 1, "global_step": 23, "checkpoint_callback_best_model_score": 0.99},
+            {"epoch": 1, "global_step": 23, "callbacks": {ModelCheckpoint: {"best_model_score": 0.99}}},
+        ),
+        (
+            {"epoch": 1, "global_step": 23, "checkpoint_callback_best_model_path": 'path'},
+            {"epoch": 1, "global_step": 23, "callbacks": {ModelCheckpoint: {"best_model_path": 'path'}}},
+        ),
+        (
+            {"epoch": 1, "global_step": 23, "early_stop_callback_wait": 2, "early_stop_callback_patience": 4},
+            {"epoch": 1, "global_step": 23, "callbacks": {EarlyStopping: {"wait_count": 2, "patience": 4}}},
+        ),
+    ],
+)
+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)
+    assert updated_checkpoint == new_checkpoint


Code review comments:
diff_hunk: @@ -0,0 +1,39 @@
+import argparse
+from shutil import copyfile
+
+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"),
+}
+
+
+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(description="Upgrade an old checkpoint to the current schema.")
+    parser.add_argument("--file", help="filepath for a checkpoint to upgrade")
+    args = parser.parse_args()
+
+    copyfile(args.file, args.file + ".bak")
path: pytorch_lightning/utilities/upgrade_checkpoint.py
text: @user1:
maybe we can write to logging.info to inform the user, and also add to the argparse description that it will create a backup file?

diff_hunk: @@ -431,25 +420,23 @@ def restore_training_state(self, checkpoint):
                 ' This is probably due to `ModelCheckpoint.save_weights_only` being set to `True`.'
             )
 
-        # TODO support more generic way for callbacks to load callback state_dicts
-        checkpoint_callbacks = [c for c in self.callbacks if isinstance(c, ModelCheckpoint)]
-        early_stopping_callbacks = [c for c in self.callbacks if isinstance(c, EarlyStopping)]
-
-        if checkpoint_callbacks:
-            if ModelCheckpoint.CHECKPOINT_STATE_BEST_SCORE in checkpoint:
-                checkpoint_callbacks[-1].best_model_score = checkpoint[ModelCheckpoint.CHECKPOINT_STATE_BEST_SCORE]
-            else:
-                # Old naming until version 0.7.6
-                rank_zero_warn(
-                    'Loading a checkpoint created with an old version of Lightning; '
-                    'this will not be supported in the future.'
-                )
-                checkpoint_callbacks[-1].best_model_score = checkpoint['checkpoint_callback_best']
-            checkpoint_callbacks[-1].best_model_path = checkpoint[ModelCheckpoint.CHECKPOINT_STATE_BEST_PATH]
+        deprecated_keys = (
+            'checkpoint_callback_best_model_score',
+            'checkpoint_callback_best_model_path',
+            'checkpoint_callback_best',
+            'early_stop_callback_wait',
+            'early_stop_callback_patience'
+        )
+        if any([key in checkpoint for key in deprecated_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."
path: pytorch_lightning/trainer/training_io.py
text: @user1:
```suggestion
                "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."
```