Repository: pytorch/ignite

Issue:
Model Ouput Transform
## 🚀 Feature

The `*_training_step` implementations used by `create_supervised_trainer` expect only a single output from the model that's passed directly to the loss function, i.e.

        x, y = prepare_batch(batch, device=device, non_blocking=non_blocking)
        y_pred = model(x)
        loss = loss_fn(y_pred, y)

But I find that commonly models return more than predictions, i.e. they may return `logits` and `embeddings`, or `logits` and `loss` or perhaps a `dictionary`

It would be nice if you could pass a transform that would apply to the output from the model in order to extract the predictions in the form required by the loss function, i.e.

        x, y = prepare_batch(batch, device=device, non_blocking=non_blocking)
        outputs = model(x)
        y_pred = model_transform(outputs)
        loss = loss_fn(y_pred, y)

I have a current work-around - it required a lambda to handle the loss, and output_transform for each metric, i.e.

    loss = functools.partial(F.cross_entropy, label_smoothing=params['train']['label_smoothing'])
    trainer = create_supervised_trainer(model, 
                                        optimizer, 
                                        # model output is tuple, need to extract first element
                                        loss_fn=lambda output, target : loss(output[0], target), 
                                        prepare_batch=prepare_batch, 
                                        device=params['train']['device'])

    evaluator = create_supervised_evaluator(model,
                                            metrics={
                                                # model output is tuple, need to extract first element
                                                'accuracy': Accuracy(output_transform=lambda output : (output[0][0], output[1])),
                                                'loss': Loss(loss, output_transform=lambda output: (output[0][0], output[1]))},
                                            prepare_batch=prepare_batch,
                                            device=params['train']['device'])

It works fine but it was a bit fiddly to make sure I'd done it correctly, it would be a lot easier if I could just pass a single function to `create_supervised_trainer` i.e. `lambda output: output[0]`

Discussion:
vfdev-5: cc @david-waterworth if case if you would like to review this PR 
vfdev-5: @guptaaryan16 thanks for the update but let's improve test code
vfdev-5: @guptaaryan16 can you please update this PR to make it merged 
guptaaryan16: Hey @vfdev-5 , sure I am making the changes by today or tomorrow EOD and I apologise for the delay in the same
guptaaryan16: Hey @vfdev-5 can you give me an idea on how to update the API doc
vfdev-5: @guptaaryan16 I pushed this PR : https://github.com/guptaaryan16/ignite/pull/1 to your master branch on how to update the tests. Could you please merge this PR quickly to unblock this PR as we are planning the release. Thanks
guptaaryan16: Okay I am working in this issue and hope to complete this today 
vfdev-5: 

> Okay I am working in this issue and hope to complete this today. Thanks for helping with the tests should I update the docs of the API

Thanks @guptaaryan16 ! Can you add `versionchanged` into the docstrings and also usage example can be helpful
guptaaryan16: @vfdev-5 I don't know why but to make the tests work I had to remove the bias variable in the tests as it gives a NoneType error
vfdev-5: > @vfdev-5 I don't know why but to make the tests work I had to remove the bias variable in the tests as it gives a NoneType error

I checked the tests locally and they were passing. Let's revert this change and see the exact error message

Patch:
diff --git a/ignite/engine/__init__.py b/ignite/engine/__init__.py
index 406c727c..cb367fd3 100644
--- a/ignite/engine/__init__.py
+++ b/ignite/engine/__init__.py
@@ -48,6 +48,7 @@ def supervised_training_step(
     device: Optional[Union[str, torch.device]] = None,
     non_blocking: bool = False,
     prepare_batch: Callable = _prepare_batch,
+    model_transform: Callable[[Any], Any] = lambda output: output,
     output_transform: Callable[[Any, Any, Any, torch.Tensor], Any] = lambda x, y, y_pred, loss: loss.item(),
     gradient_accumulation_steps: int = 1,
 ) -> Callable:
@@ -64,6 +65,8 @@ def supervised_training_step(
             with respect to the host. For other cases, this argument has no effect.
         prepare_batch: function that receives `batch`, `device`, `non_blocking` and outputs
             tuple of tensors `(batch_x, batch_y)`.
+        model_transform: function that receives the output from the model and convert it into the form as required
+            by the loss function
         output_transform: function that receives 'x', 'y', 'y_pred', 'loss' and returns value
             to be assigned to engine's state.output after each iteration. Default is returning `loss.item()`.
         gradient_accumulation_steps: Number of steps the gradients should be accumulated across.
@@ -86,6 +89,8 @@ def supervised_training_step(
     .. versionadded:: 0.4.5
     .. versionchanged:: 0.4.7
         Added Gradient Accumulation.
+    .. versionchanged:: 0.4.11
+        Added `model_transform` to transform model's output
     """
 
     if gradient_accumulation_steps <= 0:
@@ -99,7 +104,8 @@ def supervised_training_step(
             optimizer.zero_grad()
         model.train()
         x, y = prepare_batch(batch, device=device, non_blocking=non_blocking)
-        y_pred = model(x)
+        output = model(x)
+        y_pred = model_transform(output)
         loss = loss_fn(y_pred, y)
         if gradient_accumulation_steps > 1:
             loss = loss / gradient_accumulation_steps
@@ -118,6 +124,7 @@ def supervised_training_step_amp(
     device: Optional[Union[str, torch.device]] = None,
     non_blocking: bool = False,
     prepare_batch: Callable = _prepare_batch,
+    model_transform: Callable[[Any], Any] = lambda output: output,
     output_transform: Callable[[Any, Any, Any, torch.Tensor], Any] = lambda x, y, y_pred, loss: loss.item(),
     scaler: Optional["torch.cuda.amp.GradScaler"] = None,
     gradient_accumulation_steps: int = 1,
@@ -135,6 +142,8 @@ def supervised_training_step_amp(
             with respect to the host. For other cases, this argument has no effect.
         prepare_batch: function that receives `batch`, `device`, `non_blocking` and outputs
             tuple of tensors `(batch_x, batch_y)`.
+        model_transform: function that receives the output from the model and convert it into the form as required
+            by the loss function
         output_transform: function that receives 'x', 'y', 'y_pred', 'loss' and returns value
             to be assigned to engine's state.output after each iteration. Default is returning `loss.item()`.
         scaler: GradScaler instance for gradient scaling. (default: None)
@@ -160,6 +169,8 @@ def supervised_training_step_amp(
     .. versionadded:: 0.4.5
     .. versionchanged:: 0.4.7
         Added Gradient Accumulation.
+    .. versionchanged:: 0.4.11
+        Added `model_transform` to transform model's output
     """
 
     try:
@@ -179,7 +190,8 @@ def supervised_training_step_amp(
         model.train()
         x, y = prepare_batch(batch, device=device, non_blocking=non_blocking)
         with autocast(enabled=True):
-            y_pred = model(x)
+            output = model(x)
+            y_pred = model_transform(output)
             loss = loss_fn(y_pred, y)
             if gradient_accumulation_steps > 1:
                 loss = loss / gradient_accumulation_steps
@@ -204,6 +216,7 @@ def supervised_training_step_apex(
     device: Optional[Union[str, torch.device]] = None,
     non_blocking: bool = False,
     prepare_batch: Callable = _prepare_batch,
+    model_transform: Callable[[Any], Any] = lambda output: output,
     output_transform: Callable[[Any, Any, Any, torch.Tensor], Any] = lambda x, y, y_pred, loss: loss.item(),
     gradient_accumulation_steps: int = 1,
 ) -> Callable:
@@ -220,6 +233,8 @@ def supervised_training_step_apex(
             with respect to the host. For other cases, this argument has no effect.
         prepare_batch: function that receives `batch`, `device`, `non_blocking` and outputs
             tuple of tensors `(batch_x, batch_y)`.
+        model_transform: function that receives the output from the model and convert it into the form as required
+            by the loss function
         output_transform: function that receives 'x', 'y', 'y_pred', 'loss' and returns value
             to be assigned to engine's state.output after each iteration. Default is returning `loss.item()`.
         gradient_accumulation_steps: Number of steps the gradients should be accumulated across.
@@ -243,6 +258,8 @@ def supervised_training_step_apex(
     .. versionadded:: 0.4.5
     .. versionchanged:: 0.4.7
         Added Gradient Accumulation.
+    .. versionchanged:: 0.4.11
+        Added `model_transform` to transform model's output
     """
 
     try:
@@ -261,7 +278,8 @@ def supervised_training_step_apex(
             optimizer.zero_grad()
         model.train()
         x, y = prepare_batch(batch, device=device, non_blocking=non_blocking)
-        y_pred = model(x)
+        output = model(x)
+        y_pred = model_transform(output)
         loss = loss_fn(y_pred, y)
         if gradient_accumulation_steps > 1:
             loss = loss / gradient_accumulation_steps
@@ -281,6 +299,7 @@ def supervised_training_step_tpu(
     device: Optional[Union[str, torch.device]] = None,
     non_blocking: bool = False,
     prepare_batch: Callable = _prepare_batch,
+    model_transform: Callable[[Any], Any] = lambda output: output,
     output_transform: Callable[[Any, Any, Any, torch.Tensor], Any] = lambda x, y, y_pred, loss: loss.item(),
     gradient_accumulation_steps: int = 1,
 ) -> Callable:
@@ -297,6 +316,8 @@ def supervised_training_step_tpu(
             with respect to the host. For other cases, this argument has no effect.
         prepare_batch: function that receives `batch`, `device`, `non_blocking` and outputs
             tuple of tensors `(batch_x, batch_y)`.
+        model_transform: function that receives the output from the model and convert it into the form as required
+            by the loss function
         output_transform: function that receives 'x', 'y', 'y_pred', 'loss' and returns value
             to be assigned to engine's state.output after each iteration. Default is returning `loss.item()`.
         gradient_accumulation_steps: Number of steps the gradients should be accumulated across.
@@ -320,6 +341,8 @@ def supervised_training_step_tpu(
     .. versionadded:: 0.4.5
     .. versionchanged:: 0.4.7
        Added Gradient Accumulation argument for all supervised training methods.
+    .. versionchanged:: 0.4.11
+        Added `model_transform` to transform model's output
     """
     try:
         import torch_xla.core.xla_model as xm
@@ -337,7 +360,8 @@ def supervised_training_step_tpu(
             optimizer.zero_grad()
         model.train()
         x, y = prepare_batch(batch, device=device, non_blocking=non_blocking)
-        y_pred = model(x)
+        output = model(x)
+        y_pred = model_transform(output)
         loss = loss_fn(y_pred, y)
         if gradient_accumulation_steps > 1:
             loss = loss / gradient_accumulation_steps
@@ -384,6 +408,7 @@ def create_supervised_trainer(
     device: Optional[Union[str, torch.device]] = None,
     non_blocking: bool = False,
     prepare_batch: Callable = _prepare_batch,
+    model_transform: Callable[[Any], Any] = lambda output: output,
     output_transform: Callable[[Any, Any, Any, torch.Tensor], Any] = lambda x, y, y_pred, loss: loss.item(),
     deterministic: bool = False,
     amp_mode: Optional[str] = None,
@@ -403,6 +428,8 @@ def create_supervised_trainer(
             with respect to the host. For other cases, this argument has no effect.
         prepare_batch: function that receives `batch`, `device`, `non_blocking` and outputs
             tuple of tensors `(batch_x, batch_y)`.
+        model_transform: function that receives the output from the model and convert it into the form as required
+            by the loss function
         output_transform: function that receives 'x', 'y', 'y_pred', 'loss' and returns value
             to be assigned to engine's state.output after each iteration. Default is returning `loss.item()`.
         deterministic: if True, returns deterministic engine of type
@@ -496,6 +523,8 @@ def create_supervised_trainer(
 
     .. versionchanged:: 0.4.7
         Added Gradient Accumulation argument for all supervised training methods.
+    .. versionchanged:: 0.4.11
+        Added `model_transform` to transform model's output
     """
 
     device_type = device.type if isinstance(device, torch.device) else device
@@ -510,6 +539,7 @@ def create_supervised_trainer(
             device,
             non_blocking,
             prepare_batch,
+            model_transform,
             output_transform,
             _scaler,
             gradient_accumulation_steps,
@@ -522,6 +552,7 @@ def create_supervised_trainer(
             device,
             non_blocking,
             prepare_batch,
+            model_transform,
             output_transform,
             gradient_accumulation_steps,
         )
@@ -533,6 +564,7 @@ def create_supervised_trainer(
             device,
             non_blocking,
             prepare_batch,
+            model_transform,
             output_transform,
             gradient_accumulation_steps,
         )
@@ -544,6 +576,7 @@ def create_supervised_trainer(
             device,
             non_blocking,
             prepare_batch,
+            model_transform,
             output_transform,
             gradient_accumulation_steps,
         )
@@ -662,6 +695,7 @@ def create_supervised_evaluator(
     device: Optional[Union[str, torch.device]] = None,
     non_blocking: bool = False,
     prepare_batch: Callable = _prepare_batch,
+    model_transform: Callable[[Any], Any] = lambda output: output,
     output_transform: Callable[[Any, Any, Any], Any] = lambda x, y, y_pred: (y_pred, y),
     amp_mode: Optional[str] = None,
 ) -> Engine:


Test patch:
diff --git a/tests/ignite/engine/test_create_supervised.py b/tests/ignite/engine/test_create_supervised.py
index 49ac25a8..36dfe986 100644
--- a/tests/ignite/engine/test_create_supervised.py
+++ b/tests/ignite/engine/test_create_supervised.py
@@ -8,7 +8,6 @@ import pytest
 import torch
 from packaging.version import Version
 from pytest import approx
-from torch.nn import Linear
 from torch.nn.functional import mse_loss
 from torch.optim import SGD
 
@@ -26,6 +25,19 @@ from ignite.engine import (
 from ignite.metrics import MeanSquaredError
 
 
+class DummyModel(torch.nn.Module):
+    def __init__(self, output_as_list=False):
+        super(DummyModel, self).__init__()
+        self.output_as_list = output_as_list
+        self.fc = torch.nn.Linear(1, 1, bias=False)
+
+    def forward(self, x):
+        if self.output_as_list:
+            return [self.fc(x), self.fc(x)]
+
+        return self.fc(x)
+
+
 def _default_create_supervised_trainer(
     gradient_accumulation_steps: int = 1,
     model_device: Optional[str] = None,
@@ -33,13 +45,23 @@ def _default_create_supervised_trainer(
     trace: bool = False,
     amp_mode: str = None,
     scaler: Union[bool, "torch.cuda.amp.GradScaler"] = False,
+    with_model_transform: bool = False,
 ):
-    model = Linear(1, 1, bias=False)
+    if with_model_transform:
+
+        def get_first_element(output):
+            return output[0]
+
+        model = DummyModel(output_as_list=True)
+        model_transform = get_first_element
+    else:
+        model = DummyModel()
+        model_transform = None
 
     if model_device:
         model.to(model_device)
 
-    model.weight.data.zero_()
+    model.fc.weight.data.zero_()
     optimizer = SGD(model.parameters(), 0.1)
 
     if trace:
@@ -60,9 +82,9 @@ def _default_create_supervised_trainer(
         amp_mode=amp_mode,
         scaler=scaler,
         gradient_accumulation_steps=gradient_accumulation_steps,
+        model_transform=model_transform if model_transform is not None else lambda x: x,
     )
-    assert model.weight.data[0, 0].item() == approx(0.0)
-
+    assert model.fc.weight.data[0, 0].item() == approx(0.0)
     return trainer, model
 
 
@@ -73,6 +95,7 @@ def _test_create_supervised_trainer(
     trace: bool = False,
     amp_mode: str = None,
     scaler: Union[bool, "torch.cuda.amp.GradScaler"] = False,
+    with_model_transform: bool = False,
 ):
     trainer, model = _default_create_supervised_trainer(
         gradient_accumulation_steps=gradient_accumulation_steps,
@@ -81,6 +104,7 @@ def _test_create_supervised_trainer(
         trace=trace,
         amp_mode=amp_mode,
         scaler=scaler,
+        with_model_transform=with_model_transform,
     )
 
     x = torch.tensor([[0.01], [0.02], [0.03], [0.04], [0.05]])
@@ -93,17 +117,20 @@ def _test_create_supervised_trainer(
 
     @trainer.on(Events.ITERATION_COMPLETED)
     def _():
-        assert model.weight.grad != 0
+        assert model.fc.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())
         # value of loss should not be accumulated
-        loss[0] = mse_loss(model(_x), _y).item()
+        if with_model_transform:
+            loss[0] = mse_loss(model(_x)[0], _y).item()
+        else:
+            loss[0] = mse_loss(model(_x), _y).item()
 
     @trainer.on(Events.ITERATION_COMPLETED(every=gradient_accumulation_steps))
     def _():
         theta[0] -= accumulation[0] / gradient_accumulation_steps
-        assert pytest.approx(model.weight.data[0, 0].item(), abs=1.0e-5) == theta[0]
+        assert pytest.approx(model.fc.weight.data[0, 0].item(), abs=1.0e-5) == theta[0]
         assert pytest.approx(trainer.state.output[-1], abs=1e-5) == loss[0]
         accumulation[0] = loss[0] = 0.0
 
@@ -175,10 +202,16 @@ def _test_create_mocked_supervised_trainer(
                             assert training_step_mock.called
 
 
-def _test_create_supervised_trainer_wrong_accumulation(model_device=None, trainer_device=None, amp_mode=None):
+def _test_create_supervised_trainer_wrong_accumulation(
+    model_device=None, trainer_device=None, amp_mode=None, trace=False
+):
     with pytest.raises(ValueError, match="Gradient_accumulation_steps must be strictly positive."):
         _default_create_supervised_trainer(
-            gradient_accumulation_steps=0, model_device=model_device, trainer_device=trainer_device, amp_mode=amp_mode
+            gradient_accumulation_steps=0,
+            model_device=model_device,
+            trainer_device=trainer_device,
+            amp_mode=amp_mode,
+            trace=trace,
         )
 
 
@@ -188,13 +221,12 @@ def _default_create_supervised_evaluator(
     trace: bool = False,
     amp_mode: str = None,
 ):
-    model = Linear(1, 1)
+    model = DummyModel()
 
     if model_device:
         model.to(model_device)
 
-    model.weight.data.zero_()
-    model.bias.data.zero_()
+    model.fc.weight.data.zero_()
 
     if trace:
         example_input = torch.randn(1, 1)
@@ -202,8 +234,7 @@ def _default_create_supervised_evaluator(
 
     evaluator = create_supervised_evaluator(model, device=evaluator_device, amp_mode=amp_mode)
 
-    assert model.weight.data[0, 0].item() == approx(0.0)
-    assert model.bias.item() == approx(0.0)
+    assert model.fc.weight.data[0, 0].item() == approx(0.0)
 
     return model, evaluator
 
@@ -231,8 +262,7 @@ def _test_create_supervised_evaluator(
         assert y[0, 0].item() == approx(3.0)
         assert y[1, 0].item() == approx(5.0)
 
-        assert model.weight.data[0, 0].item() == approx(0.0)
-        assert model.bias.item() == approx(0.0)
+        assert model.fc.weight.data[0, 0].item() == approx(0.0)
 
     else:
         if Version(torch.__version__) >= Version("1.7.0"):
@@ -277,13 +307,12 @@ def _test_create_evaluation_step_amp(
 ):
 
     output_transform_mock = MagicMock()
-    model = Linear(1, 1)
+    model = DummyModel()
 
     if model_device:
         model.to(model_device)
 
-    model.weight.data.zero_()
-    model.bias.data.zero_()
+    model.fc.weight.data.zero_()
 
     if trace:
         example_input = torch.randn(1, 1)
@@ -313,13 +342,12 @@ def _test_create_evaluation_step(
     amp_mode: str = None,
 ):
     output_transform_mock = MagicMock()
-    model = Linear(1, 1)
+    model = DummyModel()
 
     if model_device:
         model.to(model_device)
 
-    model.weight.data.zero_()
-    model.bias.data.zero_()
+    model.fc.weight.data.zero_()
 
     if trace:
         example_input = torch.randn(1, 1)
@@ -341,25 +369,14 @@ def _test_create_evaluation_step(
     assert output_transform_mock.called
 
 
-def test_create_supervised_trainer():
-    _test_create_supervised_trainer_wrong_accumulation()
-    _test_create_supervised_trainer(gradient_accumulation_steps=1)
-    _test_create_supervised_trainer(gradient_accumulation_steps=3)
-    _test_create_mocked_supervised_trainer()
-
-
-def test_create_supervised_trainer_with_cpu():
-    _test_create_supervised_trainer_wrong_accumulation(trainer_device="cpu")
-    _test_create_supervised_trainer(gradient_accumulation_steps=1, trainer_device="cpu")
-    _test_create_supervised_trainer(gradient_accumulation_steps=3, trainer_device="cpu")
-    _test_create_mocked_supervised_trainer(trainer_device="cpu")
-
-
-def test_create_supervised_trainer_traced_with_cpu():
-    _test_create_supervised_trainer_wrong_accumulation(trainer_device="cpu")
-    _test_create_supervised_trainer(gradient_accumulation_steps=1, trainer_device="cpu", trace=True)
-    _test_create_supervised_trainer(gradient_accumulation_steps=3, trainer_device="cpu", trace=True)
-    _test_create_mocked_supervised_trainer(trainer_device="cpu", trace=True)
+@pytest.mark.parametrize("trainer_device", [None, "cpu"])
+@pytest.mark.parametrize("trace", [False, True])
+def test_create_supervised_trainer(trainer_device, trace):
+    _test_create_supervised_trainer_wrong_accumulation(trainer_device=trainer_device, trace=trace)
+    _test_create_supervised_trainer(gradient_accumulation_steps=1, trainer_device=trainer_device, trace=trace)
+    _test_create_supervised_trainer(gradient_accumulation_steps=3, trainer_device=trainer_device, trace=trace)
+    _test_create_supervised_trainer(with_model_transform=True, trainer_device=trainer_device, trace=trace)
+    _test_create_mocked_supervised_trainer(trainer_device=trainer_device, trace=trace)
 
 
 @pytest.mark.skipif(find_spec("apex"), reason="Skip if APEX")
@@ -599,9 +616,8 @@ def test_create_supervised_evaluator_amp_error(mock_torch_cuda_amp_module):
 
 
 def test_create_supervised_evaluator_with_metrics():
-    model = Linear(1, 1)
-    model.weight.data.zero_()
-    model.bias.data.zero_()
+    model = DummyModel()
+    model.fc.weight.data.zero_()
 
     evaluator = create_supervised_evaluator(model, metrics={"mse": MeanSquaredError()})