Repository: pytorch/ignite

Issue:
Improve typing hints for Engine and helper methods
## 🚀 Feature

1) Let's improve typing hints for process_function arg:
- https://github.com/pytorch/ignite/blob/8a1de70e2152936ffcb051f3bd618138130626ec/ignite/engine/engine.py#L122
We have to provide the signature: [Engine, Any] -> Any in mypy recognizable way

Also here: https://github.com/pytorch/ignite/blob/8a1de70e2152936ffcb051f3bd618138130626ec/ignite/engine/deterministic.py#L179

2) Let's improve typing hints for `prepare_batch` and `output_transform` for `create_supervised_trainer` and `create_supervised_evaluator`.

https://github.com/pytorch/ignite/blob/8a1de70e2152936ffcb051f3bd618138130626ec/ignite/engine/__init__.py#L52-L53

3) Provide example of usage of `prepare_batch` and `output_transform` in the docstring: https://pytorch.org/ignite/generated/ignite.engine.create_supervised_trainer.html


cc @Priyansi 


Discussion:
Many of the callbacks like `prepare_batch` are called with keyword arguments in many places in the source code. But `typing.Callable` can only annotated positional arguments, not keyword arguments or more complicated cases.

So if we change the annotation of the callbacks by `typing.Callable`.  `mypy` will complaint about it where there is a invokation with keyword arguments.

`typing.Protocol` is more suitable to annotate callbacks but it's available since Python 3.8, so for the moment we should skip this.

