Repository: pytorch/ignite

Issue:
Checkpoint stateful handlers and metrics
## 🚀 Feature

<!-- A clear and concise description of the feature proposal -->
Things that are attached to the `Engine` might have state that would ideally be checkpointed and restored using as part of the `Engine`'s `state_dict`.

<!-- Please outline the motivation for the proposal. Is your feature request related to a problem? e.g., I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too -->

An example is a `Checkpoint` handler when it has a `score_function`. Currently, the priorities the `Checkpoint` class stores is not saved anywhere. It is not able to recover gracefully from failure without manual intervention to parse the checkpoint path names, and directly setting the internals of the class.

<!-- A clear and concise description of what you want to happen. -->

Handlers and Metrics should have `state_dict` and `load_state_dict` methods (empty by default), and I think it should be possible for these to automatically make it into/restored from the `Engine`'s `state_dict` when it's attached to an `Engine`.


Discussion:
vfdev-5: @amatsukawa thanks for this PR too !

I would change "Fixes #966" with related to #966 I think it wont completely fix it, right ?

I agree with updating `Checkpoint` and `EarlyStopping` to be derived from `Serializable`. However, automatic inclusion of `self` with `include_itself` should be carefully tested and I think it can be too short to have this in v0.4.0...
What do you think ?


amatsukawa: I changed the description. This does support the most of the functionality I had in mind, except for the "do it automatically" part (which as we discussed elsewhere, may not be realistic after all).

I did add a test to make sure `include_self` at least seemed to work as intended. I followed a fairly simple rule that seems to me like would have a high likelihood of correctness, which is to add `self.state_dict` *right before* the call to the save handler (which is the last thing in the method). 

I'm happy to remove `include_self` from the PR for now. Its effect would be that `Checkpoint` object cannot respect the `n_saved` parameter *across resumes*.
vfdev-5: @amatsukawa I need to play a bit with the code. Maybe, a complete integration test would be nice to ensure that saving and resume (`Checkpoint.load_objects`) work as intended...
Let's take a bit of time to check all this (during 24h) and than we take a decision either keep or remove as we plan to cut the release on June 25th. Does it sound good for you ?
amatsukawa: Is there an existing integration test for things that this touches, without checkpointing? If so, I can try to just add to it or create a copy of it for checkpointing.
vfdev-5: Maybe this one is the most integrating among the others : https://github.com/pytorch/ignite/blob/master/tests/ignite/handlers/test_checkpoint.py#L713

amatsukawa: @vfdev-5 I made a version of that test you linked with new features, PTAL.
amatsukawa: Oops, @vfdev-5 I'm sorry I didn't realized you had pushed some changes to this branch. I thought I had done something wrong in merging master to cause the conflict, so I force pushed.

Could you push your changes again?
vfdev-5: Yep, sorry, I have this practice (push into existing PRs), it can be surprising :)
I'll try to merge your newer commits and mine...
vfdev-5: I think if we are done with it, we can merge it like that. No need to update the docs as API is not that brilliant.
Ideally, it would be great to have a possibilty to add somehow those stateful components like `early_stop` and `best_checkpoint` to the trainer such that when we save the trainer, we save also appended components...
Something like that 
```python
trainer = Engine(lambda e, b: None)

to_save = {"trainer": trainer}
checkpoint = ModelCheckpoint(...)

evaluator = Engine(lambda e, b: None)
early_stopping = EarlyStopping(...)
evaluator.add_event_handler(Events.COMPLETED, early_stopping)

trainer.state.early_stopping = early_stopping
trainer.state_dict_user_keys.append("early_stopping")  # this wont work but an API like that maybe more interesting vs manual addition.
```
amatsukawa: Thanks for the review!

I would vote for merging this in to support this functionality manually for 0.4, and improving usability from there. 

As future work, perhaps hook states can appear in engine states, and both engines can be checkpointed. Alternatively, #912 happens, it may be fine to hide it inside a Trainer how it is now.

Patch:
diff --git a/ignite/handlers/checkpoint.py b/ignite/handlers/checkpoint.py
index dcb2a50f..86975414 100644
--- a/ignite/handlers/checkpoint.py
+++ b/ignite/handlers/checkpoint.py
@@ -4,13 +4,14 @@ import os
 import tempfile
 import warnings
 from abc import ABCMeta, abstractmethod
-from collections import namedtuple
+from collections import OrderedDict, namedtuple
 from typing import Callable, Mapping, Optional, Union
 
 import torch
 import torch.nn as nn
 
 import ignite.distributed as idist
+from ignite.base import Serializable
 from ignite.engine import Engine, Events
 
 __all__ = ["Checkpoint", "DiskSaver", "ModelCheckpoint", "BaseSaveHandler"]
@@ -62,7 +63,7 @@ class BaseSaveHandler(metaclass=ABCMeta):
         pass
 
 
-class Checkpoint:
+class Checkpoint(Serializable):
     """Checkpoint handler can be used to periodically save and load objects which have attribute
     ``state_dict`/`load_state_dict``. This class can use specific save handlers to store on the disk or a cloud
     storage, etc. The Checkpoint handler (if used with :class:`~ignite.handlers.DiskSaver`) also handles automatically
@@ -92,6 +93,8 @@ class Checkpoint:
             Default is None, global_step based on attached engine. If provided, uses function output as global_step.
             To setup global step from another engine, please use :meth:`~ignite.handlers.global_step_from_engine`.
         archived (bool, optional): Deprecated argument as models saved by ``torch.save`` are already compressed.
+        include_self (bool): Whether to include the `state_dict` of this object in the checkpoint. If `True`, then
+            there must not be another object in ``to_save`` with key ``checkpointer``.
 
     .. _DistributedDataParallel: https://pytorch.org/docs/stable/nn.html#torch.nn.parallel.DistributedDataParallel
     .. _DataParallel: https://pytorch.org/docs/stable/nn.html#torch.nn.DataParallel
@@ -204,6 +207,7 @@ class Checkpoint:
     """
 
     Item = namedtuple("Item", ["priority", "filename"])
+    _state_dict_all_req_keys = ("saved",)
 
     def __init__(
         self,
@@ -215,6 +219,7 @@ class Checkpoint:
         n_saved: Optional[int] = 1,
         global_step_transform: Callable = None,
         archived: bool = False,
+        include_self: bool = False,
     ):
 
         if to_save is not None:  # for compatibility with ModelCheckpoint
@@ -226,6 +231,15 @@ class Checkpoint:
 
             self._check_objects(to_save, "state_dict")
 
+            if include_self:
+                if not isinstance(to_save, collections.MutableMapping):
+                    raise TypeError(
+                        "If `include_self` is True, then `to_save` must be mutable, but given {}.".format(type(to_save))
+                    )
+
+                if "checkpointer" in to_save:
+                    raise ValueError("Cannot have key 'checkpointer' if `include_self` is True: {}".format(to_save))
+
         if not (callable(save_handler) or isinstance(save_handler, BaseSaveHandler)):
             raise TypeError("Argument `save_handler` should be callable or inherit from BaseSaveHandler")
 
@@ -248,6 +262,7 @@ class Checkpoint:
         self._saved = []
         self._ext = ".pt"
         self.global_step_transform = global_step_transform
+        self.include_self = include_self
 
     @property
     def last_checkpoint(self) -> Optional[str]:
@@ -321,6 +336,10 @@ class Checkpoint:
             self._saved.append(Checkpoint.Item(priority, filename))
             self._saved.sort(key=lambda item: item[0])
 
+            if self.include_self:
+                # Now that we've updated _saved, we can add our own state_dict.
+                checkpoint["checkpointer"] = self.state_dict()
+
             try:
                 self.save_handler(checkpoint, filename, metadata)
             except TypeError:
@@ -408,6 +427,13 @@ class Checkpoint:
             else:
                 obj.load_state_dict(checkpoint[k])
 
+    def state_dict(self) -> OrderedDict:
+        return OrderedDict([("saved", [(p, f) for p, f in self._saved])])
+
+    def load_state_dict(self, state_dict: Mapping) -> None:
+        super().load_state_dict(state_dict)
+        self._saved = [Checkpoint.Item(p, f) for p, f in state_dict["saved"]]
+
 
 class DiskSaver(BaseSaveHandler):
     """Handler that saves input checkpoint on a disk.
@@ -545,6 +571,8 @@ class ModelCheckpoint(Checkpoint):
             Default is None, global_step based on attached engine. If provided, uses function output as global_step.
             To setup global step from another engine, please use :meth:`~ignite.handlers.global_step_from_engine`.
         archived (bool, optional): Deprecated argument as models saved by `torch.save` are already compressed.
+        include_self (bool): Whether to include the `state_dict` of this object in the checkpoint. If `True`, then
+            there must not be another object in ``to_save`` with key ``checkpointer``.
 
     Examples:
         >>> import os
@@ -576,6 +604,7 @@ class ModelCheckpoint(Checkpoint):
         save_as_state_dict: bool = True,
         global_step_transform: Optional[Callable] = None,
         archived: bool = False,
+        include_self: bool = False,
     ):
 
         if not save_as_state_dict:
@@ -606,6 +635,7 @@ class ModelCheckpoint(Checkpoint):
             n_saved=n_saved,
             global_step_transform=global_step_transform,
             archived=archived,
+            include_self=include_self,
         )
 
     @property
diff --git a/ignite/handlers/early_stopping.py b/ignite/handlers/early_stopping.py
index 8c7efa20..b1000b83 100644
--- a/ignite/handlers/early_stopping.py
+++ b/ignite/handlers/early_stopping.py
@@ -1,12 +1,14 @@
 import logging
-from typing import Callable
+from collections import OrderedDict
+from typing import Callable, Mapping
 
+from ignite.base import Serializable
 from ignite.engine import Engine
 
 __all__ = ["EarlyStopping"]
 
 
-class EarlyStopping:
+class EarlyStopping(Serializable):
     """EarlyStopping handler can be used to stop the training if no improvement after a given number of events.
 
     Args:
@@ -41,6 +43,11 @@ class EarlyStopping:
 
     """
 
+    _state_dict_all_req_keys = (
+        "counter",
+        "best_score",
+    )
+
     def __init__(
         self,
         patience: int,
@@ -87,3 +94,11 @@ class EarlyStopping:
         else:
             self.best_score = score
             self.counter = 0
+
+    def state_dict(self) -> OrderedDict:
+        return OrderedDict([("counter", self.counter), ("best_score", self.best_score)])
+
+    def load_state_dict(self, state_dict: Mapping) -> None:
+        super().load_state_dict(state_dict)
+        self.counter = state_dict["counter"]
+        self.best_score = state_dict["best_score"]


Test patch:
diff --git a/tests/ignite/handlers/test_checkpoint.py b/tests/ignite/handlers/test_checkpoint.py
index f69c82f5..9d6a779b 100644
--- a/tests/ignite/handlers/test_checkpoint.py
+++ b/tests/ignite/handlers/test_checkpoint.py
@@ -1,5 +1,6 @@
 import os
 import warnings
+from collections import OrderedDict
 from unittest.mock import MagicMock
 
 import pytest
@@ -8,7 +9,7 @@ import torch.nn as nn
 
 import ignite.distributed as idist
 from ignite.engine import Engine, Events, State
-from ignite.handlers import Checkpoint, DiskSaver, ModelCheckpoint
+from ignite.handlers import Checkpoint, DiskSaver, EarlyStopping, ModelCheckpoint, global_step_from_engine
 from ignite.handlers.checkpoint import BaseSaveHandler
 
 _PREFIX = "PREFIX"
@@ -63,6 +64,9 @@ def test_checkpoint_wrong_input():
     with pytest.warns(UserWarning, match=r"Argument archived is deprecated"):
         Checkpoint(to_save, lambda x: x, score_function=lambda e: 123, score_name="acc", archived=True)
 
+    with pytest.raises(ValueError, match=r"Cannot have key 'checkpointer' if `include_self` is True"):
+        Checkpoint({"checkpointer": model}, lambda x: x, include_self=True)
+
 
 def test_checkpoint_score_function_wrong_output():
     model = DummyModel()
@@ -111,6 +115,56 @@ def test_checkpoint_default():
     _test(to_save, {"model": model.state_dict(), "optimizer": optimizer.state_dict()}, "checkpoint")
 
 
+def test_checkpoint_include_self_state_dict():
+    def _test(to_save, obj, name):
+        save_handler = MagicMock(spec=BaseSaveHandler)
+
+        checkpointer = Checkpoint(to_save, save_handler=save_handler, include_self=True)
+        assert checkpointer.last_checkpoint is None
+
+        trainer = Engine(lambda e, b: None)
+        trainer.state = State(epoch=0, iteration=0)
+
+        checkpointer(trainer)
+        assert save_handler.call_count == 1
+
+        fname = "{}_0.pt".format(name)
+        obj["checkpointer"] = OrderedDict([("saved", [(0, fname)])])
+
+        metadata = {"basename": name, "score_name": None, "priority": 0}
+        save_handler.assert_called_with(obj, fname, metadata)
+
+        # Swap object, state should be maintained
+        checkpointer2 = Checkpoint(to_save, save_handler=save_handler, include_self=True)
+        checkpointer2.load_state_dict(checkpointer.state_dict())
+        assert checkpointer2.last_checkpoint == fname
+
+        trainer.state.epoch = 12
+        trainer.state.iteration = 1234
+        checkpointer2(trainer)
+        assert save_handler.call_count == 2
+        metadata["priority"] = 1234
+
+        # This delete only happens if state was restored correctly.
+        save_handler.remove.assert_called_with("{}_0.pt".format(name))
+
+        fname = "{}_1234.pt".format(name)
+        obj["checkpointer"] = OrderedDict([("saved", [(1234, fname)])])
+
+        save_handler.assert_called_with(obj, fname, metadata)
+        assert save_handler.remove.call_count == 1
+        assert checkpointer2.last_checkpoint == fname
+
+    model = DummyModel()
+    to_save = {"model": model}
+    _test(to_save, model.state_dict(), "model")
+
+    model = DummyModel()
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
+    to_save = {"model": model, "optimizer": optimizer}
+    _test(to_save, {"model": model.state_dict(), "optimizer": optimizer.state_dict()}, "checkpoint")
+
+
 def test_checkpoint_with_dp():
 
     model = DummyModel()
@@ -792,6 +846,138 @@ def test_save_model_optimizer_lr_scheduler_with_state_dict(dirname):
     _test_save_model_optimizer_lr_scheduler_with_state_dict("cpu", dirname)
 
 
+def _test_save_model_optimizer_lr_scheduler_with_validation(device, dirname, on_zero_rank=False):
+    torch.manual_seed(23)
+
+    def _build_objects(acc_list):
+
+        model = DummyModel().to(device)
+        optim = torch.optim.SGD(model.parameters(), lr=0.1)
+        lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optim, gamma=0.5)
+
+        def update_fn(engine, batch):
+            x = torch.rand((4, 1)).to(device)
+            optim.zero_grad()
+            y = model(x)
+            loss = y.pow(2.0).sum()
+            loss.backward()
+            if idist.has_xla_support:
+                import torch_xla.core.xla_model as xm
+
+                xm.optimizer_step(optim, barrier=True)
+            else:
+                optim.step()
+            lr_scheduler.step()
+
+        trainer = Engine(update_fn)
+
+        evaluator = Engine(lambda e, b: None)
+        acc_iter = iter(acc_list)
+
+        @evaluator.on(Events.EPOCH_COMPLETED)
+        def setup_result():
+            evaluator.state.metrics["accuracy"] = next(acc_iter)
+
+        @trainer.on(Events.EPOCH_COMPLETED)
+        def run_eval():
+            evaluator.run([0, 1, 2])
+
+        def score_function(engine):
+            return engine.state.metrics["accuracy"]
+
+        save_handler = DiskSaver(dirname, create_dir=True, require_empty=False)
+        early_stop = EarlyStopping(score_function=score_function, patience=2, trainer=trainer)
+        evaluator.add_event_handler(Events.COMPLETED, early_stop)
+
+        checkpointer = Checkpoint(
+            {
+                "trainer": trainer,
+                "model": model,
+                "optim": optim,
+                "lr_scheduler": lr_scheduler,
+                "early_stop": early_stop,
+            },
+            save_handler,
+            include_self=True,
+            global_step_transform=global_step_from_engine(trainer),
+        )
+        evaluator.add_event_handler(Events.COMPLETED, checkpointer)
+
+        return trainer, evaluator, model, optim, lr_scheduler, early_stop, checkpointer
+
+    trainer, evaluator, model, optim, scheduler, early, checkpointer = _build_objects([0.2, 0.3, 0.2])
+    trainer.run([0], max_epochs=3)
+
+    saved_objects = sorted(os.listdir(dirname))
+    saved_checkpoint = os.path.join(dirname, saved_objects[0])
+
+    loaded_obj = torch.load(saved_checkpoint, map_location=device)
+    for f in ["trainer", "model", "optim", "lr_scheduler", "early_stop", "checkpointer"]:
+        assert f in loaded_obj
+
+    trainer2, evaluator2, model2, optim2, scheduler2, early2, checkpointer2 = _build_objects([0.1, 0.1, 0.1])
+    Checkpoint.load_objects(
+        {
+            "trainer": trainer2,
+            "model": model2,
+            "optim": optim2,
+            "lr_scheduler": scheduler2,
+            "early_stop": early2,
+            "checkpointer": checkpointer2,
+        },
+        loaded_obj,
+    )
+    assert checkpointer2.last_checkpoint == checkpointer.last_checkpoint
+
+    model_state_dict = model.cpu().state_dict()
+    loaded_model_state_dict = model2.cpu().state_dict()
+    for key in model_state_dict.keys():
+        assert key in loaded_model_state_dict
+        model_value = model_state_dict[key]
+        loaded_model_value = loaded_model_state_dict[key]
+        assert model_value.cpu().numpy() == loaded_model_value.cpu().numpy()
+
+    optim_state_dict = optim.state_dict()
+    loaded_optimizer_state_dict = optim2.state_dict()
+    # "params" contains tensor IDs, which are different
+    del optim_state_dict["param_groups"][0]["params"]
+    del loaded_optimizer_state_dict["param_groups"][0]["params"]
+    for key in optim_state_dict.keys():
+        assert key in loaded_optimizer_state_dict
+        optim_value = optim_state_dict[key]
+        loaded_optim_value = loaded_optimizer_state_dict[key]
+        if idist.get_rank() == 0:
+            assert optim_value == loaded_optim_value
+
+    def _check_state_dict(original, loaded):
+        original_state_dict = original.state_dict()
+        loaded_state_dict = loaded.state_dict()
+        for key in original_state_dict.keys():
+            assert key in loaded_state_dict
+            original_value = original_state_dict[key]
+            loaded_value = loaded_state_dict[key]
+            assert original_value == loaded_value
+
+    _check_state_dict(trainer, trainer2)
+    _check_state_dict(scheduler, scheduler2)
+    _check_state_dict(early, early2)
+    _check_state_dict(checkpointer, checkpointer2)
+
+    trainer2.run([0], max_epochs=6)
+
+    # early stopping should have triggered
+    assert trainer2.state.epoch == 4
+
+    # If Checkpoint's state was restored correctly, it should continue to respect n_saved
+    # and delete old checkpoints, and have the correct last_checkpoint.
+    assert os.listdir(dirname) == ["checkpoint_4.pt"]
+    assert checkpointer2.last_checkpoint == "checkpoint_4.pt"
+
+
+def test_save_model_optimizer_lr_scheduler_with_validation(dirname):
+    _test_save_model_optimizer_lr_scheduler_with_validation("cpu", dirname)
+
+
 def test_checkpoint_load_objects():
 
     with pytest.raises(TypeError, match=r"Argument checkpoint should be a dictionary"):
@@ -1023,3 +1209,47 @@ def _test_tpu_saves_to_cpu_nprocs(index, dirname):
 def test_distrib_single_device_xla_nprocs(xmp_executor, dirname):
     n = int(os.environ["NUM_TPU_WORKERS"])
     xmp_executor(_test_tpu_saves_to_cpu_nprocs, args=(dirname,), nprocs=n)
+
+
+def _setup_checkpoint():
+    save_handler = MagicMock(spec=BaseSaveHandler)
+    model = DummyModel()
+    to_save = {"model": model}
+
+    checkpointer = Checkpoint(to_save, save_handler=save_handler, n_saved=None)
+    assert checkpointer.last_checkpoint is None
+
+    trainer = Engine(lambda e, b: None)
+    trainer.state = State(epoch=0, iteration=0)
+
+    checkpointer(trainer)
+    trainer.state.iteration = 10
+    checkpointer(trainer)
+    trainer.state.iteration = 20
+    checkpointer(trainer)
+    assert save_handler.call_count == 3
+    return checkpointer
+
+
+def test_checkpoint_state_dict():
+    checkpointer = _setup_checkpoint()
+    sd = checkpointer.state_dict()
+    assert "saved" in sd
+    assert isinstance(sd["saved"], list) and len(sd["saved"]) == len(checkpointer._saved)
+
+    for saved_item, true_item in zip(sd["saved"], checkpointer._saved):
+        assert saved_item[0] == true_item.priority
+        assert saved_item[1] == true_item.filename
+
+
+def test_checkpoint_load_state_dict():
+    true_checkpointer = _setup_checkpoint()
+
+    save_handler = MagicMock(spec=BaseSaveHandler)
+    model = DummyModel()
+    to_save = {"model": model}
+    checkpointer = Checkpoint(to_save, save_handler=save_handler, n_saved=None)
+
+    sd = {"saved": [(0, "model_0.pt"), (10, "model_10.pt"), (20, "model_20.pt")]}
+    checkpointer.load_state_dict(sd)
+    assert checkpointer._saved == true_checkpointer._saved
diff --git a/tests/ignite/handlers/test_early_stopping.py b/tests/ignite/handlers/test_early_stopping.py
index 3005fe31..793c922a 100644
--- a/tests/ignite/handlers/test_early_stopping.py
+++ b/tests/ignite/handlers/test_early_stopping.py
@@ -49,6 +49,31 @@ def test_simple_early_stopping():
     assert trainer.should_terminate
 
 
+def test_state_dict():
+
+    scores = iter([1.0, 0.8, 0.88])
+
+    def score_function(engine):
+        return next(scores)
+
+    trainer = Engine(do_nothing_update_fn)
+
+    h = EarlyStopping(patience=2, score_function=score_function, trainer=trainer)
+    # Call 3 times and check if stopped
+    assert not trainer.should_terminate
+    h(None)
+    assert not trainer.should_terminate
+
+    # Swap to new object, but maintain state
+    h2 = EarlyStopping(patience=2, score_function=score_function, trainer=trainer)
+    h2.load_state_dict(h.state_dict())
+
+    h2(None)
+    assert not trainer.should_terminate
+    h2(None)
+    assert trainer.should_terminate
+
+
 def test_early_stopping_on_delta():
 
     scores = iter([1.0, 2.0, 2.01, 3.0, 3.01, 3.02])