Fix state_dict handling for max_iters in Engine

ignite/base/mixins.py

@@ -1,7 +1,7 @@
 from abc import ABCMeta, abstractmethod
 from collections import OrderedDict
 from collections.abc import Mapping
-from typing import Any, TYPE_CHECKING
+from typing import Any, Tuple, TYPE_CHECKING
 
 if TYPE_CHECKING:
     from ignite.engine import Engine
@@ -28,8 +28,8 @@ def attach(self, engine: "Engine", *args: Any, **kwargs: Any) -> None:
 
 
 class Serializable:
-    _state_dict_all_req_keys: tuple = ()
-    _state_dict_one_of_opt_keys: tuple = ()
+    _state_dict_all_req_keys: Tuple[str, ...] = ()
+    _state_dict_one_of_opt_keys: Tuple[Tuple[str, ...], ...] = ((),)
 
     def state_dict(self) -> OrderedDict:
         raise NotImplementedError
@@ -43,6 +43,13 @@ def load_state_dict(self, state_dict: Mapping) -> None:
                 raise ValueError(
                     f"Required state attribute '{k}' is absent in provided state_dict '{state_dict.keys()}'"
                 )
-        opts = [k in state_dict for k in self._state_dict_one_of_opt_keys]
-        if len(opts) > 0 and ((not any(opts)) or (all(opts))):
-            raise ValueError(f"state_dict should contain only one of '{self._state_dict_one_of_opt_keys}' keys")
+
+        # Handle groups of one-of optional keys
+        for one_of_opt_keys in self._state_dict_one_of_opt_keys:
+            if len(one_of_opt_keys) > 0:
+                opts = [k in state_dict for k in one_of_opt_keys]
+                num_present = sum(opts)
+                if num_present == 0:
+                    raise ValueError(f"state_dict should contain at least one of '{one_of_opt_keys}' keys")
+                if num_present > 1:
+                    raise ValueError(f"state_dict should contain only one of '{one_of_opt_keys}' keys")

ignite/engine/engine.py

@@ -6,7 +6,7 @@
 import weakref
 from collections import defaultdict, OrderedDict
 from collections.abc import Callable, Generator, Iterable, Iterator, Mapping
-from typing import Any
+from typing import Any, Optional
 
 from torch.utils.data import DataLoader
 
@@ -128,8 +128,8 @@ def compute_mean_std(engine, batch):
 
     """
 
-    _state_dict_all_req_keys = ("epoch_length", "max_epochs")
-    _state_dict_one_of_opt_keys = ("iteration", "epoch")
+    _state_dict_all_req_keys = ("epoch_length",)
+    _state_dict_one_of_opt_keys = (("iteration", "epoch"), ("max_epochs", "max_iters"))
 
     # Flag to disable engine._internal_run as generator feature for BC
     interrupt_resume_enabled = True
@@ -707,15 +707,25 @@ def save_engine(_):
             OrderedDict:
                 a dictionary containing engine's state
 
+        .. versionchanged:: 0.5.5
+            Added support for serializing ``max_iters``.
+
         """