![](https://cdn.discordapp.com/attachments/844863443893878824/921370957401186344/unknown.png)

![](https://cdn.discordapp.com/attachments/844863443893878824/921370842267533353/unknown.png)


@louis-she would you like to fix points 1 and 3 ?
> @louis-she would you like to fix points 1 and 3 ?

Of course, i will work on it, u can assign the task to me.

Resolving patch:
diff --git a/ignite/engine/__init__.py b/ignite/engine/__init__.py
--- a/ignite/engine/__init__.py
+++ b/ignite/engine/__init__.py
@@ -48,7 +48,7 @@ def supervised_training_step(
     device: Optional[Union[str, torch.device]] = None,
     non_blocking: bool = False,
     prepare_batch: Callable = _prepare_batch,
-    output_transform: Callable = lambda x, y, y_pred, loss: loss.item(),
+    output_transform: Callable[[Any, Any, Any, torch.Tensor], Any] = lambda x, y, y_pred, loss: loss.item(),
     gradient_accumulation_steps: int = 1,
 ) -> Callable:
     """Factory function for supervised training.
@@ -117,7 +117,7 @@ def supervised_training_step_amp(
     device: Optional[Union[str, torch.device]] = None,
     non_blocking: bool = False,
     prepare_batch: Callable = _prepare_batch,
-    output_transform: Callable = lambda x, y, y_pred, loss: loss.item(),
+    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,
 ) -> Callable:
@@ -203,7 +203,7 @@ def supervised_training_step_apex(
     device: Optional[Union[str, torch.device]] = None,
     non_blocking: bool = False,
     prepare_batch: Callable = _prepare_batch,
-    output_transform: Callable = lambda x, y, y_pred, loss: loss.item(),
+    output_transform: Callable[[Any, Any, Any, torch.Tensor], Any] = lambda x, y, y_pred, loss: loss.item(),
     gradient_accumulation_steps: int = 1,
 ) -> Callable:
     """Factory function for supervised training using apex.
@@ -279,7 +279,7 @@ def supervised_training_step_tpu(
     device: Optional[Union[str, torch.device]] = None,
     non_blocking: bool = False,
     prepare_batch: Callable = _prepare_batch,
-    output_transform: Callable = lambda x, y, y_pred, loss: loss.item(),
+    output_transform: Callable[[Any, Any, Any, torch.Tensor], Any] = lambda x, y, y_pred, loss: loss.item(),
     gradient_accumulation_steps: int = 1,
 ) -> Callable:
     """Factory function for supervised training using ``torch_xla``.
@@ -381,7 +381,7 @@ def create_supervised_trainer(
     device: Optional[Union[str, torch.device]] = None,
     non_blocking: bool = False,
     prepare_batch: Callable = _prepare_batch,
-    output_transform: Callable = lambda x, y, y_pred, loss: loss.item(),
+    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,
     scaler: Union[bool, "torch.cuda.amp.GradScaler"] = False,
@@ -418,6 +418,50 @@ def create_supervised_trainer(
     Returns:
         a trainer engine with supervised update function.
 
+    Examples:
+
+        Create a trainer
+
+        .. code-block:: python
+
+            from ignite.engine import create_supervised_trainer
+            from ignite.utils import convert_tensor
+            from ignite.contrib.handlers.tqdm_logger import ProgressBar
+
+            model = ...
+            loss = ...
+            optimizer = ...
+            dataloader = ...
+
+            def prepare_batch_fn(batch, device, non_blocking):
+                x = ...  # get x from batch
+                y = ...  # get y from batch
+
+                # return a tuple of (x, y) that can be directly runned as
+                # `loss_fn(model(x), y)`
+                return (
+                    convert_tensor(x, device, non_blocking),
+                    convert_tensor(y, device, non_blocking)
+                )
+
+            def output_transform_fn(x, y, y_pred, loss):
+                # return only the loss is actually the default behavior for
+                # trainer engine, but you can return anything you want
+                return loss.item()
+
+            trainer = create_supervised_trainer(
+                model,
+                optimizer,
+                loss,
+                prepare_batch=prepare_batch_fn,
+                output_transform=output_transform_fn
+            )
+
+            pbar = ProgressBar()
+            pbar.attach(trainer, output_transform=lambda x: {"loss": x})
+
+            trainer.run(dataloader, max_epochs=5)
+
     Note:
         If ``scaler`` is True, GradScaler instance will be created internally and trainer state has attribute named
         ``scaler`` for that instance and can be used for saving and loading.
@@ -513,7 +557,7 @@ def supervised_evaluation_step(
     device: Optional[Union[str, torch.device]] = None,
     non_blocking: bool = False,
     prepare_batch: Callable = _prepare_batch,
-    output_transform: Callable = lambda x, y, y_pred: (y_pred, y),
+    output_transform: Callable[[Any, Any, Any], Any] = lambda x, y, y_pred: (y_pred, y),
 ) -> Callable:
     """
     Factory function for supervised evaluation.
@@ -561,7 +605,7 @@ def supervised_evaluation_step_amp(
     device: Optional[Union[str, torch.device]] = None,
     non_blocking: bool = False,
     prepare_batch: Callable = _prepare_batch,
-    output_transform: Callable = lambda x, y, y_pred: (y_pred, y),
+    output_transform: Callable[[Any, Any, Any], Any] = lambda x, y, y_pred: (y_pred, y),
 ) -> Callable:
     """
     Factory function for supervised evaluation using ``torch.cuda.amp``.
@@ -615,7 +659,7 @@ def create_supervised_evaluator(
     device: Optional[Union[str, torch.device]] = None,
     non_blocking: bool = False,
     prepare_batch: Callable = _prepare_batch,
-    output_transform: Callable = lambda x, y, y_pred: (y_pred, y),
+    output_transform: Callable[[Any, Any, Any], Any] = lambda x, y, y_pred: (y_pred, y),
     amp_mode: Optional[str] = None,
 ) -> Engine:
     """
diff --git a/ignite/engine/deterministic.py b/ignite/engine/deterministic.py
--- a/ignite/engine/deterministic.py
+++ b/ignite/engine/deterministic.py
@@ -176,7 +176,7 @@ class DeterministicEngine(Engine):
             in each iteration, and returns data to be stored in the engine's state.
     """
 
-    def __init__(self, process_function: Callable):
+    def __init__(self, process_function: Callable[[Engine, Any], Any]):
         super(DeterministicEngine, self).__init__(process_function)
         self.state_dict_user_keys.append("rng_states")
         self.add_event_handler(Events.STARTED, self._init_run)
diff --git a/ignite/engine/engine.py b/ignite/engine/engine.py
--- a/ignite/engine/engine.py
+++ b/ignite/engine/engine.py
@@ -120,7 +120,7 @@ def compute_mean_std(engine, batch):
     _state_dict_all_req_keys = ("epoch_length", "max_epochs")
     _state_dict_one_of_opt_keys = ("iteration", "epoch")
 
-    def __init__(self, process_function: Callable):
+    def __init__(self, process_function: Callable[["Engine", Any], Any]):
         self._event_handlers = defaultdict(list)  # type: Dict[Any, List]
         self.logger = logging.getLogger(__name__ + "." + self.__class__.__name__)
         self._process_function = process_function


Test patch: