Repository: pytorch/ignite

Issue:
Add any parameter scheduling
## 🚀 Feature

Following https://github.com/pytorch/ignite/issues/1912 we would like to enlarging the scope of `ParamScheduler` to scheduling any parameter (e.g. hyperparameters) like it is done in MONAI: Project-MONAI/MONAI#1949




Discussion:
I can try to tackle it if some help is needed!
 @danieltudosiu would you like to go on with this issue ?

Resolving patch:
diff --git a/ignite/handlers/__init__.py b/ignite/handlers/__init__.py
--- a/ignite/handlers/__init__.py
+++ b/ignite/handlers/__init__.py
@@ -7,6 +7,7 @@
 from ignite.handlers.ema_handler import EMAHandler
 from ignite.handlers.lr_finder import FastaiLRFinder
 from ignite.handlers.param_scheduler import (
+    BaseParamScheduler,
     ConcatScheduler,
     CosineAnnealingScheduler,
     CyclicalScheduler,
@@ -17,6 +18,14 @@
     PiecewiseLinear,
     create_lr_scheduler_with_warmup,
 )
+from ignite.handlers.state_param_scheduler import (
+    ExpStateScheduler,
+    LambdaStateScheduler,
+    MultiStepStateScheduler,
+    PiecewiseLinearStateScheduler,
+    StateParamScheduler,
+    StepStateScheduler,
+)
 from ignite.handlers.stores import EpochOutputStore
 from ignite.handlers.terminate_on_nan import TerminateOnNan
 from ignite.handlers.time_limit import TimeLimit
@@ -46,6 +55,13 @@
     "EMAHandler",
     "BasicTimeProfiler",
     "HandlersTimeProfiler",
+    "BaseParamScheduler",
+    "StateParamScheduler",
+    "LambdaStateScheduler",
+    "PiecewiseLinearStateScheduler",
+    "ExpStateScheduler",
+    "StepStateScheduler",
+    "MultiStepStateScheduler",
 ]
 
 
diff --git a/ignite/handlers/param_scheduler.py b/ignite/handlers/param_scheduler.py
--- a/ignite/handlers/param_scheduler.py
+++ b/ignite/handlers/param_scheduler.py
@@ -15,83 +15,26 @@
 from ignite.engine import Engine
 
 
-class ParamScheduler(metaclass=ABCMeta):
-    """An abstract class for updating an optimizer's parameter value during
+class BaseParamScheduler(metaclass=ABCMeta):
+    r"""An abstract class for updating an engine state or optimizer's parameter value during
     training.
 
     Args:
-        optimizer: torch optimizer or any object with attribute ``param_groups``
-            as a sequence.
-        param_name: name of optimizer's parameter to update.
+        param_name: name of engine state or optimizer's parameter to update.
         save_history: whether to log the parameter values to
             `engine.state.param_history`, (default=False).
-        param_group_index: optimizer's parameters group to use
 
-    Note:
-        Parameter scheduler works independently of the internal state of the attached optimizer.
-        More precisely, whatever the state of the optimizer (newly created or used by another scheduler) the scheduler
-        sets defined absolute values.
+    .. versionadded:: 0.5.0
 
-    .. versionadded:: 0.4.5
     """
 
     def __init__(
-        self,
-        optimizer: Optimizer,
-        param_name: str,
-        save_history: bool = False,
-        param_group_index: Optional[int] = None,
+        self, param_name: str, save_history: bool = False,
     ):
-
-        if not (
-            isinstance(optimizer, Optimizer)
-            or (hasattr(optimizer, "param_groups") and isinstance(optimizer.param_groups, Sequence))
-        ):
-            raise TypeError(
-                "Argument optimizer should be torch.optim.Optimizer or has attribute 'param_groups' as list/tuple, "
-                f"but given {type(optimizer)}"
-            )
-
-        self.optimizer = optimizer
-        self.param_group_index = param_group_index
         self.param_name = param_name
         self.event_index = 0
         self._save_history = save_history
-        self._state_attrs = ["event_index", "param_name", "save_history", "param_group_index"]
-
-    def __call__(self, engine: Optional[Engine], name: Optional[str] = None) -> None:
-
-        value = self.get_param()
-
-        if isinstance(value, list):
-            if len(value) != len(self.optimizer_param_groups):
-                raise ValueError(
-                    "size of value is different than optimizer_param_groups "
-                    f"{len(value)} != {len(self.optimizer_param_groups)}"
-                )
-
-            for i, param_group in enumerate(self.optimizer_param_groups):
-                param_group[self.param_name] = value[i]
-        else:
-            for i, param_group in enumerate(self.optimizer_param_groups):
-                param_group[self.param_name] = value
-
-        if name is None:
-            name = self.param_name
-
-        if self.save_history and engine:
-            if not hasattr(engine.state, "param_history") or engine.state.param_history is None:  # type: ignore
-                setattr(engine.state, "param_history", {})
-            engine.state.param_history.setdefault(name, [])  # type: ignore[attr-defined]
-            values = [pg[self.param_name] for pg in self.optimizer_param_groups]
-            engine.state.param_history[name].append(values)  # type: ignore[attr-defined]
-        self.event_index += 1
-
-    @property
-    def optimizer_param_groups(self) -> List[Dict[str, Any]]:
-        if self.param_group_index is None:
-            return self.optimizer.param_groups
-        return [self.optimizer.param_groups[self.param_group_index]]
+        self._state_attrs = ["event_index", "param_name", "save_history"]
 
     @property
     def save_history(self) -> bool:
@@ -102,11 +45,11 @@ def save_history(self, value: bool) -> None:
         self._save_history = value
 
     def state_dict(self) -> Dict[str, Any]:
-        """Returns a dictionary containing a whole state of ParamScheduler.
+        """Returns a dictionary containing a whole state of BaseParamScheduler.
 
         Returns:
             dict:
-                a dictionary containing a whole state of ParamScheduler
+                a dictionary containing a whole state of BaseParamScheduler
         """
         destination = OrderedDict()
         for name in self._state_attrs:
