Repository: pytorch/ignite

Issue:
EMAHandler (Buffer update) does lead to bad results for Mean-Teacher experiments
## 🚀 Feature

Hey there,
thanks for the increased effort in Ignite, the package just gets better by each update!

So my feature request ( or bug report ) is concerned with the recently developed [EMAHAndler](https://pytorch.org/ignite/_modules/ignite/handlers/ema_handler.html#EMAHandler).

Specifically I recently ran a lot of experiments using Ignite for Sound event classification and noticed that my results far underperform ( F1 Score of ~40 with Ignite and ~60 without ) models that do not use the included EMAHandler class.
I investigated this issue and came to the conclusion that the implementation of Ignites EMAHandler has some minor issues which lead to this behaviour.

Specifically, the lines responsible are:

https://github.com/pytorch/ignite/blob/80c6bb4c9b7b8bacfa5f39557c16065267a849ce/ignite/handlers/ema_handler.py#L174

and

https://github.com/pytorch/ignite/blob/80c6bb4c9b7b8bacfa5f39557c16065267a849ce/ignite/handlers/ema_handler.py#L182

## First problem, evaluation mode

First, I'd like to discuss the problem with ```self.ema_model.eval()```.
While this default behavior makes sense for doing EMA updates, it leads to some major problems regarding non-parameter (buffer) updates of the EMA model.
A particular example is the use of BatchNorm, which will not update its mean and variance at all during EMA training.
However BatchNorm **will** update its affine transformations during training, which might lead to a discrepancy between the EMA model and the original (student) model.

## Second problem, updating buffers

The second problem is directly connected to the first one.
Since the non-parameter updates of the EMA model are in-fact frozen in the Ignite implementation, the `buffer` variables are therefore simply copied from the Student (Original model) to the EMAModel using:

```python
        for ema_b, model_b in zip(self.ema_model.buffers(), self.model.buffers()):
            ema_b.data = model_b.data
```

While again this makes perfect sense, for variables such as `num_tracked_batches` in a BatchNorm layer, it does not make sense for their corresponding means and variances.
For example, after some EMA updates, the weights of the EMA model might not be similar to the ones of the student model. But the current implementation **forces** the EMA model to use the exact same mean's and variances from the student model, which might not be the "True" means and variances.
This then will likely lead to a problem when training for longer, since the EMA model cannot properly normalize the input data and thus somewhat diverges.


## Solutions

I suggest two solutions:

1. (Not recommended ) Also average floating point buffers, such as means and variances.

```python
for ema_b, model_b in zip(self.ema_model.buffers(),                                                                                                                                  
                          self.model.buffers()):                                                                                                                                     
    #Update buffers like running mean and variance                                                                                                                          
    if torch.is_floating_point(ema_b):                                                                                                                                      
        ema_b.mul_(1.0 - momentum).add_(model_b.data, alpha=momentum)                                                                                                       
    else:                                                                                                                                                                   
        ema_b.data = model_b.data
```

I tried this solution already in my experiments and at least from a "loss" perspective, it does outperform the standard EMA model by a large margin.
However the final results were still worse than the ones proposed in 2.:

2. Do not force `.eval()` mode and remove the buffer updates.

I'd recommend currently to follow the implementation of most mean-teacher approaches, which actually require the EMA model to be in `train()` mode, thus updating the batchnorm buffers as well as adding some noise (dropout) during training.
Thus, I'd recommend just removing the forced `self.ema_model.eval()` statement and removing the buffer updates.

What is your point of view about this issue? Did anybody maybe also encounter a similar problem?

.


Discussion:
@RicherMans Thanks a lot for detailed report and the feedback ! Sorry for spoiling lots of your experiences with our probably buggy code.

As for `self.ema_model.eval()`, a workaround rightnow is to do `ema_handler.ema_model.train()` right after creating the handler.

@sandylaker do you have any comments about buffers in EMAHandler ?

From my side, about 2 years ago with fix-match I saw that student buffers were assigned to EMA model in TF code.

In torchvision, here is how EMA is implemented:
https://github.com/pytorch/vision/blob/c890a7e75ebeaaa75ae9ace4c203b7fc145df068/references/classification/train.py#L315

https://github.com/pytorch/vision/blob/c890a7e75ebeaaa75ae9ace4c203b7fc145df068/references/classification/train.py#L47-L51

https://github.com/pytorch/vision/blob/c890a7e75ebeaaa75ae9ace4c203b7fc145df068/references/classification/utils.py#L159

They use `use_buffers=True` which corresponds to "it will compute running averages for both the parameters and the buffers of the model":

https://github.com/pytorch/pytorch/blob/master/torch/optim/swa_utils.py

Maybe we could either add this option as well if it makes sense. Another option is to take a look at https://pytorch.org/docs/stable/optim.html#stochastic-weight-averaging and have an option to use AveragedModel in the handler ? @sandylaker @RicherMans any thoughts ?
@RicherMans Thank you for raising this issue and providing the detailed analysis, and also thank @vfdev-5 for the constructive comments.

1. >As for self.ema_model.eval(), a workaround rightnow is to do ema_handler.ema_model.train() right after creating the handler.

If we replace the `self.ema_model.eval()` with `self.ema_model.train()` in the last line of the initialization method (Line 174), then we have to tell users to manually set `ema_model.eval()` when running validation with the EMA model, and also set `ema_model.train()` after validation finishes. This appears to be a little bit verbose. 

2. >They use use_buffers=True which corresponds to "it will compute running averages for both the parameters and the buffers of the model":

We can also add this option to the `EMAHandler`.

In short, the **second** suggestion sounds better in my opinion. If you agree, I can create a PR to fix this issue.


Hey @sandylaker and @vfdev-5 ,

so I'm also obviously for the second suggestion!

> They use use_buffers=True which corresponds to "it will compute running averages for both the parameters and the buffers of the model":

That sounds great!

> In short, the second suggestion sounds better in my opinion. If you agree, I can create a PR to fix this issue.

@RicherMans @sandylaker looks like we are all agree on how to improve EMAHandler. @sandylaker if you could send a PR it would be great !

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
@@ -42,17 +42,22 @@ class EMAHandler:
           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.
           warmup_iters: iterations of warmup.
+          handle_buffers: how to handle model buffers during training. There are three options: 1. "copy" means
+              copying the buffers of the online model; 2. "update" means applying EMA to the buffers of the online
+              model; 3. "ema_train" means set the EMA model to ``train`` mode and skip copying or updating the buffers.
 
     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.
+          handle_buffers: how to handle model buffers during training.
 
     Note:
-          The EMA model is already in ``eval`` mode. If model in the arguments is an ``nn.Module`` or
-          ``DistributedDataParallel``, the EMA model is an ``nn.Module`` and it is on the same device as the online
-          model. If the model is an ``nn.DataParallel``, then the EMA model is an ``nn.DataParallel``.
+          The EMA model is already in ``eval`` mode if ``handle_buffers`` is "copy" or "update". If model in the
+          arguments is an ``nn.Module`` or ``DistributedDataParallel``, the EMA model is an ``nn.Module`` and it is on
+          the same device as the online model. If the model is an ``nn.DataParallel``, then the EMA model is an
+          ``nn.DataParallel``.
 
 
     Note:
@@ -149,6 +154,7 @@ def __init__(
         momentum: float = 0.0002,
         momentum_warmup: Optional[float] = None,
         warmup_iters: Optional[int] = None,
+        handle_buffers: str = "copy",
     ) -> None:
         if not 0 < momentum < 1:
             raise ValueError(f"Invalid momentum: {momentum}")
@@ -171,16 +177,39 @@ def __init__(
         self.ema_model = deepcopy(self.model)
         for param in self.ema_model.parameters():
             param.detach_()
-        self.ema_model.eval()
+
+        if handle_buffers not in ("copy", "update", "ema_train"):
+            raise ValueError(
+                f"handle_buffers can only be one of 'copy', 'update', 'ema_train', " f"but got {handle_buffers}"
+            )
+
+        self.handle_buffers = handle_buffers
+        if self.handle_buffers == "ema_train":
+            self.ema_model.train()
+        else:
+            self.ema_model.eval()
 
     def _update_ema_model(self, engine: Engine, name: str) -> None:
         """Update weights of ema model"""
         momentum = getattr(engine.state, name)
         for ema_p, model_p in zip(self.ema_model.parameters(), self.model.parameters()):
             ema_p.mul_(1.0 - momentum).add_(model_p.data, alpha=momentum)
-        # assign the buffers
-        for ema_b, model_b in zip(self.ema_model.buffers(), self.model.buffers()):
-            ema_b.data = model_b.data
+
+        if self.handle_buffers == "update":
+            for ema_b, model_b in zip(self.ema_model.buffers(), self.model.buffers()):
+                try:
+                    ema_b.mul_(1.0 - momentum).add_(model_b.data, alpha=momentum)
+                except RuntimeError:
+                    # Handle the case where ema_b is torch.int64, torch.int32 etc.,
+                    # where a runtime error will be thrown when performing the in-place operations with floats.
+                    # In this case, just copy the data
+                    ema_b.data = model_b.data
+        elif self.handle_buffers == "copy":
+            # assign the buffers
+            for ema_b, model_b in zip(self.ema_model.buffers(), self.model.buffers()):
+                ema_b.data = model_b.data
+        else:
+            pass
 
     def attach(
         self,


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
@@ -146,33 +146,78 @@ def assert_const_momentum(engine: Engine, const_momentum):
     engine.run(range(10))
 
 
-def test_ema_buffer():
-    """Test if the tensors in buffer are also synchronized"""
+@pytest.mark.parametrize("handle_buffers", ["copy", "update", "ema_train", "invalid"])
+def test_ema_buffer(handle_buffers):
+    """Test if the tensors in buffer are also correctly updated"""
     model = nn.BatchNorm2d(2)
     model.running_mean.data.fill_(1.5)
     model.running_var.data.fill_(1.5)
-    ema_handler = EMAHandler(model)
-
-    def _bn_step_fn(engine, batch):
-        x = torch.rand(4, 2, 32, 32)
-        _ = model(x)
-        return 1
 
-    engine = Engine(_bn_step_fn)
-    ema_handler.attach(engine)
+    # manually register a buffer to test if it will be correctly updated
+    model.register_buffer("dummy_buffer", tensor=torch.tensor(1.0, dtype=torch.float32))
 
-    ema_model = ema_handler.ema_model
+    if handle_buffers == "invalid":
+        with pytest.raises(ValueError, match="handle_buffers can only"):
+            _ = EMAHandler(model, momentum=0.5, handle_buffers=handle_buffers)
+    else:
+        ema_handler = EMAHandler(model, momentum=0.5, handle_buffers=handle_buffers)
 
-    @engine.on(Events.ITERATION_COMPLETED)
-    def check_buffers():
-        assert ema_model.running_mean.allclose(model.running_mean)
-        assert ema_model.running_var.allclose(model.running_var)
+        def _bn_step_fn(engine, batch):
+            x = torch.rand(4, 2, 32, 32)
+            _ = model(x)
+            # manually increment the dummy_buffer at every step
+            model.dummy_buffer += 1.0
+            return 1
 
-    # engine will run 4 iterations
-    engine.run([0, 1], max_epochs=2)
+        engine = Engine(_bn_step_fn)
+        ema_handler.attach(engine)
 
-    assert ema_model.running_mean.allclose(model.running_mean)
-    assert ema_model.running_var.allclose(model.running_var)
+        ema_model = ema_handler.ema_model
+        if handle_buffers == "ema_train":
+            assert ema_model.training
+        else:
+            assert not ema_model.training
+
+        @engine.on(Events.ITERATION_COMPLETED)
+        def check_buffers():
+            if handle_buffers == "update":
+                # the buffers with torch.int64 data type should be directly copied
+                assert ema_model.num_batches_tracked.allclose(model.num_batches_tracked)
+
+                # buffers with floating type will be updated rather than copied
+                assert not ema_model.dummy_buffer.allclose(model.dummy_buffer)
+                assert not ema_model.running_mean.allclose(model.running_mean)
+                assert not ema_model.running_var.allclose(model.running_var)
+            elif handle_buffers == "copy":
+                # the buffers with torch.int64 data type should be directly copied
+                assert ema_model.num_batches_tracked.allclose(model.num_batches_tracked)
+
+                assert ema_model.dummy_buffer.allclose(model.dummy_buffer)
+                assert ema_model.running_mean.allclose(model.running_mean)
+                assert ema_model.running_var.allclose(model.running_var)
+            else:
+                # buffers will not be copied or EMA updated
+                assert ema_model.num_batches_tracked.allclose(torch.tensor(0, dtype=torch.int64))
+                assert ema_model.dummy_buffer.allclose(torch.tensor(1.0, dtype=torch.float32))
+
+        # engine will run 4 iterations
+        engine.run([0, 1], max_epochs=2)
+
+        if handle_buffers == "update":
+            assert ema_model.num_batches_tracked.allclose(model.num_batches_tracked)
+            assert ema_model.dummy_buffer.allclose(torch.tensor(4.0625, dtype=torch.float32))
+            assert not ema_model.dummy_buffer.allclose(model.dummy_buffer)
+            assert not ema_model.running_mean.allclose(model.running_mean)
+            assert not ema_model.running_var.allclose(model.running_var)
+        elif handle_buffers == "copy":
+            assert ema_model.num_batches_tracked.allclose(model.num_batches_tracked)
+            assert ema_model.dummy_buffer.allclose(model.dummy_buffer)
+            assert ema_model.running_mean.allclose(model.running_mean)
+            assert ema_model.running_var.allclose(model.running_var)
+        else:
+            # buffers will not be copied or EMA updated
+            assert ema_model.num_batches_tracked.allclose(torch.tensor(0, dtype=torch.int64))
+            assert ema_model.dummy_buffer.allclose(torch.tensor(1.0, dtype=torch.float32))
 
 
 def test_ema_two_handlers(get_dummy_model):