Repository: pytorch/ignite

Issue:
EMAHandler with warm-up counter-intuitive.
## ❓ Questions/Help/Support

<!-- If you have a question, please checkout other issues labelled as "question" (https://github.com/pytorch/ignite/issues?q=is%3Aissue+is%3Aopen+label%3Aquestion) -->
<!-- If your question is not listed above or you need a support or help, please provide a description-->
<!-- If you prefer, you can also reach out us on [Discussion Forum](https://discuss.pytorch.org/) adding tag "ignite". -->

Hello there,
First of all, thanks for the cool package, it's just super useful.

I recently checked out the new version of ignite, which provides a cool new feature: [EMAHandler](https://pytorch.org/ignite/_modules/ignite/handlers/ema_handler.html#EMAHandler).
The EMAHandler uses a momentum value (\alpha), which in contrast to most implementations that I've seen is the inverse decay value (i.e., 1- \alpha = decay).
Usually, when training an EMA model, the online model parameters are unreliable in the first couple epochs, thus in these cases, one might use a high \alpha value of i.e., 0.999, which leads to mainly copying the original parameters over to the EMA model.
Then after some epochs, the EMA models parameters are more reliable and thus \alpha can be set to i.e., 0.001 in order to incrementally update the model.

Thus, in most implementations, there is a linear warmup for the momentum value from i.e., 0.0 to the target decay value (i.e., 0.999).
However, in the current EMAHandler, the warmup momentum (\alpha) is required to be [strictly less than the target momentum](https://github.com/pytorch/ignite/blob/4a37e35dcb9548bb777dc86ab649ff390c6fb631/ignite/handlers/ema_handler.py#L134).
This, therefore leads to only two choices for the momentum: 

1. Start with a high \alpha value and a low (i.e., 0.0) warmup momentum, which just means that at the beginning \theta_{EMA} is used and late into training the online parameters \theta_t.
2. Don't use momentum warmup at all.

In the current EMAHandler implementation, it is impossible to start with a low \alpha (0.0) and increase that value to a high decay (0.999).

My question is, is there something I missed here or why is the implementation using warmup values strictly smaller than the target momentum?




Discussion:
Thanks for your feedback and the question @RicherMans !

I put into the copy @sandylaker for visibility as it was him who provided the PR with EMAHandler.
I'll check in details our implementation and your question and we'll get back to you asap.
@RicherMans @vfdev-5 Thank you for your feedback.

1. About the `momentum` or inverse `momentum`:

In PyTorch, momentum is preferred over the inverse one, i.e., `ema_next = model_curr * momentum + ema_curr * (1 - momentum)`. Please also check `BatchNorm2d` in PyTorch docs.

2. About the `momentum_warmup`:

This parameter denotes the initial momentum during the warmup. I.e., the momentum when `iter = 1`. Users can disable the warmup either by setting `wamup_iters=0` or `momentum_warmup = momentum`. 

3. About why `momentum_warmup` must be less than or equal to `momentum`:

My understanding is that the training is unstable in the beginning (as you also pointed out), so small updates to the model are desired. Thus, setting momentum to a small value in the beginning has similar effect as setting a small learning rate. Some implementations also adopt the same warmup strategy, e.g. [EMAHook in mmcv](https://github.com/open-mmlab/mmcv/blob/f9be3bc3336f00138b876de002477864810c5b29/mmcv/runner/hooks/ema.py#L64). However, I also wonder how the other implementations handle this wamup issue. Could you please post some examples of them? Thanks a lot.
Hey @sandylaker,
thanks for the quick response.
So my problem is actually not how the momentum is defined, 1- alpha = decay, therefore one just needs to pick one of these variables, that's alright.

> My understanding is that the training is unstable in the beginning (as you also pointed out), so small updates to the model are desired. Thus, setting momentum to a small value, in the beginning, has similar effect as setting a small learning rate. Some implementations also adopt the same warmup strategy, e.g. EMAHook in mmcv. However, I also wonder how the other implementations handle this wamup issue. Could you please post some examples of them? Thanks a lot.

I agree, so I'd like to just have this behaviour: EMA starts with a value that lets the model nearly copy the online parameters and then later goes to mainly update the EMA parameters.

For an example of some implementation of EMA, I have the original paper from CuriousAI, linked with the code [here](https://github.com/CuriousAI/mean-teacher/blob/546348ff863c998c26be4339021425df973b4a36/pytorch/main.py#L189).

The main code is :

```python
def update_ema_variables(model, ema_model, alpha, global_step):
    # Use the true average until the exponential average is more correct
    alpha = min(1 - 1 / (global_step + 1), alpha)
    for ema_param, param in zip(ema_model.parameters(), model.parameters()):
        ema_param.data.mul_(alpha).add_(1 - alpha, param.data)
```

where the parameters alpha = 0.999 and the global step equals the current iteration (cur_iter in ignite's implementation).

To illustrate my point, I just made a mock example:

```python
MOMENTUM = 0.001
DECAY=1.- MOMENTUM

def update_ema(curr_iter, decay):
    return min(1 - 1 / (curr_iter + 1), decay)

for i in range(1, 10):
    print(update_ema(i, DECAY))
```

which prints:


```bash
0.5
0.6666666666666667
0.75
0.8
0.8333333333333334
0.8571428571428572
0.875
0.8888888888888888
0.9
```

which is the expected behaviour ( 0.0 -> 0.999) .

However, to my understanding the current EMAHandler implementation does not support this type of training, since the warmup momentum needs to be smaller than the momentum.
To illustrate here is a mock example of the core code to get the momentum for some iteration in EMAHandler:

```python
from mock import Mock

someobj = Mock()

someobj.warmup_iters = 10
someobj.momentum = 0.001
someobj.momentum_warmup = 0.0


def get_momentum(curr_iter: int) -> float:
    """Get current momentum, `curr_iter` should be 1-based. When `curr_iter = 1`, `momentum =
        self.momentum_warmup`; when `curr_iter >= self.warmup_iters`, `momentum = self.momentum`"""

    if curr_iter < 1:
        raise ValueError(
            f"curr_iter should be at least 1, but got {curr_iter}.")

    denominator = max(1, someobj.warmup_iters - 1)
    momentum = someobj.momentum_warmup + (
        someobj.momentum - someobj.momentum_warmup) * (curr_iter -
                                                       1) / denominator
    return min(someobj.momentum, momentum)



for i in range(1, 10):
    print(get_momentum(i))
```

which yields:

```bash
0.0
0.00011111111111111112
0.00022222222222222223
0.0003333333333333333
0.00044444444444444447
0.0005555555555555556
0.0006666666666666666
0.0007777777777777778
0.0008888888888888889
```

Note that these values are the inverse of these above, thus they go from 1.0 to 0.999.

Therefore to my understanding the current implementation might have a problem or I didn't quite understand how to use the warmup value correctly.
As far as I understand we would just need to fix the warmup value, i.e. that it needs to be strictly greater than warmup and goes down to the `momentum` after some steps.

Please correct me if I'm mistaken.
@RicherMans thanks a lot for your analysis! In any case user can override internal warm-up behaviour as following:
```python
import ignite
print(ignite.__version__)

from ignite.engine import Engine, Events
from ignite.handlers import PiecewiseLinearStateScheduler, EMAHandler

import torch
import torch.nn as nn

import matplotlib.pyplot as plt
%matplotlib inline


model = nn.Linear(2, 1)

ema_decay_history = []

trainer = Engine(lambda e, b: ema_decay_history.append(1.0 - e.state.ema_momentum))

data = torch.rand(100, 2)

ema_handler = EMAHandler(model)
ema_handler.attach(trainer, name="ema_momentum", event=Events.ITERATION_COMPLETED)

MOMENTUM = 0.001
DECAY = 1.0 - MOMENTUM

@trainer.on(Events.ITERATION_STARTED)
def update_ema_momentum():
    curr_iter = trainer.state.iteration
    trainer.state.ema_momentum = 1.0 - min(1.0 - 1.0 / (curr_iter + 1), DECAY)


trainer.run(data, max_epochs=20)

print(ema_decay_history[:5], ema_decay_history[-5:])
> ([0.5, 0.6666666666666667, 0.75, 0.8, 0.8333333333333334],
    [0.999, 0.999, 0.999, 0.999, 0.999])

plt.figure(figsize=(10, 7))
plt.plot(ema_decay_history[:100])
```
![image](https://user-images.githubusercontent.com/2459423/139044576-25e6f16d-19dc-48f6-a97d-8470edd0d4a1.png)



But I agree that it would be better to clarify how EMAHandler's warm-up is intended to work.

EDIT: yes, EMAHandler's internal warm-up code is not intuitive for me as well and agree even a bit buggy if we need to get decay (1.0 - momentum) for go from 0.0 to a value...
Hey @sandylaker ,
thanks a lot for the fix!
However, what do you think is the current behaviour intentional?
I just mean that with the current code only allows for either:

1. Start with a low \alpha via momentum_warmup and set a high momentum value. This is I think impractical, since the model would start to first consider the unstable EMA average and then proceeds to use the online weights. At least in all my experiments, Losses and Validation metrics quickly become unstable.
2. Do not use a warmup, which likely leads to a slow convergence at the very least.

So should we change the internal code or only add this fix of yours to the docs ?
Kind regards
@RicherMans i think firstly we can try to relax conditions on the warm-up coded behaviour which is too strict IMO. We can also update the docs showing how to do any type of scheduling. We have also state param scheduler now in the codebase (which unfortunately is not yet working together with EMAHandler, https://github.com/pytorch/ignite/issues/2295) but the idea is that we could do any kind of scheduling.
@vfdev-5 
I see, that would be a much more reasonable option. Thanks guys for the code and the comments!
@vfdev-5 @RicherMans Thank you for your explanation.

Shall we change the API such that users are allowed to provide a custom function to update the momentum?  That is 
```python3
EMAHandler(model: nn.Module, momentum: Union[float, Callable])
```
1. If `momentum` is a constant float, then we use a constant momentum during the training.
2. If a momentum is a Callable that takes `curr_iter` and returns `curr_momentum`, then we use this Callable to update the momentum for training.

This design simplifies the API and allows for non-warm-up and custom warm-up configurations. 

I am willing to know your opinions.
@sandylaker 
That would actually make the entire API at the current state much more transparent and arguably slightly easier to use.
I agree to that change.
@sandylaker thanks for suggestions!
I agree that we can deprecate `momentum_warmup` and `warmup_iters` and keep only `momentum` as constant float.
I propose to enable momentum scheduling using state param scheduler and fixing the issue #2295 
Passing `momentum` as callable can be a bit more tricky for the users to figure out its signature and what it should return (it is less obvious without reading docs).

What do you think ?
@vfdev-5 I agree with you. We can provide a suggestion in the docs and suggest the users to utilize the state param scheduler to update the momentum. I think it also reduces the tests complexity of the EMAHandler.

I will send a PR in the next week after I am back from travel. Thank you guys again for raising this issue and also for the discussion.
@vfdev-5 There might be a conflict when using StateParamScheduler to schedule the momentum.

In the current implementation of EMAHandler, we create a state attribute `momentum` when we initialize and attach the handler to engine. But later when we initialize a StateParamScheduler to schedule the `momentum`, it will raise an error due to the fact that `momentum` is already defined in the engine's state.

Please check out [here](https://github.com/pytorch/ignite/blob/79678b916e1fbef702f861d71d5832e17c188351/ignite/handlers/state_param_scheduler.py#L45)

Is there any proper way to overcome this problem?
Yes, i opened an issue for that: #2295 
We have to figure out what would be the most reasonable way to fix it...

Resolving patch:
diff --git a/ignite/handlers/ema_handler.py b/ignite/handlers/ema_handler.py
--- a/ignite/handlers/ema_handler.py
+++ b/ignite/handlers/ema_handler.py
@@ -1,13 +1,31 @@
+import warnings
 from copy import deepcopy
 from typing import Optional, Union
 
 import torch.nn as nn
 
 from ignite.engine import CallableEventWithFilter, Engine, Events, EventsList
+from ignite.handlers.param_scheduler import BaseParamScheduler
+from ignite.handlers.state_param_scheduler import LambdaStateScheduler
 
 __all__ = ["EMAHandler"]
 
 
+class EMAWarmUp:
+    def __init__(self, momentum_warmup: float, warmup_iters: int, momentum: float) -> None:
+        self.momentum_warmup = momentum_warmup
+        self.warmup_iters = warmup_iters
+        self.momentum = momentum
+
+    def __call__(self, event_index: int) -> float:
+        denominator = max(1, self.warmup_iters - 1)
+        curr_momentum = self.momentum_warmup + (self.momentum - self.momentum_warmup) * (event_index - 1) / denominator
+        if self.momentum >= self.momentum_warmup:
+            return min(self.momentum, curr_momentum)
+        else:
+            return max(self.momentum, curr_momentum)
+
+
 class EMAHandler:
     r"""Exponential moving average (EMA) handler can be used to compute a smoothed version of model.
     The EMA model is updated as follows:
@@ -15,26 +33,21 @@ class EMAHandler:
     .. math:: \theta_{\text{EMA}, t+1} = (1 - \lambda) \cdot \theta_{\text{EMA}, t} + \lambda \cdot \theta_{t}
 
     where :math:`\theta_{\text{EMA}, t}` and :math:`\theta_{t}` are the EMA weights and online model weights at
-    :math:`t`-th iteration, respectively; :math:`\lambda` is the update momentum. The handler allows for linearly
-    warming up the momentum in the beginning when training process is not stable. Current momentum can be retrieved
+    :math:`t`-th iteration, respectively; :math:`\lambda` is the update momentum. Current momentum can be retrieved
     from ``Engine.state.ema_momentum``.
 
     Args:
           model: the online model for which an EMA model will be computed. If ``model`` is ``DataParallel`` or
               ``DistributedDataParallel``, the EMA smoothing will be applied to ``model.module`` .
           momentum: the update momentum after warmup phase, should be float in range :math:`\left(0, 1 \right)`.
-          momentum_warmup: the initial update momentum during warmup phase, the value should be smaller than
-              ``momentum``. Momentum will linearly increase from this value to ``momentum`` in ``warmup_iters``
-              iterations. If ``None``, no warmup will be performed.
-          warmup_iters: iterations of warmup. If ``None``, no warmup will be performed.
+          momentum_warmup: the initial update momentum during warmup phase.
+          warmup_iters: iterations of warmup.
 
     Attributes:
           ema_model: the exponential moving averaged model.
           model: the online model that is tracked by EMAHandler. It is ``model.module`` if ``model`` in
               the initialization method is an instance of ``DistributedDataParallel``.
-          momentum: the update momentum after warmup phase.
-          momentum_warmup: the initial update momentum.
-          warmup_iters: number of warmup iterations.
+          momentum: the update momentum.
 
     Note:
           The EMA model is already in ``eval`` mode. If model in the arguments is an ``nn.Module`` or
@@ -56,8 +69,7 @@ class EMAHandler:
               device = torch.device("cuda:0")
               model = nn.Linear(2, 1).to(device)
               # update the ema every 5 iterations
-              ema_handler = EMAHandler(
-                  model, momentum=0.0002, momentum_warmup=0.0001, warmup_iters=10000)
+              ema_handler = EMAHandler(model, momentum=0.0002)
               # get the ema model, which is an instance of nn.Module
               ema_model = ema_handler.ema_model
               trainer = Engine(train_step_fn)
@@ -89,6 +101,19 @@ def run_validation(engine):
 
               trainer.run(...)
 
+          The following example shows how to perform warm-up to the EMA momentum:
+
+          .. code-block:: python
+
+              device = torch.device("cuda:0")
+              model = nn.Linear(2, 1).to(device)
+              # linearly change the EMA momentum from 0.2 to 0.002 in the first 100 iterations,
+              # then keep a constant EMA momentum of 0.002 afterwards
+              ema_handler = EMAHandler(model, momentum=0.002, momentum_warmup=0.2, warmup_iters=100)
+              engine = Engine(step_fn)
+              ema_handler.attach(engine, name="ema_momentum")
+              engine.run(...)
+
           The following example shows how to attach two handlers to the same trainer:
 
           .. code-block:: python
@@ -125,25 +150,19 @@ def __init__(
         momentum_warmup: Optional[float] = None,
         warmup_iters: Optional[int] = None,
     ) -> None:
-        if momentum_warmup is not None and not 0 < momentum_warmup < 1:
-            raise ValueError(f"Invalid momentum_warmup: {momentum_warmup}")
         if not 0 < momentum < 1:
             raise ValueError(f"Invalid momentum: {momentum}")
-        if momentum_warmup is not None and not momentum_warmup <= momentum:
-            raise ValueError(
-                f"momentum_warmup should be less than or equal to momentum, but got "
-                f"momentum_warmup: {momentum_warmup} and momentum: {momentum}"
-            )
-        if warmup_iters is not None and not (isinstance(warmup_iters, int) and warmup_iters > 0):
-            raise ValueError(f"Invalid warmup_iters: {warmup_iters}")
+        self.momentum = momentum
+        self._momentum_lambda_obj: Optional[EMAWarmUp] = None
+        if momentum_warmup is not None and warmup_iters is not None:
+            self.momentum_scheduler: Optional[BaseParamScheduler] = None
+            self._momentum_lambda_obj = EMAWarmUp(momentum_warmup, warmup_iters, momentum)
+
         if not isinstance(model, nn.Module):
             raise ValueError(
                 f"model should be an instance of nn.Module or its subclasses, but got"
                 f"model: {model.__class__.__name__}"
             )
-        self.momentum_warmup = momentum_warmup
-        self.momentum = momentum
-        self.warmup_iters = warmup_iters
 
         if isinstance(model, nn.parallel.DistributedDataParallel):
             model = model.module
@@ -154,22 +173,6 @@ def __init__(
             param.detach_()
         self.ema_model.eval()
 
-    def _get_momentum(self, curr_iter: int) -> float:
-        """Get current momentum, `curr_iter` should be 1-based. When `curr_iter = 1`, `momentum =
-        self.momentum_warmup`; when `curr_iter >= self.warmup_iters`, `momentum = self.momentum`"""
-
-        # TODO: use ignite's parameter scheduling, see also GitHub issue #2090
-        if curr_iter < 1:
-            raise ValueError(f"curr_iter should be at least 1, but got {curr_iter}.")
-
-        # no warmup
-        if self.momentum_warmup is None or self.warmup_iters is None:
-            return self.momentum
-
-        denominator = max(1, self.warmup_iters - 1)
-        momentum = self.momentum_warmup + (self.momentum - self.momentum_warmup) * (curr_iter - 1) / denominator
-        return min(self.momentum, momentum)
-
     def _update_ema_model(self, engine: Engine, name: str) -> None:
         """Update weights of ema model"""
         momentum = getattr(engine.state, name)
@@ -179,36 +182,47 @@ def _update_ema_model(self, engine: Engine, name: str) -> None:
         for ema_b, model_b in zip(self.ema_model.buffers(), self.model.buffers()):
             ema_b.data = model_b.data
 
-    def _update_ema_momentum(self, engine: Engine, name: str) -> None:
-        """Update momentum in engine.state"""
-        curr_iter = engine.state.iteration
-        momentum = self._get_momentum(curr_iter)
-        setattr(engine.state, name, momentum)
-
     def attach(
         self,
         engine: Engine,
         name: str = "ema_momentum",
+        warn_if_exists: bool = True,
         event: Union[str, Events, CallableEventWithFilter, EventsList] = Events.ITERATION_COMPLETED,
     ) -> None:
         """Attach the handler to engine. After the handler is attached, the ``Engine.state`` will add an new attribute
-        with name ``name``. Then, current momentum can be retrieved by from ``Engine.state`` when the engine runs.
+        with name ``name`` if the attribute does not exist. Then, the current momentum can be retrieved from
+        ``Engine.state`` when the engine runs.
+
+
+        Note:
+            There are two cases where a momentum with name ``name`` already exists: 1. the engine has loaded its
+            state dict after resuming. In this case, there is no need to initialize the momentum again, and users
+            can set ``warn_if_exists`` to False to suppress the warning message; 2. another handler has created
+            a state attribute with the same name. In this case, users should choose another name for the ema momentum.
+
 
         Args:
             engine: trainer to which the handler will be attached.
             name: attribute name for retrieving EMA momentum from ``Engine.state``. It should be a unique name since a
                 trainer can have multiple EMA handlers.
+            warn_if_exists: if True, a warning will be thrown if the momentum with name ``name`` already exists.
             event: event when the EMA momentum and EMA model are updated.
 
         """
         if hasattr(engine.state, name):
-            raise ValueError(
-                f"Attribute: '{name}' is already in Engine.state. Thus it might be "
-                f"overridden by other EMA handlers. Please select another name."
-            )
-
-        setattr(engine.state, name, 0.0)
-
-        # first update momentum, then update ema model
-        engine.add_event_handler(event, self._update_ema_momentum, name)
+            if warn_if_exists:
+                warnings.warn(
+                    f"Attribute '{name}' already exists in Engine.state. It might because 1. the engine has loaded its "
+                    f"state dict or 2. {name} is already created by other handlers. Turn off this warning by setting"
+                    f"warn_if_exists to False.",
+                    category=UserWarning,
+                )
+        else:
+            setattr(engine.state, name, self.momentum)
+
+        if self._momentum_lambda_obj is not None:
+            self.momentum_scheduler = LambdaStateScheduler(self._momentum_lambda_obj, param_name="ema_momentum")
+
+            # first update the momentum and then update the EMA model
+            self.momentum_scheduler.attach(engine, event)
         engine.add_event_handler(event, self._update_ema_model, name)
diff --git a/ignite/handlers/state_param_scheduler.py b/ignite/handlers/state_param_scheduler.py
--- a/ignite/handlers/state_param_scheduler.py
+++ b/ignite/handlers/state_param_scheduler.py
@@ -4,7 +4,7 @@
 from typing import Any, List, Sequence, Tuple, Union
 
 from ignite.engine import CallableEventWithFilter, Engine, Events, EventsList
-from ignite.handlers import BaseParamScheduler
+from ignite.handlers.param_scheduler import BaseParamScheduler
 
 
 class StateParamScheduler(BaseParamScheduler):


Test patch:
diff --git a/tests/ignite/handlers/test_ema_handler.py b/tests/ignite/handlers/test_ema_handler.py
--- a/tests/ignite/handlers/test_ema_handler.py
+++ b/tests/ignite/handlers/test_ema_handler.py
@@ -47,18 +47,51 @@ def test_ema_invalid_momentum(get_dummy_model, momentum):
         EMAHandler(get_dummy_model(), momentum=momentum)
 
 
-@pytest.mark.parametrize("momentum_warmup", [-1, 2])
-def test_ema_invalid_momentum_warmup(get_dummy_model, momentum_warmup):
-    with pytest.raises(ValueError, match="Invalid momentum_warmup"):
-        EMAHandler(get_dummy_model, momentum_warmup=momentum_warmup)
-
+def test_has_momentum_scheduler(get_dummy_model):
+    """Test the handler has attribute `momentum_scheduler` and `_momentum_lambda_obj`"""
+    momentum_warmup = 0.0
+    warmup_iters = 10
+    ema_handler = EMAHandler(get_dummy_model(), momentum_warmup=momentum_warmup, warmup_iters=warmup_iters)
+    assert hasattr(ema_handler, "momentum_scheduler")
+    assert hasattr(ema_handler, "_momentum_lambda_obj")
+
+
+def test_ema_warmup_func(get_dummy_model):
+    """Test the built-in linear warmup function for the EMA momentum"""
+    momentum = 0.5
+    momentum_warmup_1 = 0.0
+    momentum_warmup_2 = 1.0
+    warmup_iters = 5
+
+    def check_ema_momentum(engine: Engine, momentum_warmup, final_momentum, warmup_iters):
+        if engine.state.iteration == 1:
+            assert engine.state.ema_momentum == momentum_warmup
+        elif engine.state.iteration >= 1 + warmup_iters:
+            assert engine.state.ema_momentum == final_momentum
+        else:
+            min_momentum = min(momentum, momentum_warmup)
+            max_momentum = max(momentum, momentum_warmup)
+            assert min_momentum <= engine.state.ema_momentum <= max_momentum
 
-def test_ema_invalid_momentum_start_end(get_dummy_model):
-    """Test momentum_end > momentum_start"""
-    momentum = 0.001
-    momentum_warmup = 0.1
-    with pytest.raises(ValueError, match="momentum_warmup should be less than or equal to momentum"):
-        EMAHandler(get_dummy_model(), momentum_warmup=momentum_warmup, momentum=momentum)
+    # momentum_warmup < momentum
+    model_1 = get_dummy_model()
+    engine_1 = Engine(_get_dummy_step_fn(model_1))
+    ema_handler_1 = EMAHandler(model_1, momentum, momentum_warmup_1, warmup_iters)
+    ema_handler_1.attach(engine_1)
+    engine_1.add_event_handler(
+        Events.ITERATION_COMPLETED, check_ema_momentum, momentum_warmup_1, momentum, warmup_iters
+    )
+    engine_1.run(range(10))
+
+    # momentum_warmup > momentum
+    model_2 = get_dummy_model()
+    engine_2 = Engine(_get_dummy_step_fn(model_2))
+    ema_handler_2 = EMAHandler(model_2, momentum, momentum_warmup_2, warmup_iters)
+    ema_handler_2.attach(engine_2)
+    engine_2.add_event_handler(
+        Events.ITERATION_COMPLETED, check_ema_momentum, momentum_warmup_2, momentum, warmup_iters
+    )
+    engine_2.run(range(10))
 
 
 def test_ema_invalid_model():
@@ -98,7 +131,8 @@ def test_ema_load_state_dict(get_dummy_model):
     assert ema_model.weight.data.allclose(model_1.weight.data)
 
 
-def test_ema_no_warmup_momentum(get_dummy_model):
+def test_ema_get_const_momentum(get_dummy_model):
+    """Test if momentum retrieved from the engine is constant and equal to the handler's momentum"""
     model = get_dummy_model()
     step_fn = _get_dummy_step_fn(model)
     engine = Engine(step_fn)
@@ -106,46 +140,10 @@ def test_ema_no_warmup_momentum(get_dummy_model):
     def assert_const_momentum(engine: Engine, const_momentum):
         assert engine.state.ema_momentum == const_momentum
 
-    # no momentum_warmup
-    ema_handler = EMAHandler(model, momentum=0.002, momentum_warmup=None, warmup_iters=1)
-    ema_handler.attach(engine)
-    # attach the assertion handler after ema_handler, so the momentum is first updated and then tested
-    engine.add_event_handler(Events.ITERATION_COMPLETED, assert_const_momentum, ema_handler.momentum)
-    engine.run(range(2))
-
-    # no warmup_iters
-    engine = Engine(step_fn)
-    ema_handler = EMAHandler(model, momentum=0.002, momentum_warmup=0.001, warmup_iters=None)
+    ema_handler = EMAHandler(model, momentum=0.002)
     ema_handler.attach(engine)
-    # attach the assertion handler after ema_handler, so the momentum is first updated and then tested
     engine.add_event_handler(Events.ITERATION_COMPLETED, assert_const_momentum, ema_handler.momentum)
-    engine.run(range(2))
-
-
-def test_ema_update_ema_momentum(get_dummy_model):
-    model = get_dummy_model()
-    step_fn = _get_dummy_step_fn(model)
-    engine = Engine(step_fn)
-
-    warmup_iters = 4
-    momentum_warmup = 0.1
-    momentum = 0.2
-    ema_handler = EMAHandler(model, momentum_warmup=momentum_warmup, momentum=momentum, warmup_iters=warmup_iters)
-    ema_handler.attach(engine)
-
-    # add handlers to check momentum at each iteration
-    @engine.on(Events.ITERATION_COMPLETED)
-    def assert_momentum(engine: Engine):
-        curr_iter = engine.state.iteration
-        curr_momentum = engine.state.ema_momentum
-        if curr_iter == 1:
-            assert curr_momentum == momentum_warmup
-        elif 1 < curr_iter < warmup_iters:
-            assert momentum_warmup < curr_momentum < momentum
-        else:
-            assert curr_momentum == momentum
-
-    engine.run(range(2), max_epochs=5)
+    engine.run(range(10))
 
 
 def test_ema_buffer():
@@ -180,11 +178,10 @@ def check_buffers():
 def test_ema_two_handlers(get_dummy_model):
     """Test when two EMA handlers are attached to a trainer"""
     model_1 = get_dummy_model()
-    # momentum will be constantly 0.5
-    ema_handler_1 = EMAHandler(model_1, momentum_warmup=0.5, momentum=0.5, warmup_iters=1)
+    ema_handler_1 = EMAHandler(model_1, momentum=0.5)
 
     model_2 = get_dummy_model()
-    ema_handler_2 = EMAHandler(model_2, momentum_warmup=0.5, momentum=0.5, warmup_iters=1)
+    ema_handler_2 = EMAHandler(model_2, momentum=0.5)
 
     def _step_fn(engine: Engine, batch: Any):
         model_1.weight.data.add_(1)
@@ -214,8 +211,8 @@ def _step_fn(engine: Engine, batch: Any):
 
     model_3 = get_dummy_model()
     ema_handler_3 = EMAHandler(model_3)
-    with pytest.raises(ValueError, match="Please select another name"):
-        ema_handler_3.attach(engine, "ema_momentum_2")
+    with pytest.warns(UserWarning, match="Attribute 'ema_momentum_1' already exists"):
+        ema_handler_3.attach(engine, name="ema_momentum_1")
 
 
 def _test_ema_final_weight(model, device=None, ddp=False, interval=1):
@@ -231,8 +228,7 @@ def _test_ema_final_weight(model, device=None, ddp=False, interval=1):
     step_fn = _get_dummy_step_fn(model)
     engine = Engine(step_fn)
 
-    # momentum will be constantly 0.5
-    ema_handler = EMAHandler(model, momentum_warmup=0.5, momentum=0.5, warmup_iters=1)
+    ema_handler = EMAHandler(model, momentum=0.5)
     ema_handler.attach(engine, "model", event=Events.ITERATION_COMPLETED(every=interval))
 
     # engine will run 4 iterations