@@ -118,7 +61,7 @@ def state_dict(self) -> Dict[str, Any]:
         return destination
 
     def load_state_dict(self, state_dict: Mapping) -> None:
-        """Copies parameters from :attr:`state_dict` into this ParamScheduler.
+        """Copies parameters from :attr:`state_dict` into this BaseParamScheduler.
 
         Args:
             state_dict: a dict containing parameters.
@@ -140,7 +83,7 @@ def load_state_dict(self, state_dict: Mapping) -> None:
 
     @abstractmethod
     def get_param(self) -> Union[List[float], float]:
-        """Method to get current optimizer's parameter values
+        """Method to get current parameter values
 
         Returns:
             list of params, or scalar param
@@ -148,6 +91,7 @@ def get_param(self) -> Union[List[float], float]:
         pass
 
     @classmethod
+    @abstractmethod
     def simulate_values(cls, num_events: int, **scheduler_kwargs: Any) -> List[List[int]]:
         """Method to simulate scheduled values during `num_events` events.
 
@@ -157,29 +101,8 @@ def simulate_values(cls, num_events: int, **scheduler_kwargs: Any) -> List[List[
 
         Returns:
             event_index, value
-
-        Examples:
-            .. code-block:: python
-
-                lr_values = np.array(LinearCyclicalScheduler.simulate_values(num_events=50, param_name='lr',
-                                                                             start_value=1e-1, end_value=1e-3,
-                                                                             cycle_size=10))
-
-                plt.plot(lr_values[:, 0], lr_values[:, 1], label="learning rate")
-                plt.xlabel("events")
-                plt.ylabel("values")
-                plt.legend()
         """
-        keys_to_remove = ["optimizer", "save_history"]
-        for key in keys_to_remove:
-            if key in scheduler_kwargs:
-                del scheduler_kwargs[key]
-        values = []
-        scheduler = cls(optimizer=_get_fake_optimizer(), save_history=False, **scheduler_kwargs)
-        for i in range(num_events):
-            scheduler(engine=None)
-            values.append([i, scheduler.optimizer_param_groups[0][scheduler.param_name]])
-        return values
+        pass
 
     @classmethod
     def plot_values(cls, num_events: int, **scheduler_kwargs: Mapping) -> Any:
@@ -208,7 +131,7 @@ def plot_values(cls, num_events: int, **scheduler_kwargs: Mapping) -> Any:
                                                     start_value=1e-1, end_value=1e-3, cycle_size=10))
         """
         try:
-            import matplotlib.pylab as plt
+            import matplotlib.pyplot as plt
         except ImportError:
             raise RuntimeError(
                 "This method requires matplotlib to be installed. "
@@ -223,6 +146,118 @@ def plot_values(cls, num_events: int, **scheduler_kwargs: Mapping) -> Any:
         return ax
 
 
+class ParamScheduler(BaseParamScheduler):
+    """An abstract class for updating an optimizer's parameter value during
+    training.
+
+    Args:
+        optimizer: torch optimizer or any object with attribute ``param_groups``
+            as a sequence.
+        param_name: name of optimizer's parameter to update.
+        save_history: whether to log the parameter values to
+            `engine.state.param_history`, (default=False).
+        param_group_index: optimizer's parameters group to use
+
+    Note:
+        Parameter scheduler works independently of the internal state of the attached optimizer.
+        More precisely, whatever the state of the optimizer (newly created or used by another scheduler) the scheduler
+        sets defined absolute values.
+
+
+    """
+
+    def __init__(
+        self,
+        optimizer: Optimizer,
+        param_name: str,
+        save_history: bool = False,
+        param_group_index: Optional[int] = None,
+    ):
+        super(ParamScheduler, self).__init__(param_name, save_history)
+        if not (
+            isinstance(optimizer, Optimizer)
+            or (hasattr(optimizer, "param_groups") and isinstance(optimizer.param_groups, Sequence))
+        ):
+            raise TypeError(
+                "Argument optimizer should be torch.optim.Optimizer or has attribute 'param_groups' as list/tuple, "
+                f"but given {type(optimizer)}"
+            )
+
+        self.optimizer = optimizer
+        self.param_group_index = param_group_index
+        self._state_attrs += ["param_group_index"]
+
+    def __call__(self, engine: Optional[Engine], name: Optional[str] = None) -> None:
+
+        value = self.get_param()
+
+        if isinstance(value, list):
+            if len(value) != len(self.optimizer_param_groups):
+                raise ValueError(
+                    "size of value is different than optimizer_param_groups "
+                    f"{len(value)} != {len(self.optimizer_param_groups)}"
+                )
+
+            for i, param_group in enumerate(self.optimizer_param_groups):
+                param_group[self.param_name] = value[i]
+        else:
+            for i, param_group in enumerate(self.optimizer_param_groups):
+                param_group[self.param_name] = value
+
+        if name is None:
+            name = self.param_name
+
+        if self.save_history and engine:
+            if not hasattr(engine.state, "param_history") or engine.state.param_history is None:  # type: ignore
+                setattr(engine.state, "param_history", {})
+            engine.state.param_history.setdefault(name, [])  # type: ignore[attr-defined]
+            values = [pg[self.param_name] for pg in self.optimizer_param_groups]
+            engine.state.param_history[name].append(values)  # type: ignore[attr-defined]
+        self.event_index += 1
+
+    @property
+    def optimizer_param_groups(self) -> List[Dict[str, Any]]:
+        if self.param_group_index is None:
+            return self.optimizer.param_groups
+        return [self.optimizer.param_groups[self.param_group_index]]
+
+    @classmethod
+    def simulate_values(cls, num_events: int, **scheduler_kwargs: Any) -> List[List[int]]:
+        """Method to simulate scheduled values during `num_events` events.
+
+        Args:
+            num_events: number of events during the simulation.
+            scheduler_kwargs: parameter scheduler configuration kwargs.
+
+        Returns:
+            event_index, value
+
+        Examples:
+
+        .. code-block:: python
+
+            lr_values = np.array(LinearCyclicalScheduler.simulate_values(num_events=50, param_name='lr',
+                                                                         start_value=1e-1, end_value=1e-3,
+                                                                         cycle_size=10))
+
+            plt.plot(lr_values[:, 0], lr_values[:, 1], label="learning rate")
+            plt.xlabel("events")
+            plt.ylabel("values")
+            plt.legend()
+
+        """
+        keys_to_remove = ["optimizer", "save_history"]
+        for key in keys_to_remove:
+            if key in scheduler_kwargs:
+                del scheduler_kwargs[key]
+        values = []
+        scheduler = cls(optimizer=_get_fake_optimizer(), save_history=False, **scheduler_kwargs)
+        for i in range(num_events):
+            scheduler(engine=None)
+            values.append([i, scheduler.optimizer_param_groups[0][scheduler.param_name]])
+        return values
+
+
 class CyclicalScheduler(ParamScheduler):
     """An abstract class for updating an optimizer's parameter value over a
     cycle of some size.
