Repository: pytorch/ignite

Issue:
Make EpochMetric work in DDP
## 🚀 Feature

Currently, metrics like `EpochMetric` which store complete history of predictions and targets, are not supporting DDP setting. 
Idea is to perform [gather](https://pytorch.org/docs/stable/distributed.html#torch.distributed.gather) on zero rank process in `compute` method to accumulate all predictions and targets across all the processes and then execute metric function...
We may also need to warn about high memory consumption of this approach.

cc @sdesrozis 


Discussion:
Maybe this feature could be disabled by user if he can do better.
I agree

Resolving patch:
diff --git a/ignite/metrics/epoch_metric.py b/ignite/metrics/epoch_metric.py
--- a/ignite/metrics/epoch_metric.py
+++ b/ignite/metrics/epoch_metric.py
@@ -1,9 +1,11 @@
 import warnings
-from typing import Callable, Sequence
+from typing import Callable, Sequence, Union
 
 import torch
 
-from ignite.metrics.metric import Metric
+import ignite.distributed as idist
+from ignite.exceptions import NotComputableError
+from ignite.metrics.metric import Metric, reinit__is_reduced
 
 __all__ = ["EpochMetric"]
 
@@ -18,10 +20,10 @@ class EpochMetric(Metric):
         Current implementation stores all input data (output and target) in as tensors before computing a metric.
         This can potentially lead to a memory error if the input data is larger than available RAM.
 
-    .. warning::
-
-        Current implementation does not work with distributed computations. Results are not gather across all devices
-        and computed results are valid for a single device only.
+        In distributed configuration, all stored data (output and target) is mutually collected across all processes
+        using all gather collective operation. This can potentially lead to a memory error.
+        Compute method executes ``compute_fn`` on zero rank process only and final result is broadcasted to
+        all processes.
 
     - ``update`` must receive output of the form ``(y_pred, y)`` or ``{'y_pred': y_pred, 'y': y}``.
 
@@ -30,19 +32,29 @@ class EpochMetric(Metric):
 
     Args:
         compute_fn (callable): a callable with the signature (`torch.tensor`, `torch.tensor`) takes as the input
-            `predictions` and `targets` and returns a scalar.
+            `predictions` and `targets` and returns a scalar. Input tensors will be on specified ``device``
+            (see arg below).
         output_transform (callable, optional): a callable that is used to transform the
             :class:`~ignite.engine.engine.Engine`'s ``process_function``'s output into the
             form expected by the metric. This can be useful if, for example, you have a multi-output model and
             you want to compute the metric with respect to one of the outputs.
         check_compute_fn (bool): if True, ``compute_fn`` is run on the first batch of data to ensure there are no
             issues. If issues exist, user is warned that there might be an issue with the ``compute_fn``.
+            Default, True.
+        device (str or torch.device, optional): optional device specification for internal storage.
 
     Warnings:
-        EpochMetricWarning: User is warned that there are issues with compute_fn on a batch of data processed.
+        EpochMetricWarning: User is warned that there are issues with ``compute_fn`` on a batch of data processed.
+        To disable the warning, set ``check_compute_fn=False``.
     """
 
-    def __init__(self, compute_fn: Callable, output_transform: Callable = lambda x: x, check_compute_fn: bool = True):
+    def __init__(
+        self,
+        compute_fn: Callable,
+        output_transform: Callable = lambda x: x,
+        check_compute_fn: bool = True,
+        device: Union[str, torch.device] = torch.device("cpu"),
+    ):
 
         if not callable(compute_fn):
             raise TypeError("Argument compute_fn should be callable.")
@@ -52,8 +64,9 @@ def __init__(self, compute_fn: Callable, output_transform: Callable = lambda x:
         self.compute_fn = compute_fn
         self._check_compute_fn = check_compute_fn
 
-        super(EpochMetric, self).__init__(output_transform=output_transform, device="cpu")
+        super(EpochMetric, self).__init__(output_transform=output_transform, device=device)
 
+    @reinit__is_reduced
     def reset(self) -> None:
         self._predictions = []
         self._targets = []
@@ -74,30 +87,32 @@ def _check_type(self, output):
         y_pred, y = output
         if len(self._predictions) < 1:
             return
-        dtype_preds = self._predictions[-1].type()
-        if dtype_preds != y_pred.type():
+        dtype_preds = self._predictions[-1].dtype
+        if dtype_preds != y_pred.dtype:
             raise ValueError(
                 "Incoherent types between input y_pred and stored predictions: "
-                "{} vs {}".format(dtype_preds, y_pred.type())
+                "{} vs {}".format(dtype_preds, y_pred.dtype)
             )
 
-        dtype_targets = self._targets[-1].type()
-        if dtype_targets != y.type():
+        dtype_targets = self._targets[-1].dtype
+        if dtype_targets != y.dtype:
             raise ValueError(
-                "Incoherent types between input y and stored targets: " "{} vs {}".format(dtype_targets, y.type())
+                "Incoherent types between input y and stored targets: " "{} vs {}".format(dtype_targets, y.dtype)
             )
 
+    @reinit__is_reduced
     def update(self, output: Sequence[torch.Tensor]) -> None:
         self._check_shape(output)
-        y_pred, y = output
+        y_pred, y = output[0].detach(), output[1].detach()
+
         if y_pred.ndimension() == 2 and y_pred.shape[1] == 1:
             y_pred = y_pred.squeeze(dim=-1)
 
         if y.ndimension() == 2 and y.shape[1] == 1:
             y = y.squeeze(dim=-1)
 
-        y_pred = y_pred.cpu().clone()
-        y = y.cpu().clone()
+        y_pred = y_pred.clone().to(self._device)
+        y = y.clone().to(self._device)
 
         self._check_type((y_pred, y))
         self._predictions.append(y_pred)
@@ -111,9 +126,30 @@ def update(self, output: Sequence[torch.Tensor]) -> None:
                 warnings.warn("Probably, there can be a problem with `compute_fn`:\n {}.".format(e), EpochMetricWarning)
 
     def compute(self) -> None:
+        if len(self._predictions) < 1 or len(self._targets) < 1:
+            raise NotComputableError("EpochMetric must have at least one example before it can be computed.")
+
         _prediction_tensor = torch.cat(self._predictions, dim=0)
         _target_tensor = torch.cat(self._targets, dim=0)
-        return self.compute_fn(_prediction_tensor, _target_tensor)
+
+        ws = idist.get_world_size()
+
+        if ws > 1 and not self._is_reduced:
+            # All gather across all processes
+            _prediction_tensor = idist.all_gather(_prediction_tensor)
+            _target_tensor = idist.all_gather(_target_tensor)
+        self._is_reduced = True
+
+        result = 0.0
+        if idist.get_rank() == 0:
+            # Run compute_fn on zero rank only
+            result = self.compute_fn(_prediction_tensor, _target_tensor)
+
+        if ws > 1:
+            # broadcast result to all processes
+            result = idist.broadcast(result, src=0)
+
+        return result
 
 
 class EpochMetricWarning(UserWarning):
diff --git a/ignite/metrics/metric.py b/ignite/metrics/metric.py
--- a/ignite/metrics/metric.py
+++ b/ignite/metrics/metric.py
@@ -122,7 +122,7 @@ class Metric(metaclass=ABCMeta):
             form expected by the metric. This can be useful if, for example, you have a multi-output model and
             you want to compute the metric with respect to one of the outputs.
             By default, metrics require the output as ``(y_pred, y)`` or ``{'y_pred': y_pred, 'y': y}``.
-        device (str of torch.device, optional): optional device specification for internal storage.
+        device (str or torch.device, optional): optional device specification for internal storage.
 
     """
 
@@ -458,14 +458,17 @@ def another_wrapper(self: Metric, *args, **kwargs) -> Callable:
                 raise RuntimeError(
                     "Decorator sync_all_reduce should be used on ignite.metric.Metric class methods only"
                 )
-
+            ws = idist.get_world_size()
             if len(attrs) > 0 and not self._is_reduced:
-                for attr in attrs:
-                    t = getattr(self, attr, None)
-                    if t is not None and idist.get_world_size() > 1:
-                        t = idist.all_reduce(t)
-                        self._is_reduced = True
-                        setattr(self, attr, t)
+                if ws > 1:
+                    for attr in attrs:
+                        t = getattr(self, attr, None)
+                        if t is not None:
+                            t = idist.all_reduce(t)
+                            self._is_reduced = True
+                            setattr(self, attr, t)
+                else:
+                    self._is_reduced = True
 
             return func(self, *args, **kwargs)
 


Test patch:
diff --git a/tests/ignite/metrics/test_epoch_metric.py b/tests/ignite/metrics/test_epoch_metric.py
--- a/tests/ignite/metrics/test_epoch_metric.py
+++ b/tests/ignite/metrics/test_epoch_metric.py
@@ -4,8 +4,9 @@
 import torch
 
 import ignite.distributed as idist
+from ignite.engine import Engine
 from ignite.metrics import EpochMetric
-from ignite.metrics.epoch_metric import EpochMetricWarning
+from ignite.metrics.epoch_metric import EpochMetricWarning, NotComputableError
 
 
 def test_epoch_metric_wrong_setup_or_input():
@@ -51,6 +52,12 @@ def compute_fn(y_preds, y_targets):
         output2 = (torch.rand(4, 3), torch.randint(0, 2, size=(4, 3)).to(torch.int32))
         em.update(output2)
 
+    with pytest.raises(
+        NotComputableError, match="EpochMetric must have at least one example before it can be computed"
+    ):
+        em = EpochMetric(compute_fn)
+        em.compute()
+
 
 def test_epoch_metric():
     def compute_fn(y_preds, y_targets):
@@ -136,14 +143,6 @@ def compute_fn(y_preds, y_targets):
         em.update(output1)
 
 
-def _test_warning():
-    def compute_fn(y_preds, y_targets):
-        return 0.0
-
-    with pytest.warns(RuntimeWarning, match="EpochMetric class does not support distributed setting"):
-        EpochMetric(compute_fn)
-
-
 def test_check_compute_fn():
     def compute_fn(y_preds, y_targets):
         raise Exception
@@ -159,47 +158,80 @@ def compute_fn(y_preds, y_targets):
     em.update(output1)
 
 
-@pytest.mark.distributed
-@pytest.mark.skipif(not idist.has_native_dist_support, reason="Skip if no native dist support")
-@pytest.mark.skipif(torch.cuda.device_count() < 1, reason="Skip if no GPU")
-def test_distrib_gpu(local_rank, distributed_context_single_node_nccl):
-    # Perform some ops otherwise, next tests fail
+def _test_distrib_integration(device=None):
+
+    if device is None:
+        device = idist.device()
+
+    rank = idist.get_rank()
+    torch.manual_seed(12)
+
+    n_iters = 60
+    s = 16
+    n_classes = 7
+
+    offset = n_iters * s
+    y_true = torch.randint(0, n_classes, size=(offset * idist.get_world_size(),), device=device)
+    y_preds = torch.rand(offset * idist.get_world_size(), n_classes, device=device)
 
-    if idist.get_world_size() > 1:
-        _test_warning()
+    def update(engine, i):
+        return (
+            y_preds[i * s + rank * offset : (i + 1) * s + rank * offset, :],
+            y_true[i * s + rank * offset : (i + 1) * s + rank * offset],
+        )
 
-    device = "cuda:{}".format(local_rank)
+    engine = Engine(update)
 
-    y = torch.rand(10, 12, device=device)
-    y = idist.all_gather(y)
-    assert isinstance(y, torch.Tensor)
+    def assert_data_fn(all_preds, all_targets):
+        assert all_preds.equal(y_preds), "{} vs {}".format(all_preds.shape, y_preds.shape)
+        assert all_targets.equal(y_true), "{} vs {}".format(all_targets.shape, y_true.shape)
+        return (all_preds.argmax(dim=1) == all_targets).sum().item()
+
+    ep_metric = EpochMetric(assert_data_fn, check_compute_fn=False, device=device)
+    ep_metric.attach(engine, "epm")
+
+    data = list(range(n_iters))
+    engine.run(data=data, max_epochs=3)
+    assert engine.state.metrics["epm"] == (y_preds.argmax(dim=1) == y_true).sum().item()
 
 
 @pytest.mark.distributed
 @pytest.mark.skipif(not idist.has_native_dist_support, reason="Skip if no native dist support")
-def test_distrib_cpu(local_rank, distributed_context_single_node_gloo):
-    # Perform some ops otherwise, next tests fail
+@pytest.mark.skipif(torch.cuda.device_count() < 1, reason="Skip if no GPU")
+def test_distrib_gpu(distributed_context_single_node_nccl):
+    _test_distrib_integration(device="cuda")
 
-    if idist.get_world_size() > 1:
-        _test_warning()
 
-    device = "cpu"
+@pytest.mark.distributed
+@pytest.mark.skipif(not idist.has_native_dist_support, reason="Skip if no native dist support")
+def test_distrib_cpu(distributed_context_single_node_gloo):
+    _test_distrib_integration(device="cpu")
 
-    y = torch.rand(10, 12, device=device)
-    y = idist.all_gather(y)
-    assert isinstance(y, torch.Tensor)
 
+@pytest.mark.tpu
+@pytest.mark.skipif("NUM_TPU_WORKERS" in os.environ, reason="Skip if NUM_TPU_WORKERS is in env vars")
+@pytest.mark.skipif(not idist.has_xla_support, reason="Skip if no PyTorch XLA package")
+def test_distrib_single_device_xla():
+    _test_distrib_integration()
 
-@pytest.mark.multinode_distributed
-@pytest.mark.skipif(not idist.has_native_dist_support, reason="Skip if no native dist support")
-@pytest.mark.skipif("MULTINODE_DISTRIB" not in os.environ, reason="Skip if not multi-node distributed")
-def test_multinode_distrib_cpu(distributed_context_multi_node_gloo):
 
-    _test_warning()
+def _test_distrib_xla_nprocs(index):
+    _test_distrib_integration()
 
 
-@pytest.mark.multinode_distributed
-@pytest.mark.skipif(not idist.has_native_dist_support, reason="Skip if no native dist support")
-@pytest.mark.skipif("GPU_MULTINODE_DISTRIB" not in os.environ, reason="Skip if not multi-node distributed")
-def test_multinode_distrib_gpu(distributed_context_multi_node_nccl):
-    _test_warning()
+@pytest.mark.tpu
+@pytest.mark.skipif("NUM_TPU_WORKERS" not in os.environ, reason="Skip if no NUM_TPU_WORKERS in env vars")
+@pytest.mark.skipif(not idist.has_xla_support, reason="Skip if no PyTorch XLA package")
+def test_distrib_xla_nprocs(xmp_executor):
+    n = int(os.environ["NUM_TPU_WORKERS"])
+    xmp_executor(_test_distrib_xla_nprocs, args=(), nprocs=n)
+
+
+@pytest.mark.distributed
+@pytest.mark.skipif(not idist.has_hvd_support, reason="Skip if no Horovod dist support")
+@pytest.mark.skipif("WORLD_SIZE" in os.environ, reason="Skip if launched as multiproc")
+def test_distrib_hvd(gloo_hvd_executor):
+
+    nproc = 4 if not torch.cuda.is_available() else torch.cuda.device_count()
+
+    gloo_hvd_executor(_test_distrib_integration, (None,), np=nproc, do_init=True)