Repository: pytorch/ignite

Issue:
Initialize engine.state before calling run
## 🚀 Feature

Idea is to improve flexibility of Engine and simplify user code by initializing `engine.state` before calling `run`.

What we gain:
```python
trainer = Engine(...)
trainer.state.alpha = 0.0
trainer.state.beta = 1.0
```



Discussion:
vfdev-5: > I don't understand yet why lr_finder implies to save trainer but it's a detail :)

because when lr finder is run, engine is run too. Thus we need to come back to the initial state.

Patch:
diff --git a/ignite/contrib/handlers/lr_finder.py b/ignite/contrib/handlers/lr_finder.py
index a6804ea3..c0a19c20 100644
--- a/ignite/contrib/handlers/lr_finder.py
+++ b/ignite/contrib/handlers/lr_finder.py
@@ -310,7 +310,9 @@ class FastaiLRFinder:
         # store to_save
         with tempfile.TemporaryDirectory() as tmpdirname:
             obj = {k: o.state_dict() for k, o in to_save.items()}
-            cache_filepath = Path(tmpdirname) / "ignite_lr_finder_cache.pt.tar"
+            # add trainer
+            obj["trainer"] = trainer.state_dict()
+            cache_filepath = Path(tmpdirname) / "ignite_lr_finder_cache.pt"
             torch.save(obj, cache_filepath.as_posix())
 
             optimizer = to_save["optimizer"]
@@ -336,9 +338,10 @@ class FastaiLRFinder:
             self._detach(trainer)
             # restore to_save and reset trainer's state
             obj = torch.load(cache_filepath.as_posix())
-            trainer.state = None
+            trainer.load_state_dict(obj["trainer"])
             for k, o in obj.items():
-                to_save[k].load_state_dict(o)
+                if k in to_save:
+                    to_save[k].load_state_dict(o)
 
 
 class _ExponentialLR(_LRScheduler):
diff --git a/ignite/contrib/handlers/param_scheduler.py b/ignite/contrib/handlers/param_scheduler.py
index ecbe9121..f129ad41 100644
--- a/ignite/contrib/handlers/param_scheduler.py
+++ b/ignite/contrib/handlers/param_scheduler.py
@@ -51,7 +51,7 @@ class ParamScheduler(metaclass=ABCMeta):
             name = self.param_name
 
         if self.save_history:
-            if not hasattr(engine.state, "param_history"):
+            if not hasattr(engine.state, "param_history") or engine.state.param_history is None:
                 setattr(engine.state, "param_history", {})
             engine.state.param_history.setdefault(name, [])
             values = [pg[self.param_name] for pg in self.optimizer_param_groups]
diff --git a/ignite/engine/engine.py b/ignite/engine/engine.py
index 73857fd1..8983afc3 100644
--- a/ignite/engine/engine.py
+++ b/ignite/engine/engine.py
@@ -5,7 +5,7 @@ import warnings
 import weakref
 from collections import OrderedDict, defaultdict
 from collections.abc import Mapping
-from typing import Any, Callable, Iterable, List, Optional, Tuple
+from typing import Any, Callable, Iterable, List, Optional
 
 from ignite._utils import _to_hours_mins_secs
 from ignite.base import Serializable
@@ -24,7 +24,8 @@ class Engine(Serializable):
 
     Attributes:
         state (State): object that is used to pass internal and user-defined state between event handlers.
-            It is created and reset on every :meth:`~ignite.engine.Engine.run`.
+            It is created with the engine and its attributes (e.g. `state.iteration`, `state.epoch` etc) are reset
+            on every :meth:`~ignite.engine.Engine.run`.
         last_event_name (Events): last event name triggered by the engine.
 
     Examples:
@@ -125,7 +126,7 @@ class Engine(Serializable):
         self.last_event_name = None
         self.should_terminate = False
         self.should_terminate_single_epoch = False
-        self.state = None
+        self.state = State()
         self._state_dict_user_keys = []
         self._allowed_events = []
 
@@ -195,6 +196,8 @@ class Engine(Serializable):
             self._allowed_events.append(e)
             if event_to_attr and e in event_to_attr:
                 State.event_to_attr[e] = event_to_attr[e]
+        # we need to update state attributes associated with new custom events
+        self.state._update_attrs()
 
     def _handler_wrapper(self, handler: Callable, event_name: Any, event_filter: Callable) -> Callable:
         # signature of the following wrapper will be inspected during registering to check if engine is necessary
@@ -468,8 +471,6 @@ class Engine(Serializable):
                 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])