diff --git a/ignite/handlers/state_param_scheduler.py b/ignite/handlers/state_param_scheduler.py
new file mode 100644
--- /dev/null
+++ b/ignite/handlers/state_param_scheduler.py
@@ -0,0 +1,398 @@
+import numbers
+from bisect import bisect_right
+from typing import Any, List, Sequence, Tuple, Union
+
+from ignite.engine import CallableEventWithFilter, Engine, Events, EventsList
+from ignite.handlers import BaseParamScheduler
+
+
+class StateParamScheduler(BaseParamScheduler):
+    """An abstract class for updating an engine state parameter values during training.
+
+    Args:
+        param_name: name of parameter to update.
+        save_history: whether to log the parameter values to
+            `engine.state.param_history`, (default=False).
+
+    Note:
+        Parameter scheduler works independently of the internal state of the attached engine.
+        More precisely, whatever the state of the engine (newly created or used by another scheduler) the scheduler
+        sets defined absolute values.
+
+    .. versionadded:: 0.5.0
+
+    """
+
+    def __init__(
+        self, param_name: str, save_history: bool = False,
+    ):
+        super(StateParamScheduler, self).__init__(param_name, save_history)
+
+    def attach(
+        self,
+        engine: Engine,
+        event: Union[str, Events, CallableEventWithFilter, EventsList] = Events.ITERATION_COMPLETED,
+    ) -> None:
+        """Attach the handler to the engine. Once the handler is attached, the ``Engine.state`` will have a new
+        attribute with the name ``param_name``. Then the current value of the parameter can be retrieved from
+        ``Engine.state`` when the engine is running.
+
+        Args:
+            engine: trainer to which the handler will be attached.
+            event: trigger ``param_name`` value update.
+
+        """
+        if hasattr(engine.state, self.param_name):
+            raise ValueError(
+                f"Attribute: '{self.param_name}' is already defined in the Engine.state."
+                f"This may be a conflict between multiple StateParameterScheduler handlers."
+                f"Please choose another name."
+            )
+
+        if self.save_history:
+            if not hasattr(engine.state, "param_history") or engine.state.param_history is None:  # type: ignore
+                setattr(engine.state, "param_history", {})
+                engine.state.param_history.setdefault(self.param_name, [])  # type: ignore[attr-defined]
+
+        engine.add_event_handler(event, self)
+
+    def __call__(self, engine: Engine) -> None:
+        self.event_index += 1
+        value = self.get_param()
+        setattr(engine.state, self.param_name, value)
+        if self.save_history:
+            engine.state.param_history[self.param_name].append(value)  # type: ignore[attr-defined]
+
+    @classmethod
+    def simulate_values(cls, num_events: int, **scheduler_kwargs: Any) -> List[List[int]]:
+        """Method to simulate scheduled engine state parameter values during `num_events` events.
+
+        Args:
+            num_events: number of events during the simulation.
+            scheduler_kwargs: parameter scheduler configuration kwargs.
+
+        Returns:
+            event_index, value
+
+        Examples:
+
+        .. code-block:: python
+
+            import matplotlib.pyplot as plt
+            import numpy as np
+
+            step_state_param_values = np.array(
+                StepStateScheduler.simulate_values(
+                    num_events=20, param_name="step_scheduled_param", initial_value=10, gamma=0.99, step_size=5
+                )
+            )
+
+            plt.plot(step_state_param_values[:, 0], step_state_param_values[:, 1], label="learning rate")
+            plt.xlabel("events")
+            plt.ylabel("values")
+            plt.legend()
+
+        """
+        for key in ["save_history"]:
+            if key in scheduler_kwargs:
+                del scheduler_kwargs[key]
+        values = []
+        scheduler = cls(save_history=False, **scheduler_kwargs)
+        engine = Engine(lambda e, b: None)
+        for i in range(num_events):
+            scheduler(engine=engine)
+            values.append([i, getattr(engine.state, scheduler_kwargs["param_name"])])
+        return values
+
+
+class LambdaStateScheduler(StateParamScheduler):
+    """Update a parameter during training by using a user defined callable object.
+        User defined callable object is taking an event index as input and returns parameter value.
+
+    Args:
+        lambda_obj: user defined callable object.
+        param_name: name of parameter to update.
+        save_history: whether to log the parameter values to
+            `engine.state.param_history`, (default=False).
+
+    Examples:
+
+        .. code-block:: python
+
+            ...
+            engine = Engine(train_step)
+
+            class LambdaState:
+                def __init__(self, initial_value, gamma):
+                    self.initial_value = initial_value
+                    self.gamma = gamma
+
+                def __call__(self, event_index):
+                    return self.initial_value * self.gamma ** (event_index % 9)
+
+
+            param_scheduler = LambdaStateScheduler(
+                param_name="param",
+                lambda_obj=LambdaState(10, 0.99),
+            )
+
+            param_scheduler.attach(engine, Events.EPOCH_COMPLETED)
+
+            # basic handler to print scheduled state parameter
+            engine.add_event_handler(Events.EPOCH_COMPLETED, lambda _ : print(engine.state.param))
+
+            engine.run([0] * 8, max_epochs=2)
+
+    .. versionadded:: 0.5.0
+
+    """
+
+    def __init__(self, lambda_obj: Any, param_name: str, save_history: bool = False):
+        super(LambdaStateScheduler, self).__init__(param_name, save_history)
+
+        if not callable(lambda_obj):
+            raise ValueError("Expected lambda_obj to be callable.")
+
+        self.lambda_obj = lambda_obj
+        self._state_attrs += ["lambda_obj"]
+
+    def get_param(self) -> Union[List[float], float]:
+        return self.lambda_obj(self.event_index)
+
+
+class PiecewiseLinearStateScheduler(StateParamScheduler):
+    """Piecewise linear state parameter scheduler.
+
+    Args:
+        milestones_values: list of tuples (event index, parameter value)
+            represents milestones and parameter values. Milestones should be increasing integers.
+        param_name: name of parameter to update.
+        save_history: whether to log the parameter values to
+            `engine.state.param_history`, (default=False).
+
+    Examples:
+
+        .. code-block:: python
+
+            ...
+            engine = Engine(train_step)
+
+            param_scheduler = PiecewiseLinearStateScheduler(
+                param_name="param",
+                milestones_values=[(10, 0.5), (20, 0.45), (21, 0.3), (30, 0.1), (40, 0.1)]
+            )
+
+            param_scheduler.attach(engine, Events.EPOCH_COMPLETED)
+
+            # basic handler to print scheduled state parameter
+            engine.add_event_handler(Events.EPOCH_COMPLETED, lambda _ : print(engine.state.param))
+
+            engine.run([0] * 8, max_epochs=40)
+            #
+            # Sets the state parameter value to 0.5 over the first 10 iterations, then decreases linearly from 0.5 to
+            # 0.45 between 10th and 20th iterations. Next there is a jump to 0.3 at the 21st iteration and the state
+            # parameter value decreases linearly from 0.3 to 0.1 between 21st and 30th iterations and remains 0.1 until
+            # the end of the iterations.
+            #
+
+    .. versionadded:: 0.5.0
+    """
+
+    def __init__(
+        self, milestones_values: List[Tuple[int, float]], param_name: str, save_history: bool = False,
+    ):
+        super(PiecewiseLinearStateScheduler, self).__init__(param_name, save_history)
+
+        if not isinstance(milestones_values, Sequence):
+            raise TypeError(
+                f"Argument milestones_values should be a list or tuple, but given {type(milestones_values)}"
+            )
+        if len(milestones_values) < 1:
+            raise ValueError(
+                f"Argument milestones_values should be with at least one value, but given {milestones_values}"
+            )
+
+        values = []  # type: List[float]
+        milestones = []  # type: List[int]
+        for pair in milestones_values:
+            if not isinstance(pair, tuple) or len(pair) != 2:
+                raise ValueError("Argument milestones_values should be a list of pairs (milestone, param_value)")
+            if not isinstance(pair[0], numbers.Integral):
+                raise TypeError(f"Value of a milestone should be integer, but given {type(pair[0])}")
+            if len(milestones) > 0 and pair[0] < milestones[-1]:
+                raise ValueError(
+                    f"Milestones should be increasing integers, but given {pair[0]} is smaller "
+                    f"than the previous milestone {milestones[-1]}"
+                )
+            milestones.append(pair[0])
+            values.append(pair[1])
+
+        self.values = values
+        self.milestones = milestones
+        self._index = 0
+        self._state_attrs += ["values", "milestones", "_index"]
+
+    def _get_start_end(self) -> Tuple[int, int, float, float]:
+        if self.milestones[0] > self.event_index:
+            return self.event_index - 1, self.event_index, self.values[0], self.values[0]
+        elif self.milestones[-1] <= self.event_index:
+            return (self.event_index, self.event_index + 1, self.values[-1], self.values[-1])
+        elif self.milestones[self._index] <= self.event_index < self.milestones[self._index + 1]:
+            return (
+                self.milestones[self._index],
+                self.milestones[self._index + 1],
+                self.values[self._index],
+                self.values[self._index + 1],
+            )
+        else:
+            self._index += 1
+            return self._get_start_end()
+
+    def get_param(self) -> Union[List[float], float]:
+        start_index, end_index, start_value, end_value = self._get_start_end()
+        return start_value + (end_value - start_value) * (self.event_index - start_index) / (end_index - start_index)
+
+
+class ExpStateScheduler(StateParamScheduler):
+    """Update a parameter during training by using exponential function.
+    The function decays the parameter value by gamma every step.
+    Based on the closed form of ExponentialLR from PyTorch
+    https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.ExponentialLR.html
+
+    Args:
+        initial_value: Starting value of the parameter.
+        gamma: Multiplicative factor of parameter value decay.
+        param_name: name of parameter to update.
+        save_history: whether to log the parameter values to
+            `engine.state.param_history`, (default=False).
+
+    Examples:
+
+        .. code-block:: python
+
+            ...
+            engine = Engine(train_step)
+
+            param_scheduler = ExpStateScheduler(
+                param_name="param", initial_value=10, gamma=0.99
+            )
+
+            param_scheduler.attach(engine, Events.EPOCH_COMPLETED)
+
+            # basic handler to print scheduled state parameter
+            engine.add_event_handler(Events.EPOCH_COMPLETED, lambda _ : print(engine.state.param))
+
+            engine.run([0] * 8, max_epochs=2)
+
+    .. versionadded:: 0.5.0
+
+    """
+
+    def __init__(
+        self, initial_value: float, gamma: float, param_name: str, save_history: bool = False,
+    ):
+        super(ExpStateScheduler, self).__init__(param_name, save_history)
+        self.initial_value = initial_value
+        self.gamma = gamma
+        self._state_attrs += ["initial_value", "gamma"]
+
+    def get_param(self) -> Union[List[float], float]:
+        return self.initial_value * self.gamma ** self.event_index
+
+
+class StepStateScheduler(StateParamScheduler):
+    """Update a parameter during training by using a step function.
+    This function decays the parameter value by gamma every step_size.
+    Based on StepLR from PyTorch.
+    https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.StepLR.html
+
+    Args:
+        initial_value: Starting value of the parameter.
+        gamma: Multiplicative factor of parameter value decay.
+        step_size: Period of parameter value decay.
+        param_name: name of parameter to update.
+        save_history: whether to log the parameter values to
+            `engine.state.param_history`, (default=False).
+
+    Examples:
+
+        .. code-block:: python
+
+            ...
+            engine = Engine(train_step)
+
+            param_scheduler = StepStateScheduler(
+                param_name="param", initial_value=10, gamma=0.99, step_size=5
+            )
+
+            param_scheduler.attach(engine, Events.EPOCH_COMPLETED)
+
+            # basic handler to print scheduled state parameter
+            engine.add_event_handler(Events.EPOCH_COMPLETED, lambda _ : print(engine.state.param))
+
+            engine.run([0] * 8, max_epochs=10)
+
+    .. versionadded:: 0.5.0
+
+    """
+
+    def __init__(
+        self, initial_value: float, gamma: float, step_size: int, param_name: str, save_history: bool = False,
+    ):
+        super(StepStateScheduler, self).__init__(param_name, save_history)
+        self.initial_value = initial_value
+        self.gamma = gamma
+        self.step_size = step_size
+        self._state_attrs += ["initial_value", "gamma", "step_size"]
+
+    def get_param(self) -> Union[List[float], float]:
+        return self.initial_value * self.gamma ** (self.event_index // self.step_size)
+
+
+class MultiStepStateScheduler(StateParamScheduler):
+    """Update a parameter during training by using a multi step function.
+    The function decays the parameter value by gamma once the number of steps reaches one of the milestones.
+    Based on MultiStepLR from PyTorch.
+    https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.MultiStepLR.html
+
+    Args:
+        initial_value: Starting value of the parameter.
+        gamma: Multiplicative factor of parameter value decay.
+        milestones: List of step indices. Must be increasing.
+        param_name: name of parameter to update.
+        save_history: whether to log the parameter values to
+            `engine.state.param_history`, (default=False).
+
+    Examples:
+
+        .. code-block:: python
+
+            ...
+            engine = Engine(train_step)
+
+            param_scheduler = MultiStepStateScheduler(
+                param_name="param", initial_value=10, gamma=0.99, milestones=[3, 6],
+            )
+
+            param_scheduler.attach(engine, Events.EPOCH_COMPLETED)
+
+            # basic handler to print scheduled state parameter
+            engine.add_event_handler(Events.EPOCH_COMPLETED, lambda _ : print(engine.state.param))
+
+            engine.run([0] * 8, max_epochs=10)
+
+    .. versionadded:: 0.5.0
+
+    """
+
+    def __init__(
+        self, initial_value: float, gamma: float, milestones: List[int], param_name: str, save_history: bool = False,
+    ):
+        super(MultiStepStateScheduler, self).__init__(param_name, save_history)
+        self.initial_value = initial_value
+        self.gamma = gamma
+        self.milestones = milestones
+        self._state_attrs += ["initial_value", "gamma", "milestones"]
+
+    def get_param(self) -> Union[List[float], float]:
+        return self.initial_value * self.gamma ** bisect_right(self.milestones, self.event_index)


Test patch:
diff --git a/tests/ignite/handlers/test_param_scheduler.py b/tests/ignite/handlers/test_param_scheduler.py
--- a/tests/ignite/handlers/test_param_scheduler.py
+++ b/tests/ignite/handlers/test_param_scheduler.py
@@ -899,7 +899,7 @@ def save_lr(engine):
     )
 
     with pytest.raises(RuntimeError, match=r"This method requires matplotlib to be installed."):
