Repository: pytorch/ignite

Issue:
GradsScalarHandler logs 0 gradients if default update function is used
## 🐛 Bug description

Logging the gradients per epoch / iteration is a useful way to debug an under-performing model. Ignite provides an easy-to-use tensorboard_logger handler, an example accessible from `ignite.contrib.handlers.tensorboard_logger.GradsScalarHandler`. However, the default `update` function used by Engines generated by `create_supervised_trainer` zero the gradients before terminating, causing the handler to log zeroed out gradients all the time. 

**Steps to reproduce:**
My code is too complicated at the momemt to provide a clear insight, and I am limited by time to provide a minimal (not-)working example, so I will provide abstracted steps.

1. Generate an `Engine` / `DeterministicEngine` on an arbitrary problem by the `create_supervised_trainer` method.
2. Establish a `TensorboardLogger` and a .
3. Attach a `GradsScalarHandler` / your choice of a gradient logger. Also log the training loss or some other metric.
4. Start the training, check tensorboard and see the constant-0 gradient norms / gradients, despite the losses/metrics implying some sort of improvement/learning takes place.

**Solution proposal**

Preserving the gradients until epoch end is tricky, but not required for my purposes. If we are OK with using  `Events.ITERATION_COMPLETED` as a cue to log gradients, then we can simply modify the default `update` functions as follows:

(assuming  `engine.state.iteration` counts from 1).
```
    def update(engine: Engine, batch: Sequence[torch.Tensor]) -> Union[Any, Tuple[torch.Tensor]]:
        if (engine.state.iteration + 1) % gradient_accumulation_steps == 0:
            optimizer.zero_grad()
        model.train()
        x, y = prepare_batch(batch, device=device, non_blocking=non_blocking)
        y_pred = model(x)
        loss = loss_fn(y_pred, y)
        if gradient_accumulation_steps > 1:
            loss = loss / gradient_accumulation_steps
        loss.backward()
        if engine.state.iteration % gradient_accumulation_steps == 0:
            optimizer.step()
        return output_transform(x, y, y_pred, loss)
```

This way, upon completion of `update` and at the moment of `Events.ITERATION_COMPLETED` firing, there will be some non-zero gradients available to be logged.

## Environment (latest version of Ignite still has the same bug)

- PyTorch Version (e.g., 1.4): 1.10.1
- Ignite Version (e.g., 0.3.0): 0.4.7
- OS (e.g., Linux):
- How you installed Ignite (`conda`, `pip`, source): conda
- Python version: 3.9.7
- Any other relevant information:

GradsScalarHandler logs 0 gradients if default update function is used
## 🐛 Bug description

Logging the gradients per epoch / iteration is a useful way to debug an under-performing model. Ignite provides an easy-to-use tensorboard_logger handler, an example accessible from `ignite.contrib.handlers.tensorboard_logger.GradsScalarHandler`. However, the default `update` function used by Engines generated by `create_supervised_trainer` zero the gradients before terminating, causing the handler to log zeroed out gradients all the time. 

**Steps to reproduce:**
My code is too complicated at the momemt to provide a clear insight, and I am limited by time to provide a minimal (not-)working example, so I will provide abstracted steps.

1. Generate an `Engine` / `DeterministicEngine` on an arbitrary problem by the `create_supervised_trainer` method.
2. Establish a `TensorboardLogger` and a .
3. Attach a `GradsScalarHandler` / your choice of a gradient logger. Also log the training loss or some other metric.
4. Start the training, check tensorboard and see the constant-0 gradient norms / gradients, despite the losses/metrics implying some sort of improvement/learning takes place.

**Solution proposal**

Preserving the gradients until epoch end is tricky, but not required for my purposes. If we are OK with using  `Events.ITERATION_COMPLETED` as a cue to log gradients, then we can simply modify the default `update` functions as follows:

(assuming  `engine.state.iteration` counts from 1).
```
    def update(engine: Engine, batch: Sequence[torch.Tensor]) -> Union[Any, Tuple[torch.Tensor]]:
        if (engine.state.iteration + 1) % gradient_accumulation_steps == 0:
            optimizer.zero_grad()
        model.train()
        x, y = prepare_batch(batch, device=device, non_blocking=non_blocking)
        y_pred = model(x)
        loss = loss_fn(y_pred, y)
        if gradient_accumulation_steps > 1:
            loss = loss / gradient_accumulation_steps
        loss.backward()
        if engine.state.iteration % gradient_accumulation_steps == 0:
            optimizer.step()
        return output_transform(x, y, y_pred, loss)
```

This way, upon completion of `update` and at the moment of `Events.ITERATION_COMPLETED` firing, there will be some non-zero gradients available to be logged.

## Environment (latest version of Ignite still has the same bug)

- PyTorch Version (e.g., 1.4): 1.10.1
- Ignite Version (e.g., 0.3.0): 0.4.7
- OS (e.g., Linux):
- How you installed Ignite (`conda`, `pip`, source): conda
- Python version: 3.9.7
- Any other relevant information:



Discussion:

Resolving patch:
diff --git a/ignite/engine/__init__.py b/ignite/engine/__init__.py
--- a/ignite/engine/__init__.py
+++ b/ignite/engine/__init__.py
@@ -95,6 +95,8 @@ def supervised_training_step(
         )
 
     def update(engine: Engine, batch: Sequence[torch.Tensor]) -> Union[Any, Tuple[torch.Tensor]]:
+        if (engine.state.iteration - 1) % gradient_accumulation_steps == 0:
+            optimizer.zero_grad()
         model.train()
         x, y = prepare_batch(batch, device=device, non_blocking=non_blocking)
         y_pred = model(x)
@@ -104,7 +106,6 @@ def update(engine: Engine, batch: Sequence[torch.Tensor]) -> Union[Any, Tuple[to
         loss.backward()
         if engine.state.iteration % gradient_accumulation_steps == 0:
             optimizer.step()
-            optimizer.zero_grad()
         return output_transform(x, y, y_pred, loss)
 
     return update
@@ -173,6 +174,8 @@ def supervised_training_step_amp(
         )
 
     def update(engine: Engine, batch: Sequence[torch.Tensor]) -> Union[Any, Tuple[torch.Tensor]]:
+        if (engine.state.iteration - 1) % gradient_accumulation_steps == 0:
+            optimizer.zero_grad()
         model.train()
         x, y = prepare_batch(batch, device=device, non_blocking=non_blocking)
         with autocast(enabled=True):
@@ -185,12 +188,10 @@ def update(engine: Engine, batch: Sequence[torch.Tensor]) -> Union[Any, Tuple[to
             if engine.state.iteration % gradient_accumulation_steps == 0:
                 scaler.step(optimizer)
                 scaler.update()
-                optimizer.zero_grad()
         else:
             loss.backward()
             if engine.state.iteration % gradient_accumulation_steps == 0:
                 optimizer.step()
-                optimizer.zero_grad()
         return output_transform(x, y, y_pred, loss)
 
     return update
@@ -256,6 +257,8 @@ def supervised_training_step_apex(
         )
 
     def update(engine: Engine, batch: Sequence[torch.Tensor]) -> Union[Any, Tuple[torch.Tensor]]:
+        if (engine.state.iteration - 1) % gradient_accumulation_steps == 0:
+            optimizer.zero_grad()
         model.train()
         x, y = prepare_batch(batch, device=device, non_blocking=non_blocking)
         y_pred = model(x)
@@ -266,7 +269,6 @@ def update(engine: Engine, batch: Sequence[torch.Tensor]) -> Union[Any, Tuple[to
             scaled_loss.backward()
         if engine.state.iteration % gradient_accumulation_steps == 0:
             optimizer.step()
-            optimizer.zero_grad()
         return output_transform(x, y, y_pred, loss)
 
     return update
@@ -331,6 +333,8 @@ def supervised_training_step_tpu(
         )
 
     def update(engine: Engine, batch: Sequence[torch.Tensor]) -> Union[Any, Tuple[torch.Tensor]]:
+        if (engine.state.iteration - 1) % gradient_accumulation_steps == 0:
+            optimizer.zero_grad()
         model.train()
         x, y = prepare_batch(batch, device=device, non_blocking=non_blocking)
         y_pred = model(x)
@@ -340,7 +344,6 @@ def update(engine: Engine, batch: Sequence[torch.Tensor]) -> Union[Any, Tuple[to
         loss.backward()
         if engine.state.iteration % gradient_accumulation_steps == 0:
             xm.optimizer_step(optimizer, barrier=True)
-            optimizer.zero_grad()
         return output_transform(x, y, y_pred, loss)
 
     return update


Test patch:
diff --git a/tests/ignite/engine/test_create_supervised.py b/tests/ignite/engine/test_create_supervised.py
--- a/tests/ignite/engine/test_create_supervised.py
+++ b/tests/ignite/engine/test_create_supervised.py
@@ -93,6 +93,7 @@ def _test_create_supervised_trainer(
 
     @trainer.on(Events.ITERATION_COMPLETED)
     def _():
+        assert model.weight.grad != 0
         _x, _y = trainer.state.batch
         _x, _y = _x.to(model_device), _y.to(model_device)
         accumulation[0] += 0.2 * _x.item() * (theta[0] * _x.item() - _y.item())