-        keys: tuple[str, ...] = self._state_dict_all_req_keys + (self._state_dict_one_of_opt_keys[0],)
+        keys: tuple[str, ...] = self._state_dict_all_req_keys
+        # We add iteration by default to get exact measure of progress
+        keys += ("iteration",)
+        # Include either max_epochs or max_iters based on which was originally set
+        if self.state.max_iters is not None:
+            keys += ("max_iters",)
+        else:
+            keys += ("max_epochs",)
         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`, `max_epochs` and `epoch_length`.
+        State dictionary should contain keys: `iteration` or `epoch`, `max_epochs` or `max_iters`, and `epoch_length`.
         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.
 
@@ -726,15 +736,20 @@ def load_state_dict(self, state_dict: Mapping) -> None:
 
         .. code-block:: python
 
-            # Restore from the 4rd epoch
+            # Restore from the 4th epoch
             state_dict = {"epoch": 3, "max_epochs": 100, "epoch_length": len(data_loader)}
             # or 500th iteration
             # state_dict = {"iteration": 499, "max_epochs": 100, "epoch_length": len(data_loader)}
+            # or with max_iters
+            # state_dict = {"iteration": 499, "max_iters": 1000, "epoch_length": len(data_loader)}
 
             trainer = Engine(...)
             trainer.load_state_dict(state_dict)
             trainer.run(data)
 
+        .. versionchanged:: 0.5.5
+            Added support for restoring from a state dict containing ``max_iters`` instead of ``max_epochs``.
+
         """
         super().load_state_dict(state_dict)
 
@@ -743,17 +758,15 @@ def load_state_dict(self, state_dict: Mapping) -> None:
                 raise ValueError(
                     f"Required user state attribute '{k}' is absent in provided state_dict '{state_dict.keys()}'"
                 )
-        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])
+        self.state.epoch_length = state_dict["epoch_length"]
 
         if "iteration" in state_dict:
             self.state.iteration = state_dict["iteration"]
             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:
+        else:  # epoch is in state_dict
             self.state.epoch = state_dict["epoch"]
             if self.state.epoch_length is None:
                 raise ValueError(
@@ -762,6 +775,36 @@ def load_state_dict(self, state_dict: Mapping) -> None:
                 )
             self.state.iteration = self.state.epoch_length * self.state.epoch
 
+        # Set max_epochs or max_iters with validation
+        max_epochs_value = state_dict.get("max_epochs", None)
+        max_iters_value = state_dict.get("max_iters", None)
+
+        # Validate max_epochs if present
+        if max_epochs_value is not None:
+            if max_epochs_value < 1:
+                raise ValueError("max_epochs in state_dict is invalid. Please, set a correct max_epochs positive value")
+            if max_epochs_value < self.state.epoch:
+                raise ValueError(
+                    "max_epochs in state_dict should be larger than or equal to the current epoch "
+                    f"defined in the state: {max_epochs_value} vs {self.state.epoch}. "
+                )
+            self.state.max_epochs = max_epochs_value
+        else:
+            self.state.max_epochs = None
+
+        # Validate max_iters if present
+        if max_iters_value is not None:
+            if max_iters_value < 1:
+                raise ValueError("max_iters in state_dict is invalid. Please, set a correct max_iters positive value")
+            if max_iters_value < self.state.iteration:
+                raise ValueError(
+                    "max_iters in state_dict should be larger than or equal to the current iteration "
+                    f"defined in the state: {max_iters_value} vs {self.state.iteration}. "
+                )
+            self.state.max_iters = max_iters_value
+        else:
+            self.state.max_iters = None
+
     @staticmethod
     def _is_done(state: State) -> bool:
         is_done_iters = state.max_iters is not None and state.iteration >= state.max_iters
@@ -773,6 +816,36 @@ def _is_done(state: State) -> bool:
         is_done_epochs = state.max_epochs is not None and state.epoch >= state.max_epochs
         return is_done_iters or is_done_count or is_done_epochs
 
+    def _check_and_set_max_epochs(self, max_epochs: Optional[int] = None) -> None:
+        """Validate and set max_epochs with proper checks."""
+        if max_epochs is not None:
+            if max_epochs < 1:
+                raise ValueError("Argument max_epochs is invalid. Please, set a correct max_epochs positive value")
+            # Only validate if training is actually done - allow resuming interrupted training
+            if self.state.max_epochs is not None and max_epochs < self.state.epoch:
+                raise ValueError(
+                    "Argument max_epochs should be greater than or equal to the start "
+                    f"epoch defined in the state: {max_epochs} vs {self.state.epoch}. "
+                    "Please, set engine.state.max_epochs = None "
+                    "before calling engine.run() in order to restart the training from the beginning."
+                )
+            self.state.max_epochs = max_epochs
+
+    def _check_and_set_max_iters(self, max_iters: Optional[int] = None) -> None:
+        """Validate and set max_iters with proper checks."""
+        if max_iters is not None:
+            if max_iters < 1:
+                raise ValueError("Argument max_iters is invalid. Please, set a correct max_iters positive value")
+            # Only validate if training is actually done - allow resuming interrupted training
+            if (self.state.max_iters is not None) and max_iters < self.state.iteration:
+                raise ValueError(
+                    "Argument max_iters should be greater than or equal to the start "
+                    f"iteration defined in the state: {max_iters} vs {self.state.iteration}. "
+                    "Please, set engine.state.max_iters = None "
+                    "before calling engine.run() in order to restart the training from the beginning."
+                )
+            self.state.max_iters = max_iters
+
     def set_data(self, data: Iterable | DataLoader) -> None:
         """Method to set data. After calling the method the next batch passed to `processing_function` is
         from newly provided data. Please, note that epoch length is not modified.
@@ -871,43 +944,44 @@ def switch_batch(engine):
         if data is not None and not isinstance(data, Iterable):
             raise TypeError("Argument data should be iterable")
 
+        if max_epochs is not None and max_iters is not None:
+            raise ValueError(
+                "Arguments max_iters and max_epochs are mutually exclusive.Please provide only max_epochs or max_iters."
+            )
+
         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:
-                    raise ValueError(
-                        "Argument max_epochs should be greater than or equal to the start "
-                        f"epoch defined in the state: {max_epochs} vs {self.state.epoch}. "
-                        "Please, set engine.state.max_epochs = None "
-                        "before calling engine.run() in order to restart the training from the beginning."
-                    )
-                self.state.max_epochs = max_epochs
-            if epoch_length is not None:
-                if epoch_length != self.state.epoch_length:
-                    raise ValueError(
-                        "Argument epoch_length should be same as in the state, "
-                        f"but given {epoch_length} vs {self.state.epoch_length}"
-                    )
+            self._check_and_set_max_epochs(max_epochs)
+
+        if self.state.max_iters is not None:
+            self._check_and_set_max_iters(max_iters)
+
+        # Check if we need to create new state or resume
+        # Create new state if:
+        # 1. No termination params set (first run), OR
+        # 2. Training is done AND generator is None
+        should_create_new_state = (self.state.max_epochs is None and self.state.max_iters is None) or (
+            self._is_done(self.state) and self._internal_run_generator is None
+        )
 
-        if self.state.max_epochs is None or (self._is_done(self.state) and self._internal_run_generator is None):
+        if should_create_new_state:
             # Create new state
-            if epoch_length is None:
-                if data is None:
-                    raise ValueError("epoch_length should be provided if data is None")
+            if data is None and epoch_length is None and self.state.epoch_length is None:
+                raise ValueError("epoch_length should be provided if data is None")
 
-                epoch_length = self._get_data_length(data)
-                if epoch_length is not None and epoch_length < 1:
-                    raise ValueError("Input data has zero size. Please provide non-empty data")
+            # Set epoch_length for new state
+            if epoch_length is None:
+                # Try to get from data first, then fall back to existing state
+                if data is not None:
+                    epoch_length = self._get_data_length(data)
+                if epoch_length is None and self.state.epoch_length is not None:
+                    epoch_length = self.state.epoch_length
+            if epoch_length is not None and epoch_length < 1:
+                raise ValueError("Input data has zero size. Please provide non-empty data")
 
             if max_iters is None:
                 if max_epochs is None:
                     max_epochs = 1
             else:
-                if max_epochs is not None:
-                    raise ValueError(
-                        "Arguments max_iters and max_epochs are mutually exclusive."
-                        "Please provide only max_epochs or max_iters."
-                    )
                 if epoch_length is not None:
                     max_epochs = math.ceil(max_iters / epoch_length)
 
@@ -918,18 +992,41 @@ def switch_batch(engine):
             self.state.epoch_length = epoch_length
             # Reset generator if previously used
             self._internal_run_generator = None
-            self.logger.info(f"Engine run starting with max_epochs={max_epochs}.")
+
+            if self.state.max_epochs is not None:
+                self.logger.info(f"Engine run starting with max_epochs={self.state.max_epochs}.")
+            else:
+                self.logger.info(f"Engine run starting with max_iters={self.state.max_iters}.")
         else:
-            self.logger.info(
-                f"Engine run resuming from iteration {self.state.iteration}, "
-                f"epoch {self.state.epoch} until {self.state.max_epochs} epochs"
-            )
+            if self.state.epoch_length is not None:
+                if epoch_length is not None and epoch_length != self.state.epoch_length:
+                    raise ValueError(
+                        "Argument epoch_length should be same as in the state, "
+                        f"but given {epoch_length} vs {self.state.epoch_length}"
+                    )
+            else:
+                if epoch_length is None and data is not None:
+                    epoch_length = self._get_data_length(data)
+                if epoch_length is not None:
+                    if epoch_length < 1:
+                        raise ValueError("Input data has zero size. Please provide non-empty data")
+                    self.state.epoch_length = epoch_length
+
+            if self.state.max_epochs is not None:
+                self.logger.info(
+                    f"Engine run resuming from iteration {self.state.iteration}, "
+                    f"epoch {self.state.epoch} until {self.state.max_epochs} epochs"
+                )
+            else:
+                self.logger.info(
+                    f"Engine run resuming from iteration {self.state.iteration}, "
+                    f"epoch {self.state.epoch} until {self.state.max_iters} iterations"
+                )
+
             if self.state.epoch_length is None and data is None:
                 raise ValueError("epoch_length should be provided if data is None")
 
             if self.should_terminate:
-                # If engine was terminated and now is resuming from terminated state
-                # we need to initialize iter_counter as 0
                 self._init_iter = 0
 
         if self._dataloader_iter is None: