FIX: accumulation should not effect value of loss

Author: louis-she

ignite/engine/__init__.py

@@ -102,11 +102,11 @@ def update(engine: Engine, batch: Sequence[torch.Tensor]) -> Union[Any, Tuple[to
         y_pred = model(x)
         loss = loss_fn(y_pred, y)
         if gradient_accumulation_steps > 1:
-            loss = loss / gradient_accumulation_steps
+            loss = loss / gradient_accumulation_steps  # fix this
         loss.backward()
         if engine.state.iteration % gradient_accumulation_steps == 0:
             optimizer.step()
-        return output_transform(x, y, y_pred, loss)
+        return output_transform(x, y, y_pred, loss * gradient_accumulation_steps)
 
     return update
 
@@ -192,7 +192,7 @@ def update(engine: Engine, batch: Sequence[torch.Tensor]) -> Union[Any, Tuple[to
             loss.backward()
             if engine.state.iteration % gradient_accumulation_steps == 0:
                 optimizer.step()
-        return output_transform(x, y, y_pred, loss)
+        return output_transform(x, y, y_pred, loss * gradient_accumulation_steps)
 
     return update
 
@@ -269,7 +269,7 @@ 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()
-        return output_transform(x, y, y_pred, loss)
+        return output_transform(x, y, y_pred, loss * gradient_accumulation_steps)
 
     return update
 
@@ -340,7 +340,7 @@ def update(engine: Engine, batch: Sequence[torch.Tensor]) -> Union[Any, Tuple[to
         y_pred = model(x)
         loss = loss_fn(y_pred, y)
         if gradient_accumulation_steps > 1:
-            loss = loss / gradient_accumulation_steps
+            loss = loss / gradient_accumulation_steps  # fix this
         loss.backward()
         if engine.state.iteration % gradient_accumulation_steps == 0:
             xm.optimizer_step(optimizer, barrier=True)

tests/ignite/engine/test_create_supervised.py

@@ -97,8 +97,8 @@ def _():
         _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())
-        # loss is not accumulated !
-        loss[0] = mse_loss(model(_x), _y).item() / gradient_accumulation_steps
+        # value of loss should not be accumulated
+        loss[0] = mse_loss(model(_x), _y).item()
 
     @trainer.on(Events.ITERATION_COMPLETED(every=gradient_accumulation_steps))
     def _():