-        with patch.dict("sys.modules", {"matplotlib.pylab": None}):
+        with patch.dict("sys.modules", {"matplotlib.pyplot": None}):
             _test(
                 PiecewiseLinear,
                 optimizer=optimizer,
diff --git a/tests/ignite/handlers/test_state_param_scheduler.py b/tests/ignite/handlers/test_state_param_scheduler.py
new file mode 100644
--- /dev/null
+++ b/tests/ignite/handlers/test_state_param_scheduler.py
@@ -0,0 +1,405 @@
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from ignite.engine import Engine, Events
+from ignite.handlers.state_param_scheduler import (
+    ExpStateScheduler,
+    LambdaStateScheduler,
+    MultiStepStateScheduler,
+    PiecewiseLinearStateScheduler,
+    StepStateScheduler,
+)
+
+config1 = (3, [(2, 0), (5, 10)], True, [0.0, 0.0, 3.3333333333333335])
+expected_hist2 = [0.0] * 10 + [float(i) for i in range(1, 11)] + [10.0] * 10
+config2 = (30, [(10, 0), (20, 10)], True, expected_hist2)
+
+
+@pytest.mark.parametrize(
+    "max_epochs, milestones_values,  save_history, expected_param_history", [config1, config2],
+)
+def test_pwlinear_scheduler_linear_increase_history(
+    max_epochs, milestones_values, save_history, expected_param_history
+):
+    # Testing linear increase
+    engine = Engine(lambda e, b: None)
+    pw_linear_step_parameter_scheduler = PiecewiseLinearStateScheduler(
+        param_name="pwlinear_scheduled_param", milestones_values=milestones_values, save_history=save_history,
+    )
+    pw_linear_step_parameter_scheduler.attach(engine, Events.EPOCH_COMPLETED)
+    engine.run([0] * 8, max_epochs=max_epochs)
+    expected_param_history = expected_param_history
+    assert hasattr(engine.state, "param_history")
+    state_param = engine.state.param_history["pwlinear_scheduled_param"]
+    assert len(state_param) == len(expected_param_history)
+    assert state_param == expected_param_history
+
+    state_dict = pw_linear_step_parameter_scheduler.state_dict()
+    pw_linear_step_parameter_scheduler.load_state_dict(state_dict)
+
+
+@pytest.mark.parametrize(
+    "max_epochs, milestones_values", [(3, [(3, 12), (5, 10)]), (5, [(10, 12), (20, 10)]),],
+)
+def test_pwlinear_scheduler_step_constant(max_epochs, milestones_values):
+    # Testing step_constant
+    engine = Engine(lambda e, b: None)
+    linear_state_parameter_scheduler = PiecewiseLinearStateScheduler(
+        param_name="pwlinear_scheduled_param", milestones_values=milestones_values
+    )
+    linear_state_parameter_scheduler.attach(engine, Events.EPOCH_COMPLETED)
+    engine.run([0] * 8, max_epochs=max_epochs)
+    torch.testing.assert_allclose(getattr(engine.state, "pwlinear_scheduled_param"), milestones_values[0][1])
+
+    state_dict = linear_state_parameter_scheduler.state_dict()
+    linear_state_parameter_scheduler.load_state_dict(state_dict)
+
+
+@pytest.mark.parametrize(
+    "max_epochs, milestones_values, expected_val",
+    [(2, [(0, 0), (3, 10)], 6.666666666666667), (10, [(0, 0), (20, 10)], 5.0),],
+)
+def test_pwlinear_scheduler_linear_increase(max_epochs, milestones_values, expected_val):
+    # Testing linear increase
+    engine = Engine(lambda e, b: None)
+    linear_state_parameter_scheduler = PiecewiseLinearStateScheduler(
+        param_name="pwlinear_scheduled_param", milestones_values=milestones_values
+    )
+    linear_state_parameter_scheduler.attach(engine, Events.EPOCH_COMPLETED)
+    engine.run([0] * 8, max_epochs=max_epochs)
+    torch.testing.assert_allclose(getattr(engine.state, "pwlinear_scheduled_param"), expected_val, atol=0.001, rtol=0.0)
+
+    state_dict = linear_state_parameter_scheduler.state_dict()
+    linear_state_parameter_scheduler.load_state_dict(state_dict)
+
+
+@pytest.mark.parametrize(
+    "max_epochs, milestones_values,", [(3, [(0, 0), (3, 10)]), (40, [(0, 0), (20, 10)])],
+)
+def test_pwlinear_scheduler_max_value(
+    max_epochs, milestones_values,
+):
+    # Testing max_value
+    engine = Engine(lambda e, b: None)
+    linear_state_parameter_scheduler = PiecewiseLinearStateScheduler(
+        param_name="linear_scheduled_param", milestones_values=milestones_values,
+    )
+    linear_state_parameter_scheduler.attach(engine, Events.EPOCH_COMPLETED)
+    engine.run([0] * 8, max_epochs=max_epochs)
+    torch.testing.assert_allclose(getattr(engine.state, "linear_scheduled_param"), milestones_values[-1][1])
+
+    state_dict = linear_state_parameter_scheduler.state_dict()
+    linear_state_parameter_scheduler.load_state_dict(state_dict)
+
+
+def test_piecewiselinear_asserts():
+
+    with pytest.raises(TypeError, match=r"Argument milestones_values should be a list or tuple"):
+        PiecewiseLinearStateScheduler(param_name="linear_scheduled_param", milestones_values=None)
+
+    with pytest.raises(ValueError, match=r"Argument milestones_values should be with at least one value"):
+        PiecewiseLinearStateScheduler(param_name="linear_scheduled_param", milestones_values=[])
+
+    with pytest.raises(ValueError, match=r"Argument milestones_values should be a list of pairs"):
+        PiecewiseLinearStateScheduler(param_name="linear_scheduled_param", milestones_values=[(0.5,)])
+
+    with pytest.raises(ValueError, match=r"Argument milestones_values should be a list of pairs"):
+        PiecewiseLinearStateScheduler(param_name="linear_scheduled_param", milestones_values=[(10, 0.5), (0.6,)])
+
+    with pytest.raises(ValueError, match=r"Milestones should be increasing integers"):
+        PiecewiseLinearStateScheduler(param_name="linear_scheduled_param", milestones_values=[(10, 0.5), (5, 0.6)])
+
+    with pytest.raises(TypeError, match=r"Value of a milestone should be integer"):
+        PiecewiseLinearStateScheduler(param_name="linear_scheduled_param", milestones_values=[(0.5, 1)])
+
+
+@pytest.mark.parametrize(
+    "max_epochs, initial_value, gamma", [(3, 10, 0.99), (40, 5, 0.98)],
+)
+def test_exponential_scheduler(max_epochs, initial_value, gamma):
+    engine = Engine(lambda e, b: None)
+    exp_state_parameter_scheduler = ExpStateScheduler(
+        param_name="exp_scheduled_param", initial_value=initial_value, gamma=gamma
+    )
+    exp_state_parameter_scheduler.attach(engine, Events.EPOCH_COMPLETED)
+    engine.run([0] * 8, max_epochs=max_epochs)
+    torch.testing.assert_allclose(getattr(engine.state, "exp_scheduled_param"), initial_value * gamma ** max_epochs)
+
+    state_dict = exp_state_parameter_scheduler.state_dict()
+    exp_state_parameter_scheduler.load_state_dict(state_dict)
+
+
+@pytest.mark.parametrize(
+    "max_epochs, initial_value, gamma, step_size", [(3, 10, 0.99, 5), (40, 5, 0.98, 22)],
+)
+def test_step_scheduler(
+    max_epochs, initial_value, gamma, step_size,
+):
+    engine = Engine(lambda e, b: None)
+    step_state_parameter_scheduler = StepStateScheduler(
+        param_name="step_scheduled_param", initial_value=initial_value, gamma=gamma, step_size=step_size
+    )
+    step_state_parameter_scheduler.attach(engine, Events.EPOCH_COMPLETED)
+    engine.run([0] * 8, max_epochs=max_epochs)
+    torch.testing.assert_allclose(
+        getattr(engine.state, "step_scheduled_param"), initial_value * gamma ** (max_epochs // step_size)
+    )
+
+    state_dict = step_state_parameter_scheduler.state_dict()
+    step_state_parameter_scheduler.load_state_dict(state_dict)
+
+
+from bisect import bisect_right
+
+
+@pytest.mark.parametrize(
+    "max_epochs, initial_value, gamma, milestones", [(3, 10, 0.99, [3, 6]), (40, 5, 0.98, [3, 6, 9, 10, 11])],
+)
+def test_multistep_scheduler(
+    max_epochs, initial_value, gamma, milestones,
+):
+    engine = Engine(lambda e, b: None)
+    multi_step_state_parameter_scheduler = MultiStepStateScheduler(
+        param_name="multistep_scheduled_param", initial_value=initial_value, gamma=gamma, milestones=milestones,
+    )
+    multi_step_state_parameter_scheduler.attach(engine, Events.EPOCH_COMPLETED)
+    engine.run([0] * 8, max_epochs=max_epochs)
+    torch.testing.assert_allclose(
+        getattr(engine.state, "multistep_scheduled_param"),
+        initial_value * gamma ** bisect_right(milestones, max_epochs),
+    )
+
+    state_dict = multi_step_state_parameter_scheduler.state_dict()
+    multi_step_state_parameter_scheduler.load_state_dict(state_dict)
+
+
+def test_custom_scheduler():
+
+    engine = Engine(lambda e, b: None)
+
+    class LambdaState:
+        def __init__(self, initial_value, gamma):
+            self.initial_value = initial_value
+            self.gamma = gamma
+
+        def __call__(self, event_index):
+            return self.initial_value * self.gamma ** (event_index % 9)
+
+    lambda_state_parameter_scheduler = LambdaStateScheduler(
+        param_name="custom_scheduled_param", lambda_obj=LambdaState(initial_value=10, gamma=0.99),
+    )
+    lambda_state_parameter_scheduler.attach(engine, Events.EPOCH_COMPLETED)
+    engine.run([0] * 8, max_epochs=2)
+    torch.testing.assert_allclose(
+        getattr(engine.state, "custom_scheduled_param"), LambdaState(initial_value=10, gamma=0.99)(2)
+    )
+    engine.run([0] * 8, max_epochs=20)
+    torch.testing.assert_allclose(
+        getattr(engine.state, "custom_scheduled_param"), LambdaState(initial_value=10, gamma=0.99)(20)
+    )
+
+    state_dict = lambda_state_parameter_scheduler.state_dict()
+    lambda_state_parameter_scheduler.load_state_dict(state_dict)
+
+
+def test_custom_scheduler_asserts():
+    class LambdaState:
+        def __init__(self, initial_value, gamma):
+            self.initial_value = initial_value
+            self.gamma = gamma
+
+    with pytest.raises(ValueError, match=r"Expected lambda_obj to be callable."):
+        lambda_state_parameter_scheduler = LambdaStateScheduler(
+            param_name="custom_scheduled_param", lambda_obj=LambdaState(initial_value=10, gamma=0.99),
+        )
+
+
+config1 = (
+    PiecewiseLinearStateScheduler,
+    {"param_name": "linear_scheduled_param", "milestones_values": [(3, 12), (5, 10)]},
+)
+config2 = (ExpStateScheduler, {"param_name": "exp_scheduled_param", "initial_value": 10, "gamma": 0.99})
+config3 = (
+    MultiStepStateScheduler,
+    {"param_name": "multistep_scheduled_param", "initial_value": 10, "gamma": 0.99, "milestones": [3, 6]},
+)
+
+
+class LambdaState:
+    def __init__(self, initial_value, gamma):
+        self.initial_value = initial_value
+        self.gamma = gamma
+
+    def __call__(self, event_index):
+        return self.initial_value * self.gamma ** (event_index % 9)
+
+
+config4 = (
+    LambdaStateScheduler,
+    {"param_name": "custom_scheduled_param", "lambda_obj": LambdaState(initial_value=10, gamma=0.99)},
+)
+
+
+@pytest.mark.parametrize("scheduler_cls,scheduler_kwargs", [config1, config2, config3, config4])
+def test_simulate_and_plot_values(scheduler_cls, scheduler_kwargs):
+
+    import matplotlib
+
+    matplotlib.use("Agg")
+
+    def _test(scheduler_cls, scheduler_kwargs):
+        event = Events.EPOCH_COMPLETED
+        max_epochs = 2
+        data = [0] * 10
+
+        scheduler = scheduler_cls(**scheduler_kwargs)
+        trainer = Engine(lambda engine, batch: None)
+        scheduler.attach(trainer, event)
+        trainer.run(data, max_epochs=max_epochs)
+
+        # launch plot values
+        scheduler_cls.plot_values(num_events=len(data) * max_epochs, **scheduler_kwargs)
+
+    _test(scheduler_cls, scheduler_kwargs)
+
+
+@pytest.mark.parametrize("scheduler_cls,scheduler_kwargs", [config1, config2, config3, config4])
+def test_state_param_asserts(scheduler_cls, scheduler_kwargs):
+    import re
+
+    def _test(scheduler_cls, scheduler_kwargs):
+        scheduler = scheduler_cls(**scheduler_kwargs)
+        with pytest.raises(
+            ValueError,
+            match=r"Attribute: '"
+            + re.escape(scheduler_kwargs["param_name"])
+            + "' is already defined in the Engine.state.This may be a conflict between multiple StateParameterScheduler"
+            + " handlers.Please choose another name.",
+        ):
+
+            trainer = Engine(lambda engine, batch: None)
+            event = Events.EPOCH_COMPLETED
+            max_epochs = 2
+            data = [0] * 10
+            scheduler.attach(trainer, event)
+            trainer.run(data, max_epochs=max_epochs)
+            scheduler.attach(trainer, event)
+
+    _test(scheduler_cls, scheduler_kwargs)
+
+
+def test_torch_save_load():
+
+    lambda_state_parameter_scheduler = LambdaStateScheduler(
+        param_name="custom_scheduled_param", lambda_obj=LambdaState(initial_value=10, gamma=0.99),
+    )
+
+    torch.save(lambda_state_parameter_scheduler, "dummy_lambda_state_parameter_scheduler.pt")
+    loaded_lambda_state_parameter_scheduler = torch.load("dummy_lambda_state_parameter_scheduler.pt")
+
+    engine1 = Engine(lambda e, b: None)
+    lambda_state_parameter_scheduler.attach(engine1, Events.EPOCH_COMPLETED)
+    engine1.run([0] * 8, max_epochs=2)
+    torch.testing.assert_allclose(
+        getattr(engine1.state, "custom_scheduled_param"), LambdaState(initial_value=10, gamma=0.99)(2)
+    )
+
+    engine2 = Engine(lambda e, b: None)
+    loaded_lambda_state_parameter_scheduler.attach(engine2, Events.EPOCH_COMPLETED)
+    engine2.run([0] * 8, max_epochs=2)
+    torch.testing.assert_allclose(
+        getattr(engine2.state, "custom_scheduled_param"), LambdaState(initial_value=10, gamma=0.99)(2)
+    )
+
+    torch.testing.assert_allclose(
+        getattr(engine1.state, "custom_scheduled_param"), getattr(engine2.state, "custom_scheduled_param")
+    )
+
+
+def test_simulate_and_plot_values_no_matplotlib():
+    def _test(scheduler_cls, **scheduler_kwargs):
+        event = Events.EPOCH_COMPLETED
+        max_epochs = 2
+        data = [0] * 10
+
+        scheduler = scheduler_cls(**scheduler_kwargs)
+        trainer = Engine(lambda engine, batch: None)
+        scheduler.attach(trainer, event)
+        trainer.run(data, max_epochs=max_epochs)
+
+        # launch plot values
+        scheduler_cls.plot_values(num_events=len(data) * max_epochs, **scheduler_kwargs)
+
+    with pytest.raises(RuntimeError, match=r"This method requires matplotlib to be installed."):
+        with patch.dict("sys.modules", {"matplotlib.pyplot": None}):
+            _test(
+                MultiStepStateScheduler,
+                param_name="multistep_scheduled_param",
+                initial_value=10,
+                gamma=0.99,
+                milestones=[3, 6],
+            )
+
+
+def test_docstring_examples():
+    # LambdaStateScheduler
+
+    engine = Engine(lambda e, b: None)
+
+    class LambdaState:
+        def __init__(self, initial_value, gamma):
+            self.initial_value = initial_value
+            self.gamma = gamma
+
+        def __call__(self, event_index):
+            return self.initial_value * self.gamma ** (event_index % 9)
+
+    param_scheduler = LambdaStateScheduler(param_name="param", lambda_obj=LambdaState(10, 0.99),)
+
+    param_scheduler.attach(engine, Events.EPOCH_COMPLETED)
+
+    engine.run([0] * 8, max_epochs=2)
+
+    # PiecewiseLinearStateScheduler
+
+    engine = Engine(lambda e, b: None)
+
+    param_scheduler = PiecewiseLinearStateScheduler(
+        param_name="param", milestones_values=[(10, 0.5), (20, 0.45), (21, 0.3), (30, 0.1), (40, 0.1)]
+    )
+
+    param_scheduler.attach(engine, Events.EPOCH_COMPLETED)
+
+    engine.run([0] * 8, max_epochs=40)
+
+    # ExpStateScheduler
+
+    engine = Engine(lambda e, b: None)
+
+    param_scheduler = ExpStateScheduler(param_name="param", initial_value=10, gamma=0.99)
+
+    param_scheduler.attach(engine, Events.EPOCH_COMPLETED)
+
+    engine.run([0] * 8, max_epochs=2)
+
+    # StepStateScheduler
+
+    engine = Engine(lambda e, b: None)
+
+    param_scheduler = StepStateScheduler(param_name="param", initial_value=10, gamma=0.99, step_size=5)
+
+    param_scheduler.attach(engine, Events.EPOCH_COMPLETED)
+
+    engine.run([0] * 8, max_epochs=10)
+
+    # MultiStepStateScheduler
+
+    engine = Engine(lambda e, b: None)
+
+    param_scheduler = MultiStepStateScheduler(param_name="param", initial_value=10, gamma=0.99, milestones=[3, 6],)
+
+    param_scheduler.attach(engine, Events.EPOCH_COMPLETED)
+
+    engine.run([0] * 8, max_epochs=10)