@@ -481,6 +482,8 @@ class Engine(Serializable):
         `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.
 
+        This method does not remove any custom attributs added by user.
+
         Args:
             state_dict (Mapping): a dict with parameters
 
@@ -505,16 +508,23 @@ class Engine(Serializable):
                         k, state_dict.keys()
                     )
                 )
-
-        self.state = State(max_epochs=state_dict["max_epochs"], epoch_length=state_dict["epoch_length"], metrics={},)
+        self.state.max_epochs = state_dict["max_epochs"]
+        self.state.epoch_length = state_dict["epoch_length"]
         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"]
-            self.state.epoch = self.state.iteration // self.state.epoch_length
+            self.state.epoch = 0
+            if self.state.epoch_length is not None:
+                self.state.epoch = self.state.iteration // self.state.epoch_length
         elif "epoch" in state_dict:
             self.state.epoch = state_dict["epoch"]
+            if self.state.epoch_length is None:
+                raise ValueError(
+                    "If epoch is provided in the state dict, epoch_length should not be None. "
+                    "Input state_dict: {}".format(state_dict)
+                )
             self.state.iteration = self.state.epoch_length * self.state.epoch
 
     @staticmethod
@@ -610,7 +620,7 @@ class Engine(Serializable):
                 "Please, use torch.manual_seed or ignite.utils.manual_seed"
             )
 
-        if self.state is not None:
+        if self.state.max_epochs is not None:
             # Check and apply overridden parameters
             if max_epochs is not None:
                 if max_epochs < self.state.epoch:
@@ -627,7 +637,7 @@ class Engine(Serializable):
                         )
                     )
 
-        if self.state is None or self._is_done(self.state):
+        if self.state.max_epochs is None or self._is_done(self.state):
             # Create new state
             if max_epochs is None:
                 max_epochs = 1
@@ -637,7 +647,10 @@ class Engine(Serializable):
                     if epoch_length < 1:
                         raise ValueError("Input data has zero size. Please provide non-empty data")
 
-            self.state = State(iteration=0, epoch=0, max_epochs=max_epochs, epoch_length=epoch_length)
+            self.state.iteration = 0
+            self.state.epoch = 0
+            self.state.max_epochs = max_epochs
+            self.state.epoch_length = epoch_length
             self.logger.info("Engine run starting with max_epochs={}.".format(max_epochs))
         else:
             self.logger.info(
diff --git a/ignite/engine/events.py b/ignite/engine/events.py
index 82af5a4e..7e9e02ea 100644
--- a/ignite/engine/events.py
+++ b/ignite/engine/events.py
@@ -310,6 +310,9 @@ class State:
         for k, v in kwargs.items():
             setattr(self, k, v)
 
+        self._update_attrs()
+
+    def _update_attrs(self):
         for value in self.event_to_attr.values():
             if not hasattr(self, value):
                 setattr(self, value, 0)


Test patch:
diff --git a/tests/ignite/contrib/handlers/test_lr_finder.py b/tests/ignite/contrib/handlers/test_lr_finder.py
index 8de6d45f..cf22659b 100644
--- a/tests/ignite/contrib/handlers/test_lr_finder.py
+++ b/tests/ignite/contrib/handlers/test_lr_finder.py
@@ -136,8 +136,9 @@ def test_model_optimizer_reset(lr_finder, to_save, dummy_engine, dataloader):
     init_model_sd = copy.deepcopy(model.state_dict())
     init_trainer_sd = copy.deepcopy(dummy_engine.state_dict())
 
-    with lr_finder.attach(dummy_engine, to_save=to_save, diverge_th=float("inf")) as trainer_with_finder:
-        trainer_with_finder.run(dataloader)
+    with pytest.warns(UserWarning, match=r"Run completed without loss diverging"):
+        with lr_finder.attach(dummy_engine, to_save=to_save, diverge_th=float("inf")) as trainer_with_finder:
+            trainer_with_finder.run(dataloader)
 
     assert init_optimizer_sd == optimizer.state_dict()
     assert init_model_sd == model.state_dict()
@@ -167,18 +168,23 @@ def assert_output_sizes(lr_finder, dummy_engine):
 
 def test_num_iter_is_none(lr_finder, to_save, dummy_engine, dataloader):
 
-    with lr_finder.attach(dummy_engine, to_save=to_save, diverge_th=float("inf")) as trainer_with_finder:
-        trainer_with_finder.run(dataloader)
-        assert_output_sizes(lr_finder, dummy_engine)
-        assert dummy_engine.state.iteration == len(dataloader)
+    with pytest.warns(UserWarning, match=r"Run completed without loss diverging"):
+        with lr_finder.attach(dummy_engine, to_save=to_save, diverge_th=float("inf")) as trainer_with_finder:
+            trainer_with_finder.run(dataloader)
+            assert_output_sizes(lr_finder, dummy_engine)
+            assert dummy_engine.state.iteration == len(dataloader)
 
 
 def test_num_iter_is_enough(lr_finder, to_save, dummy_engine, dataloader):
-    with lr_finder.attach(dummy_engine, to_save=to_save, num_iter=50, diverge_th=float("inf")) as trainer_with_finder:
-        trainer_with_finder.run(dataloader)
-        assert_output_sizes(lr_finder, dummy_engine)
-        # -1 because it terminates when state.iteration > num_iter
-        assert dummy_engine.state.iteration - 1 == 50
+
+    with pytest.warns(UserWarning, match=r"Run completed without loss diverging"):
+        with lr_finder.attach(
+            dummy_engine, to_save=to_save, num_iter=50, diverge_th=float("inf")
+        ) as trainer_with_finder:
+            trainer_with_finder.run(dataloader)
+            assert_output_sizes(lr_finder, dummy_engine)
+            # -1 because it terminates when state.iteration > num_iter
+            assert dummy_engine.state.iteration - 1 == 50
 
 
 def test_num_iter_is_not_enough(lr_finder, to_save, dummy_engine, dataloader):
diff --git a/tests/ignite/contrib/handlers/test_param_scheduler.py b/tests/ignite/contrib/handlers/test_param_scheduler.py
index f7b74c19..499ce437 100644
--- a/tests/ignite/contrib/handlers/test_param_scheduler.py
+++ b/tests/ignite/contrib/handlers/test_param_scheduler.py
@@ -384,6 +384,8 @@ def test_concat_scheduler_two_schedulers():
             assert lrs == pytest.approx([v for i, v in simulated_values])
             concat_scheduler.load_state_dict(state_dict)
 
+            trainer.state.param_history = None
+
     _test(duration_vals_as_np_int=False)
     _test(duration_vals_as_np_int=True)
 
@@ -456,6 +458,8 @@ def test_concat_scheduler_two_linear():
         assert lrs == pytest.approx([v for i, v in simulated_values])
         concat_scheduler.load_state_dict(state_dict)
 
+        trainer.state.param_history = None
+
 
 def test_concat_scheduler_3_schedulers():
     tensor = torch.zeros([1], requires_grad=True)
@@ -527,6 +531,8 @@ def test_concat_scheduler_3_schedulers():
         assert lrs == pytest.approx([v for i, v in simulated_values])
         concat_scheduler.load_state_dict(state_dict)
 
+        trainer.state.param_history = None
+
 
 def test_save_param_history():
     tensor = torch.zeros([1], requires_grad=True)
@@ -1022,6 +1028,8 @@ def test_create_lr_scheduler_with_warmup_on_combined_scheduler():
                 param_history = trainer.state.param_history["lr"]
                 assert lrs == pytest.approx([v[0] for v in param_history])
 
+                trainer.state.param_history = None
+
             scheduler.load_state_dict(state_dict)
 
     _test(save_history=False)
diff --git a/tests/ignite/engine/test_engine_state_dict.py b/tests/ignite/engine/test_engine_state_dict.py
index bb4d485f..a9413853 100644
--- a/tests/ignite/engine/test_engine_state_dict.py
+++ b/tests/ignite/engine/test_engine_state_dict.py
@@ -11,7 +11,10 @@ from tests.ignite.engine import BatchChecker, EpochCounter, IterationCounter
 def test_state_dict():
     engine = Engine(lambda e, b: 1)
     sd = engine.state_dict()
-    assert isinstance(sd, Mapping) and len(sd) == 0
+    assert isinstance(sd, Mapping) and len(sd) == 3
+    assert "iteration" in sd and sd["iteration"] == 0
+    assert "max_epochs" in sd and sd["max_epochs"] is None
+    assert "epoch_length" in sd and sd["epoch_length"] is None
 
     def _test(state):
         engine.state = state
@@ -76,6 +79,10 @@ def test_load_state_dict_asserts():
     with pytest.raises(ValueError, match=r"Required user state attribute"):
         engine.load_state_dict({"max_epochs": 100, "epoch_length": 120, "iteration": 12})
 
+    engine = Engine(lambda e, b: 1)
+    with pytest.raises(ValueError, match=r"If epoch is provided in the state dict, epoch_length should not be None"):
+        engine.load_state_dict({"max_epochs": 100, "epoch": 2, "epoch_length": None})
+
 
 def test_load_state_dict():
     engine = Engine(lambda e, b: 1)
@@ -148,6 +155,12 @@ def test_load_state_dict_with_params_overriding_integration():
         engine.run(data, epoch_length=90)
 
 
+def test_empty_state_dict_load_state_dict():
+    engine = Engine(lambda e, b: 1)
+    sd = engine.state_dict()
+    engine.load_state_dict(sd)
+
+
 def test_continue_training():
     # Tests issue : https://github.com/pytorch/ignite/issues/993
     max_epochs = 2
@@ -234,3 +247,23 @@ def test_epoch_length():
 
     _test_as_iter(data, 1, num_iters)
     _test_as_iter(data, 2, num_iters // 2)
+
+
+def test_state_custom_attrs_init():
+    def _test(with_load_state_dict=False):
+        engine = Engine(lambda e, b: None)
+        engine.state.alpha = 0.0
+        engine.state.beta = 1.0
+
+        if with_load_state_dict:
+            engine.load_state_dict({"iteration": 3, "max_epochs": 5, "epoch_length": 5})
+
+        @engine.on(Events.STARTED | Events.EPOCH_STARTED | Events.EPOCH_COMPLETED | Events.COMPLETED)
+        def check_custom_attr():
+            assert hasattr(engine.state, "alpha") and engine.state.alpha == 0.0
+            assert hasattr(engine.state, "beta") and engine.state.beta == 1.0
+
+        engine.run([0, 1, 2, 3, 4], max_epochs=5)
+
+    _test()
+    _test(with_load_state_dict=True)