Repository: pytorch/ignite

Issue:
add automatic mixed precision in create_supervised_evaluator
## 🚀 Feature

Like #1235, add automatic mixed precision (both torch.cuda.amp) in `create_supervised_evaluator`.
PR that can reference: #1589 
Related issue: #1666



Discussion:
How shall we call them? `supervised_evaluation_step_*` or `supervised_validation_step_*`?
`supervised_evaluation_step_*` is good, IMO  
@ydcjeff  Are you currently working on this one? If not I'd like to give it a shot
 
@Moh-Yakoub I am not working on this.
Feel free to go ahead

Resolving patch:
diff --git a/ignite/engine/__init__.py b/ignite/engine/__init__.py
--- a/ignite/engine/__init__.py
+++ b/ignite/engine/__init__.py
@@ -25,6 +25,8 @@
     "supervised_training_step_amp",
     "supervised_training_step_apex",
     "supervised_training_step_tpu",
+    "supervised_evaluation_step",
+    "supervised_evaluation_step_amp",
 ]
 
 
@@ -412,6 +414,107 @@ def create_supervised_trainer(
     return trainer
 
 
+def supervised_evaluation_step(
+    model: torch.nn.Module,
+    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),
+) -> Callable:
+    """
+    Factory function for supervised evaluation.
+
+    Args:
+        model: the model to train.
+        device: device type specification (default: None).
+            Applies to batches after starting the engine. Model *will not* be moved.
+        non_blocking: if True and this copy is between CPU and GPU, the copy may occur asynchronously
+            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)`.
+        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.
+
+    Returns:
+        Inference function.
+
+    Note:
+        `engine.state.output` for this engine is defind by `output_transform` parameter and is
+        a tuple of `(batch_pred, batch_y)` by default.
+
+    .. warning::
+
+        The internal use of `device` has changed.
+        `device` will now *only* be used to move the input data to the correct device.
+        The `model` should be moved by the user before creating an optimizer.
+
+    .. versionadded:: 0.5.0
+    """
+
+    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)
+            y_pred = model(x)
+            return output_transform(x, y, y_pred)
+
+    return evaluate_step
+
+
+def supervised_evaluation_step_amp(
+    model: torch.nn.Module,
+    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),
+) -> Callable:
+    """
+    Factory function for supervised evaluation using ``torch.cuda.amp``.
+
+    Args:
+        model: the model to train.
+        device: device type specification (default: None).
+            Applies to batches after starting the engine. Model *will not* be moved.
+        non_blocking: if True and this copy is between CPU and GPU, the copy may occur asynchronously
+            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)`.
+        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.
+
+    Returns:
+        Inference function.
+
+    Note:
+        `engine.state.output` for this engine is defind by `output_transform` parameter and is
+        a tuple of `(batch_pred, batch_y)` by default.
+
+    .. warning::
+
+        The internal use of `device` has changed.
+        `device` will now *only* be used to move the input data to the correct device.
+        The `model` should be moved by the user before creating an optimizer.
+
+    .. versionadded:: 0.5.0
+    """
+    try:
+        from torch.cuda.amp import autocast
+    except ImportError:
+        raise ImportError("Please install torch>=1.6.0 to use amp_mode='amp'.")
+
+    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)
+            with autocast(enabled=True):
+                y_pred = model(x)
+            return output_transform(x, y, y_pred)
+
+    return evaluate_step
+
+
 def create_supervised_evaluator(
     model: torch.nn.Module,
     metrics: Optional[Dict[str, Metric]] = None,
@@ -419,6 +522,7 @@ def create_supervised_evaluator(
     non_blocking: bool = False,
     prepare_batch: Callable = _prepare_batch,
     output_transform: Callable = lambda x, y, y_pred: (y_pred, y),
+    amp_mode: Optional[str] = None,
 ) -> Engine:
     """
     Factory function for creating an evaluator for supervised models.
@@ -435,6 +539,8 @@ def create_supervised_evaluator(
         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.
+        amp_mode: can be ``amp``, model will be casted to float16 using
+            `torch.cuda.amp <https://pytorch.org/docs/stable/amp.html>`_
 
     Returns:
         an evaluator engine with supervised inference function.
@@ -454,17 +560,21 @@ def create_supervised_evaluator(
         - `PyTorch Documentation <https://pytorch.org/docs/stable/optim.html#constructing-it>`_
 
         - `PyTorch's Explanation <https://github.com/pytorch/pytorch/issues/7844#issuecomment-503713840>`_
+
+    .. versionchanged:: 0.5.0
+        - Added ``amp_mode`` argument for automatic mixed precision.
     """
-    metrics = metrics or {}
+    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
+    mode, _ = _check_arg(on_tpu, amp_mode, None)
 
-    def _inference(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)
-            y_pred = model(x)
-            return output_transform(x, y, y_pred)
+    metrics = metrics or {}
+    if mode == "amp":
+        evaluate_step = supervised_evaluation_step_amp(model, device, non_blocking, prepare_batch, output_transform)
+    else:
+        evaluate_step = supervised_evaluation_step(model, device, non_blocking, prepare_batch, output_transform)
 
-    evaluator = Engine(_inference)
+    evaluator = Engine(evaluate_step)
 
     for name, metric in metrics.items():
         metric.attach(evaluator, name)


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
@@ -12,7 +12,13 @@
 from torch.optim import SGD
 
 import ignite.distributed as idist
-from ignite.engine import create_supervised_evaluator, create_supervised_trainer, supervised_training_step_tpu
+from ignite.engine import (
+    create_supervised_evaluator,
+    create_supervised_trainer,
+    supervised_evaluation_step,
+    supervised_evaluation_step_amp,
+    supervised_training_step_tpu,
+)
 from ignite.metrics import MeanSquaredError
 
 
@@ -79,7 +85,10 @@ def _test_create_supervised_trainer(
 
 
 def _test_create_supervised_evaluator(
-    model_device: Optional[str] = None, evaluator_device: Optional[str] = None, trace: bool = False
+    model_device: Optional[str] = None,
+    evaluator_device: Optional[str] = None,
+    trace: bool = False,
+    amp_mode: str = None,
 ):
     model = Linear(1, 1)
 
@@ -93,7 +102,7 @@ def _test_create_supervised_evaluator(
         example_input = torch.randn(1, 1)
         model = torch.jit.trace(model, example_input)
 
-    evaluator = create_supervised_evaluator(model, device=evaluator_device)
+    evaluator = create_supervised_evaluator(model, device=evaluator_device, amp_mode=amp_mode)
 
     x = torch.tensor([[1.0], [2.0]])
     y = torch.tensor([[3.0], [5.0]])
@@ -259,6 +268,18 @@ def test_create_supervised_evaluator_on_cuda_with_model_on_cpu():
     _test_create_supervised_evaluator(evaluator_device="cuda")
 
 
+@pytest.mark.skipif(LooseVersion(torch.__version__) < LooseVersion("1.6.0"), reason="Skip if < 1.6.0")
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Skip if no GPU")
+def test_create_supervised_evaluator_on_cuda_amp():
+    model_device = evaluator_device = "cuda"
+    _test_create_supervised_evaluator(model_device=model_device, evaluator_device=evaluator_device, amp_mode="amp")
+
+
+def test_create_supervised_evaluator_amp_error(mock_torch_cuda_amp_module):
+    with pytest.raises(ImportError, match="Please install torch>=1.6.0 to use amp_mode='amp'."):
+        _test_create_supervised_evaluator(amp_mode="amp")
+
+
 def test_create_supervised_evaluator_with_metrics():
     model = Linear(1, 1)
     model.weight.data.zero_()