Repository: pytorch/ignite

Issue:
Add state_dict user keys
## 🚀 Feature

There can be cases when user would like to save some custom state attributes when storing the engine. It can be helpful if `engine.state_dict` could return in addition user defined keys.
Proposed API:
```python
engine = Engine(...)
engine.state_dict_user_keys.append("user_value")
...

@engine.on(Events.STARTED)
def init_user_value(_):
     engine.state.user_value = ...

@engine.on(...)
def save_engine(_):
    state_dict = engine.state_dict()
    assert "user_value" in state_dict
    torch.save(state_dict, "/tmp/engine.pt")
```



Discussion:
It is not wanted to save automatically all what is added in `state` ?
No, as we can have non-serializable or other large objects that are just for exchange purposes... 

Resolving patch:
diff --git a/ignite/engine/engine.py b/ignite/engine/engine.py
--- a/ignite/engine/engine.py
+++ b/ignite/engine/engine.py
@@ -5,7 +5,7 @@
 import weakref
 import random
 import warnings
-from typing import Union, Optional, Callable, Iterable, Iterator, Any, Tuple
+from typing import Union, Optional, Callable, Iterable, Iterator, Any, Tuple, List
 
 import torch
 
@@ -127,6 +127,7 @@ def __init__(self, process_function: Callable):
         self.should_terminate = False
         self.should_terminate_single_epoch = False
         self.state = None
+        self._state_dict_user_keys = []
         self._allowed_events = []
 
         self._dataloader_iter = None
@@ -504,24 +505,49 @@ def _handle_exception(self, e: Exception) -> None:
         else:
             raise e
 
+    @property
+    def state_dict_user_keys(self) -> List:
+        return self._state_dict_user_keys
+
     def state_dict(self) -> OrderedDict:
-        """Returns a dictionary containing engine's state: "seed", "epoch_length", "max_epochs" and "iteration"
+        """Returns a dictionary containing engine's state: "seed", "epoch_length", "max_epochs" and "iteration" and
+        other state values defined by `engine.state_dict_user_keys`
+
+        .. code-block:: python
+
+            engine = Engine(...)
+            engine.state_dict_user_keys.append("alpha")
+            engine.state_dict_user_keys.append("beta")
+            ...
+
+            @engine.on(Events.STARTED)
+            def init_user_value(_):
+                 engine.state.alpha = 0.1
+                 engine.state.beta = 1.0
+
+            @engine.on(Events.COMPLETED)
+            def save_engine(_):
+                state_dict = engine.state_dict()
+                assert "alpha" in state_dict and "beta" in state_dict
+                torch.save(state_dict, "/tmp/engine.pt")
 
         Returns:
-            dict:
+            OrderedDict:
                 a dictionary containing engine's state
 
         """
         if self.state is None:
             return OrderedDict()
         keys = self._state_dict_all_req_keys + (self._state_dict_one_of_opt_keys[0],)
+        keys += tuple(self._state_dict_user_keys)
         return OrderedDict([(k, getattr(self.state, k)) for k in keys])
 
     def load_state_dict(self, state_dict: Mapping) -> None:
         """Setups engine from `state_dict`.
 
         State dictionary should contain keys: `iteration` or `epoch` and `max_epochs`, `epoch_length` and
-        `seed`. Iteration and epoch values are 0-based: the first iteration or epoch is zero.
+        `seed`. If `engine.state_dict_user_keys` contains keys, they should be also present in the state dictionary.
+        Iteration and epoch values are 0-based: the first iteration or epoch is zero.
 
         Args:
             state_dict (Mapping): a dict with parameters
@@ -547,6 +573,13 @@ def load_state_dict(self, state_dict: Mapping) -> None:
                 raise ValueError(
                     "Required state attribute '{}' is absent in provided state_dict '{}'".format(k, state_dict.keys())
                 )
+        for k in self._state_dict_user_keys:
+            if k not in state_dict:
+                raise ValueError(
+                    "Required user state attribute '{}' is absent in provided state_dict '{}'".format(
+                        k, state_dict.keys()
+                    )
+                )
 
         opts = [k in state_dict for k in self._state_dict_one_of_opt_keys]
         if (not any(opts)) or (all(opts)):
@@ -558,6 +591,8 @@ def load_state_dict(self, state_dict: Mapping) -> None:
             epoch_length=state_dict["epoch_length"],
             metrics={},
         )
