Repository: pytorch/ignite

Issue:
Update cifar10 example, validation with amp
Currently, we provide AMP option for the training, but model's evaluation still run on plain f32. Recently, we merged [transformers example](https://github.com/pytorch/ignite/tree/master/examples/contrib/transformers) where evaluation takes amp into account : https://github.com/pytorch/ignite/blob/master/examples/contrib/transformers/main.py#L367

The idea is to make cifar10 example compatible as transformers example:

Replace 
https://github.com/pytorch/ignite/blob/7753eabc5af8bed2cca4713805613adfbc460a8e/examples/contrib/cifar10/main.py#L86-L87

by 
https://github.com/pytorch/ignite/blob/7753eabc5af8bed2cca4713805613adfbc460a8e/examples/contrib/transformers/main.py#L88-L89

and create similar `create_evaluator` helper method.


Discussion:
I'd like to pick this one up, I have a question though.

Is the idea to update that method invocation https://github.com/pytorch/ignite/blob/master/examples/contrib/cifar10/main.py#L86 to pass AMP into evaluation as well? 

Since the evaluator https://github.com/pytorch/ignite/blob/master/ignite/engine/__init__.py#L370 doesn't have the param we have the option with either 
1. Add an optional param for the method.
2. Just create a custom evaluator method for the sake of the example.

I am leaning towards 1 as it offers more generalized solution, any preference on that?
I updated the description to make it more explicit.
I'd say, let's do 2nd option for this issue. But I agree in a long term we can replace that with another solution (`create_supervised_evaluator` could accept `amp_mode` argument) that could join recently merged PR for AMP: https://github.com/pytorch/ignite/pull/1589 (cc @ydcjeff)


Resolving patch:
diff --git a/examples/contrib/cifar10/main.py b/examples/contrib/cifar10/main.py
--- a/examples/contrib/cifar10/main.py
+++ b/examples/contrib/cifar10/main.py
@@ -12,7 +12,7 @@
 import ignite.distributed as idist
 from ignite.contrib.engines import common
 from ignite.contrib.handlers import PiecewiseLinear
-from ignite.engine import Engine, Events, create_supervised_evaluator
+from ignite.engine import Engine, Events
 from ignite.handlers import Checkpoint, DiskSaver, global_step_from_engine
 from ignite.metrics import Accuracy, Loss
 from ignite.utils import manual_seed, setup_logger
@@ -83,8 +83,8 @@ def training(local_rank, config):
 
     # We define two evaluators as they wont have exactly similar roles:
     # - `evaluator` will save the best model based on validation score
-    evaluator = create_supervised_evaluator(model, metrics=metrics, device=device, non_blocking=True)
-    train_evaluator = create_supervised_evaluator(model, metrics=metrics, device=device, non_blocking=True)
+    evaluator = create_evaluator(model, metrics=metrics, config=config)
+    train_evaluator = create_evaluator(model, metrics=metrics, config=config)
 
     def run_validation(engine):
         epoch = trainer.state.epoch
@@ -193,6 +193,8 @@ def run(
     del config["spawn_kwargs"]
 
     spawn_kwargs["nproc_per_node"] = nproc_per_node
+    if backend == "xla-tpu" and with_amp:
+        raise RuntimeError("The value of with_amp should be False if backend is xla")
 
     with idist.Parallel(backend=backend, **spawn_kwargs) as parallel:
 
@@ -349,6 +351,33 @@ def train_step(engine, batch):
     return trainer
 
 
+def create_evaluator(model, metrics, config, tag="val"):
+    with_amp = config["with_amp"]
+    device = idist.device()
+
+    @torch.no_grad()
+    def evaluate_step(engine: Engine, batch):
+        model.eval()
+        x, y = batch[0], batch[1]
+        if x.device != device:
+            x = x.to(device, non_blocking=True)
+            y = y.to(device, non_blocking=True)
+
+        with autocast(enabled=with_amp):
+            output = model(x)
+        return output, y
+
+    evaluator = Engine(evaluate_step)
+
+    for name, metric in metrics.items():
+        metric.attach(evaluator, name)
+
+    if idist.get_rank() == 0 and (not config["with_clearml"]):
+        common.ProgressBar(desc=f"Evaluation ({tag})", persist=False).attach(evaluator)
+
+    return evaluator
+
+
 def get_save_handler(config):
     if config["with_clearml"]:
         from ignite.contrib.handlers.clearml_logger import ClearMLSaver


Test patch: