Repository: pytorch/ignite

Issue:
Add a feature "support multi params to call forward method"?
## 🚀 Feature

<!-- A clear and concise description of the feature proposal -->
<!-- Please outline the motivation for the proposal. Is your feature request related to a problem? e.g., I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too -->
<!-- A clear and concise description of what you want to happen. -->
<!-- A clear and concise description of any alternative solutions or features you've considered, if any. -->
<!-- Add any other context or screenshots about the feature request here. -->

I think in some cases we may need pass multi params to the model. For example,
if the model has two or more inputs or some states (LSTM, etc.).
It is easy to implement this feature. The source code is here:
https://github.com/pytorch/ignite/blob/34a707e53785cf8a524589f33a570a7516fe064e/ignite/engine/__init__.py#L107C26-L107C26, and we only need to replace 
```python
output = model(x)
```
with
```python
if isinstance(x, tuple):
   output = model(*x)
else:
   output = model(x)
```
. Then multi params can be pass by `prepare_batch`, and this change is compatible with single param.

Discussion:
vfdev-5: @invoker-bot thanks for the PR, please add also a test for this feature into https://github.com/pytorch/ignite/blob/master/tests/ignite/engine/test_create_supervised.py
invoker-bot: > @invoker-bot thanks for the PR, please add also a test for this feature into https://github.com/pytorch/ignite/blob/master/tests/ignite/engine/test_create_supervised.py

I have made these changes, please check it.
vfdev-5: @invoker-bot please run code style formatting script to fix CI issues:
```bash
bash tests/run_code_style.bash install
bash tests/run_code_style.bash fmt
```
vfdev-5: @invoker-bot can you please address the comment such that the PR can be merged and will be included to the next release?
vfdev-5: @invoker-bot thanks, please also check above comments:
- https://github.com/pytorch/ignite/pull/3074#discussion_r1341340450
- https://github.com/pytorch/ignite/pull/3074#discussion_r1341340629

Patch:
diff --git a/ignite/engine/__init__.py b/ignite/engine/__init__.py
index 472b90bd..60d6f769 100644
--- a/ignite/engine/__init__.py
+++ b/ignite/engine/__init__.py
@@ -51,6 +51,7 @@ def supervised_training_step(
     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,
+    model_fn: Callable[[torch.nn.Module, Any], Any] = lambda model, x: model(x),
 ) -> Callable:
     """Factory function for supervised training.
 
@@ -71,6 +72,8 @@ def supervised_training_step(
             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.
             (default: 1 (means no gradient accumulation))
+        model_fn: the model function that receives `model` and `x`, and returns `y_pred`.
+
     Returns:
         Callable: update function.
 
@@ -91,6 +94,8 @@ def supervised_training_step(
         Added Gradient Accumulation.
     .. versionchanged:: 0.4.11
         Added `model_transform` to transform model's output
+    .. versionchanged:: 0.4.13
+        Added `model_fn` to customize model's application on the sample
     """
 
     if gradient_accumulation_steps <= 0:
@@ -104,7 +109,7 @@ def supervised_training_step(
             optimizer.zero_grad()
         model.train()
         x, y = prepare_batch(batch, device=device, non_blocking=non_blocking)
-        output = model(x)
+        output = model_fn(model, x)
         y_pred = model_transform(output)
         loss = loss_fn(y_pred, y)
         if gradient_accumulation_steps > 1:
@@ -128,6 +133,7 @@ def supervised_training_step_amp(
     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,
+    model_fn: Callable[[torch.nn.Module, Any], Any] = lambda model, x: model(x),
 ) -> Callable:
     """Factory function for supervised training using ``torch.cuda.amp``.
 
@@ -149,6 +155,7 @@ def supervised_training_step_amp(
         scaler: GradScaler instance for gradient scaling. (default: None)
         gradient_accumulation_steps: Number of steps the gradients should be accumulated across.
             (default: 1 (means no gradient accumulation))
+        model_fn: the model function that receives `model` and `x`, and returns `y_pred`.
 
     Returns:
         Callable: update function
@@ -171,6 +178,8 @@ def supervised_training_step_amp(
         Added Gradient Accumulation.
     .. versionchanged:: 0.4.11
         Added `model_transform` to transform model's output
+    .. versionchanged:: 0.4.13
+        Added `model_fn` to customize model's application on the sample
     """
 
     try:
@@ -190,7 +199,7 @@ def supervised_training_step_amp(
         model.train()
         x, y = prepare_batch(batch, device=device, non_blocking=non_blocking)
         with autocast(enabled=True):
-            output = model(x)
+            output = model_fn(model, x)
             y_pred = model_transform(output)
             loss = loss_fn(y_pred, y)
             if gradient_accumulation_steps > 1:
@@ -219,6 +228,7 @@ def supervised_training_step_apex(
     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,
+    model_fn: Callable[[torch.nn.Module, Any], Any] = lambda model, x: model(x),
 ) -> Callable:
     """Factory function for supervised training using apex.
 
@@ -239,6 +249,7 @@ def supervised_training_step_apex(
             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.
             (default: 1 (means no gradient accumulation))
+        model_fn: the model function that receives `model` and `x`, and returns `y_pred`.
 
     Returns:
         Callable: update function.
@@ -260,6 +271,8 @@ def supervised_training_step_apex(
         Added Gradient Accumulation.
     .. versionchanged:: 0.4.11
         Added `model_transform` to transform model's output
+    .. versionchanged:: 0.4.13
+        Added `model_fn` to customize model's application on the sample
     """
 
     try:
@@ -278,7 +291,7 @@ def supervised_training_step_apex(
             optimizer.zero_grad()
         model.train()
         x, y = prepare_batch(batch, device=device, non_blocking=non_blocking)
-        output = model(x)
+        output = model_fn(model, x)
         y_pred = model_transform(output)
         loss = loss_fn(y_pred, y)
         if gradient_accumulation_steps > 1:
@@ -302,6 +315,7 @@ def supervised_training_step_tpu(
     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,
+    model_fn: Callable[[torch.nn.Module, Any], Any] = lambda model, x: model(x),
 ) -> Callable:
     """Factory function for supervised training using ``torch_xla``.
 
@@ -322,6 +336,7 @@ def supervised_training_step_tpu(
             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.
             (default: 1 (means no gradient accumulation))
+        model_fn: the model function that receives `model` and `x`, and returns `y_pred`.
 
     Returns:
         Callable: update function.
@@ -343,6 +358,8 @@ def supervised_training_step_tpu(
        Added Gradient Accumulation argument for all supervised training methods.
     .. versionchanged:: 0.4.11
         Added `model_transform` to transform model's output
+    .. versionchanged:: 0.4.13
+        Added `model_fn` to customize model's application on the sample
     """
     try:
         import torch_xla.core.xla_model as xm
@@ -360,7 +377,7 @@ def supervised_training_step_tpu(
             optimizer.zero_grad()
         model.train()
         x, y = prepare_batch(batch, device=device, non_blocking=non_blocking)
-        output = model(x)
+        output = model_fn(model, x)
         y_pred = model_transform(output)
         loss = loss_fn(y_pred, y)
         if gradient_accumulation_steps > 1:
@@ -414,6 +431,7 @@ def create_supervised_trainer(
     amp_mode: Optional[str] = None,
     scaler: Union[bool, "torch.cuda.amp.GradScaler"] = False,
     gradient_accumulation_steps: int = 1,
+    model_fn: Callable[[torch.nn.Module, Any], Any] = lambda model, x: model(x),
 ) -> Engine:
     """Factory function for creating a trainer for supervised models.
 
@@ -444,6 +462,7 @@ def create_supervised_trainer(
             (default: False)
         gradient_accumulation_steps: Number of steps the gradients should be accumulated across.
             (default: 1 (means no gradient accumulation))
+        model_fn: the model function that receives `model` and `x`, and returns `y_pred`.
 
     Returns:
         a trainer engine with supervised update function.
@@ -525,6 +544,8 @@ def create_supervised_trainer(
         Added Gradient Accumulation argument for all supervised training methods.
     .. versionchanged:: 0.4.11
         Added ``model_transform`` to transform model's output
+    .. versionchanged:: 0.4.13
+        Added `model_fn` to customize model's application on the sample
     """
 
     device_type = device.type if isinstance(device, torch.device) else device
@@ -543,6 +564,7 @@ def create_supervised_trainer(
             output_transform,
             _scaler,
             gradient_accumulation_steps,
+            model_fn,
         )
     elif mode == "apex":
         _update = supervised_training_step_apex(
@@ -555,6 +577,7 @@ def create_supervised_trainer(
             model_transform,
             output_transform,
             gradient_accumulation_steps,
+            model_fn,
         )
     elif mode == "tpu":
         _update = supervised_training_step_tpu(
@@ -567,6 +590,7 @@ def create_supervised_trainer(
             model_transform,
             output_transform,
             gradient_accumulation_steps,
+            model_fn,
         )
     else:
         _update = supervised_training_step(
@@ -579,6 +603,7 @@ def create_supervised_trainer(
             model_transform,
             output_transform,
             gradient_accumulation_steps,
+            model_fn,
         )
 
     trainer = Engine(_update) if not deterministic else DeterministicEngine(_update)
@@ -595,6 +620,7 @@ def supervised_evaluation_step(
     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),
+    model_fn: Callable[[torch.nn.Module, Any], Any] = lambda model, x: model(x),
 ) -> Callable:
     """
     Factory function for supervised evaluation.
@@ -612,6 +638,7 @@ def supervised_evaluation_step(
         output_transform: function that receives 'x', 'y', 'y_pred' and returns value
             to be assigned to engine's state.output after each iteration. Default is returning `(y_pred, y,)` which fits
             output expected by metrics. If you change it you should use `output_transform` in metrics.
+        model_fn: the model function that receives `model` and `x`, and returns `y_pred`.
 
     Returns:
         Inference function.
@@ -629,13 +656,15 @@ def supervised_evaluation_step(
     .. versionadded:: 0.4.5
     .. versionchanged:: 0.4.12
         Added ``model_transform`` to transform model's output
+    .. versionchanged:: 0.4.13
+        Added `model_fn` to customize model's application on the sample
     """
 
     def evaluate_step(engine: Engine, batch: Sequence[torch.Tensor]) -> Union[Any, Tuple[torch.Tensor]]:
         model.eval()
         with torch.no_grad():
             x, y = prepare_batch(batch, device=device, non_blocking=non_blocking)
-            output = model(x)
+            output = model_fn(model, x)
             y_pred = model_transform(output)
             return output_transform(x, y, y_pred)
 
@@ -649,6 +678,7 @@ def supervised_evaluation_step_amp(
     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),
+    model_fn: Callable[[torch.nn.Module, Any], Any] = lambda model, x: model(x),
 ) -> Callable:
     """
     Factory function for supervised evaluation using ``torch.cuda.amp``.
@@ -666,6 +696,7 @@ def supervised_evaluation_step_amp(
         output_transform: function that receives 'x', 'y', 'y_pred' and returns value
             to be assigned to engine's state.output after each iteration. Default is returning `(y_pred, y,)` which fits
             output expected by metrics. If you change it you should use `output_transform` in metrics.
+        model_fn: the model function that receives `model` and `x`, and returns `y_pred`.
 
     Returns:
         Inference function.
@@ -683,6 +714,8 @@ def supervised_evaluation_step_amp(
     .. versionadded:: 0.4.5
     .. versionchanged:: 0.4.12
         Added ``model_transform`` to transform model's output
+    .. versionchanged:: 0.4.13
+        Added `model_fn` to customize model's application on the sample
     """
     try:
         from torch.cuda.amp import autocast
@@ -694,7 +727,7 @@ def supervised_evaluation_step_amp(
         with torch.no_grad():
             x, y = prepare_batch(batch, device=device, non_blocking=non_blocking)
             with autocast(enabled=True):
-                output = model(x)
+                output = model_fn(model, x)
                 y_pred = model_transform(output)
             return output_transform(x, y, y_pred)
 
@@ -710,6 +743,7 @@ def create_supervised_evaluator(
     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,
+    model_fn: Callable[[torch.nn.Module, Any], Any] = lambda model, x: model(x),
 ) -> Engine:
     """
     Factory function for creating an evaluator for supervised models.
@@ -730,6 +764,7 @@ def create_supervised_evaluator(
             output expected by metrics. If you change it you should use `output_transform` in metrics.
         amp_mode: can be ``amp``, model will be casted to float16 using
             `torch.cuda.amp <https://pytorch.org/docs/stable/amp.html>`_
+        model_fn: the model function that receives `model` and `x`, and returns `y_pred`.
 
     Returns:
         an evaluator engine with supervised inference function.
@@ -754,6 +789,8 @@ def create_supervised_evaluator(
         Added ``amp_mode`` argument for automatic mixed precision.
     .. versionchanged:: 0.4.12
         Added ``model_transform`` to transform model's output
+    .. versionchanged:: 0.4.13
+        Added `model_fn` to customize model's application on the sample
     """
     device_type = device.type if isinstance(device, torch.device) else device
     on_tpu = "xla" in device_type if device_type is not None else False
@@ -768,6 +805,7 @@ def create_supervised_evaluator(
             prepare_batch=prepare_batch,
             model_transform=model_transform,
             output_transform=output_transform,
+            model_fn=model_fn,
         )
     else:
         evaluate_step = supervised_evaluation_step(
@@ -777,6 +815,7 @@ def create_supervised_evaluator(
             prepare_batch=prepare_batch,
             model_transform=model_transform,
             output_transform=output_transform,
+            model_fn=model_fn,
         )
 
     evaluator = Engine(evaluate_step)


Test patch:
diff --git a/tests/ignite/engine/test_create_supervised.py b/tests/ignite/engine/test_create_supervised.py
index a13a95c9..8d001a8d 100644
--- a/tests/ignite/engine/test_create_supervised.py
+++ b/tests/ignite/engine/test_create_supervised.py
@@ -31,11 +31,13 @@ class DummyModel(torch.nn.Module):
         self.output_as_list = output_as_list
         self.fc = torch.nn.Linear(1, 1, bias=False)
 
-    def forward(self, x):
+    def forward(self, x, bias=None):
+        if bias is None:
+            bias = 0.0
         if self.output_as_list:
-            return self.fc(x), self.fc(x)
+            return self.fc(x) + bias, self.fc(x) + bias
 
-        return self.fc(x)
+        return self.fc(x) + bias
 
 
 def _default_create_supervised_trainer(
@@ -46,6 +48,7 @@ def _default_create_supervised_trainer(
     amp_mode: str = None,
     scaler: Union[bool, "torch.cuda.amp.GradScaler"] = False,
     with_model_transform: bool = False,
+    with_model_fn: bool = False,
 ):
     if with_model_transform:
 
@@ -65,8 +68,8 @@ def _default_create_supervised_trainer(
     optimizer = SGD(model.parameters(), 0.1)
 
     if trace:
-        example_input = torch.randn(1)
-        model = torch.jit.trace(model, example_input)
+        example_inputs = (torch.randn(1), torch.randn(1)) if with_model_fn else torch.randn(1)
+        model = torch.jit.trace(model, example_inputs)
 
     if amp_mode == "apex" and model_device == trainer_device == "cuda":
         from apex import amp
@@ -83,6 +86,9 @@ def _default_create_supervised_trainer(
         scaler=scaler,
         gradient_accumulation_steps=gradient_accumulation_steps,
         model_transform=model_transform if model_transform is not None else lambda x: x,
+        model_fn=(lambda model, x: model(x, torch.tensor([0.01], device=model_device)))
+        if with_model_fn
+        else (lambda model, x: model(x)),
     )
     assert model.fc.weight.data[0, 0].item() == approx(0.0)
     return trainer, model
@@ -96,6 +102,7 @@ def _test_create_supervised_trainer(
     amp_mode: str = None,
     scaler: Union[bool, "torch.cuda.amp.GradScaler"] = False,
     with_model_transform: bool = False,
+    with_model_fn: bool = False,
 ):
     trainer, model = _default_create_supervised_trainer(
         gradient_accumulation_steps=gradient_accumulation_steps,
@@ -105,10 +112,13 @@ def _test_create_supervised_trainer(
         amp_mode=amp_mode,
         scaler=scaler,
         with_model_transform=with_model_transform,
+        with_model_fn=with_model_fn,
     )
 
     x = torch.tensor([[0.01], [0.02], [0.03], [0.04], [0.05]])
     y = torch.tensor([[0.015], [0.025], [0.035], [0.045], [0.055]])
+    if with_model_fn:
+        y += 0.01
     data = [(_x, _y) for _x, _y in zip(x, y)]
 
     theta = [0.0]
@@ -120,12 +130,14 @@ def _test_create_supervised_trainer(
         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())
+        bias = 0.01 if with_model_fn else 0.0
+        accumulation[0] += 0.2 * _x.item() * (theta[0] * _x.item() - (_y.item() - bias))
         # value of loss should not be accumulated
+        _y_pred = model(_x, torch.tensor([bias], device=model_device)) if with_model_fn else model(_x)
         if with_model_transform:
-            loss[0] = mse_loss(model(_x)[0], _y).item()
-        else:
-            loss[0] = mse_loss(model(_x), _y).item()
+            _y_pred = _y_pred[0]
+
+        loss[0] = mse_loss(_y_pred, _y).item()
 
     @trainer.on(Events.ITERATION_COMPLETED(every=gradient_accumulation_steps))
     def _():
@@ -218,6 +230,7 @@ def _default_create_supervised_evaluator(
     trace: bool = False,
     amp_mode: str = None,
     with_model_transform: bool = False,
+    with_model_fn: bool = False,
 ):
     if with_model_transform:
 
@@ -236,14 +249,17 @@ def _default_create_supervised_evaluator(
     model.fc.weight.data.zero_()
 
     if trace:
-        example_input = torch.randn(1, 1)
-        model = torch.jit.trace(model, example_input)
+        example_inputs = (torch.randn(1), torch.randn(1)) if with_model_fn else torch.randn(1)
+        model = torch.jit.trace(model, example_inputs)
 
     evaluator = create_supervised_evaluator(
         model,
         device=evaluator_device,
         amp_mode=amp_mode,
         model_transform=model_transform if model_transform is not None else lambda x: x,
+        model_fn=(lambda model, x: model(x, torch.tensor([0.01], device=model_device)))
+        if with_model_fn
+        else (lambda model, x: model(x)),
     )
 
     assert model.fc.weight.data[0, 0].item() == approx(0.0)
@@ -257,6 +273,7 @@ def _test_create_supervised_evaluator(
     trace: bool = False,
     amp_mode: str = None,
     with_model_transform: bool = False,
+    with_model_fn: bool = False,
 ):
     model, evaluator = _default_create_supervised_evaluator(
         model_device=model_device,
@@ -264,16 +281,21 @@ def _test_create_supervised_evaluator(
         trace=trace,
         amp_mode=amp_mode,
         with_model_transform=with_model_transform,
+        with_model_fn=with_model_fn,
     )
     x = torch.tensor([[1.0], [2.0]])
     y = torch.tensor([[3.0], [5.0]])
+    if with_model_fn:
+        y += 0.01
     data = [(x, y)]
 
     if model_device == evaluator_device or ((model_device == "cpu") ^ (evaluator_device == "cpu")):
         state = evaluator.run(data)
 
         y_pred, y = state.output
-
+        if with_model_fn:
+            y_pred -= 0.01
+            y -= 0.01
         assert y_pred[0, 0].item() == approx(0.0)
         assert y_pred[1, 0].item() == approx(0.0)
         assert y[0, 0].item() == approx(3.0)
@@ -392,6 +414,7 @@ def test_create_supervised_trainer(trainer_device, 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_supervised_trainer(with_model_fn=True, trainer_device=trainer_device, trace=trace)
     _test_create_mocked_supervised_trainer(trainer_device=trainer_device, trace=trace)
 
 
@@ -576,6 +599,8 @@ def test_create_supervised_trainer_on_cuda_with_model_on_cpu():
 
 def test_create_supervised_evaluator():
     _test_create_supervised_evaluator()
+    _test_create_supervised_evaluator(with_model_transform=True)
+    _test_create_supervised_evaluator(with_model_fn=True)
     _test_mocked_supervised_evaluator()
 
     # older versions didn't have the autocast method so we skip the test for older builds