+        for k in self._state_dict_user_keys:
+            setattr(self.state, k, state_dict[k])
 
         if "iteration" in state_dict:
             self.state.iteration = state_dict["iteration"]


Test patch:
diff --git a/tests/ignite/engine/test_resume_run.py b/tests/ignite/engine/test_resume_run.py
--- a/tests/ignite/engine/test_resume_run.py
+++ b/tests/ignite/engine/test_resume_run.py
@@ -27,6 +27,27 @@ def _test(state):
     _test(State(seed=0, epoch=5, epoch_length=1000, max_epochs=100))
 
 
+def test_state_dict_with_user_keys():
+    engine = Engine(lambda e, b: 1)
+    engine.state_dict_user_keys.append("alpha")
+    engine.state_dict_user_keys.append("beta")
+
+    def _test(state):
+        engine.state = state
+        sd = engine.state_dict()
+        assert isinstance(sd, Mapping) and len(sd) == len(engine._state_dict_all_req_keys) + 1 + len(
+            engine.state_dict_user_keys
+        )
+        assert sd["seed"] == engine.state.seed
+        assert sd["iteration"] == engine.state.iteration
+        assert sd["epoch_length"] == engine.state.epoch_length
+        assert sd["max_epochs"] == engine.state.max_epochs
+        assert sd["alpha"] == engine.state.alpha
+        assert sd["beta"] == engine.state.beta
+
+    _test(State(seed=0, iteration=500, epoch_length=1000, max_epochs=100, alpha=0.01, beta="Good"))
+
+
 def test_state_dict_integration():
     engine = Engine(lambda e, b: 1)
     data = list(range(100))
@@ -54,6 +75,11 @@ def test_load_state_dict_asserts():
     with pytest.raises(ValueError, match=r"state_dict should contain only one of"):
         engine.load_state_dict({"seed": 0, "max_epochs": 100, "epoch_length": 120, "iteration": 12, "epoch": 123})
 
+    engine = Engine(lambda e, b: 1)
+    engine.state_dict_user_keys.append("alpha")
+    with pytest.raises(ValueError, match=r"Required user state attribute "):
+        engine.load_state_dict({"seed": 0, "max_epochs": 100, "epoch_length": 120, "iteration": 12, "epoch": 123})
+
 
 def test_load_state_dict():
     engine = Engine(lambda e, b: 1)
@@ -72,6 +98,26 @@ def _test(sd):
     _test({"seed": 0, "max_epochs": 100, "epoch_length": 120, "epoch": 5})
 
 
+def test_load_state_dict_with_user_keys():
+    engine = Engine(lambda e, b: 1)
+    engine.state_dict_user_keys.append("alpha")
+    engine.state_dict_user_keys.append("beta")
+
+    def _test(sd):
+        engine.load_state_dict(sd)
+        assert sd["seed"] == engine.state.seed
+        if "iteration" in sd:
+            assert sd["iteration"] == engine.state.iteration
+        elif "epoch" in sd:
+            assert sd["epoch"] == engine.state.epoch
+        assert sd["epoch_length"] == engine.state.epoch_length
+        assert sd["max_epochs"] == engine.state.max_epochs
+        assert sd["alpha"] == engine.state.alpha
+        assert sd["beta"] == engine.state.beta
+
+    _test({"seed": 0, "max_epochs": 100, "epoch_length": 120, "iteration": 123, "alpha": 0.1, "beta": "abc"})
+
+
 def test_load_state_dict_integration(counter_factory):
     engine = Engine(lambda e, b: 1)
 
@@ -84,6 +130,29 @@ def test_load_state_dict_integration(counter_factory):
     engine.run(data)
 
 
+def test_state_dict_with_user_keys_integration(dirname, counter_factory):
+    engine = Engine(lambda e, b: 1)
+    engine.state_dict_user_keys.append("alpha")
+
+    @engine.on(Events.STARTED)
+    def init_user_values(_):
+        engine.state.alpha = 0.1
+
+    fp = os.path.join(dirname, "engine.pt")
+
+    @engine.on(Events.COMPLETED)
+    def save_engine(_):
+        state_dict = engine.state_dict()
+        assert "alpha" in state_dict
+        torch.save(state_dict, fp)
+
+    engine.run([0, 1])
+
+    assert os.path.exists(fp)
+    state_dict = torch.load(fp)
+    assert "alpha" in state_dict and state_dict["alpha"] == 0.1
+
+
 class BatchChecker:
     def __init__(self, data, init_counter=0):
         self.counter = init_counter