Repository: pytorch/ignite

Issue:
Support list of tensors in Metrics
## 🚀 Feature

Related to discussion from here: https://github.com/Project-MONAI/MONAI/issues/2220#issuecomment-849250700

The feature request is to support list of tensors or scalars as input for metrics:
a) classification problem:
```python
acc = Accuracy()
output = (
    [torch.rand(10), torch.rand(10), torch.rand(10)],  # predictions
    [4, 5, 1],  # ground truth
)
acc.update(output)
print(acc.compute())
```

b) segmentation problem:
```python
acc = Accuracy()
output = (
    [torch.rand(10, 32, 32), torch.rand(10, 40, 30), torch.rand(10, 30, 24)],  # predictions
    [torch.randint(0, 10, size=(32, 32)), torch.randint(0, 10, size=(40, 30)),  torch.randint(0, 10, size=(30, 24))],  # ground truth
)
acc.update(output)
print(acc.compute())
```

c) Same for API with Engine.

cc @Nic-Ma 


Discussion:
sdesrozis: Just one remark but it looks good !
Nic-Ma: Hi @vfdev-5 @sdesrozis ,

Thanks for the quick PR, looks very good to me.
We also added support for `list of tensors` to MONAI last week: https://github.com/Project-MONAI/MONAI/blob/dev/monai/metrics/metric.py#L61
I have a quick question about the metrics support for **dict**:
Is it possible to make the global variable `required_output_keys` of ignite Metric to be an arg of `Metric.__init__()`?
We already received much feedback that our workflow programs have too many lambda functions and hard to understand. I am trying to reduce the `output_transform= lambda x: ...` in the code. If we can explicitly specify the `keys` of `y_pred` and `y`, maybe it's easier to use and understand?
CC @wyli

Thanks.

vfdev-5: Hi @Nic-Ma , 

> Is it possible to make the global variable required_output_keys of ignite Metric to be an arg of Metric.__init__()?

thanks for the feedback, `required_output_keys` is a class attribute and I think it would be better to set `required_output_keys` globally instead of for each metric: 
```python
from ignite.engine import Engine
from ignite.metrics import Accuracy, Metric

Metric.required_output_keys = ("my_y_pred", "my_y_true")

acc = Accuracy()
acc.required_output_keys

engine = Engine(lambda e, b: None)
engine.state.output = {"my_y_pred": torch.rand(4, 10), "my_y_true": torch.randint(0, 10, size=(4, ))}

acc.iteration_completed(engine)
acc.compute()
> 0.125
```
what do you think ?
Nic-Ma: Hi @vfdev-5 ,

Thanks for your proposal and sample code!
It can solve our problem, but maybe a little bit "hacking", especially when setting ignite args in the **JSON config** of NVIDIA Clara-Train SDK.
Our use case is something like:
```
"trainer": {
    "name": "SupervisedTrainer",
    "args": {
        ... ...
    }
},
post_transforms: [
    {
        "name": "CopyItemsD":
        "args": {"keys": "pred", "times": 1, "names": ["pred_inverted"]}
    },
    {
        "name": "InvertD",
        "args": {"keys": "pred_inverted"}
    },
    {
        "name": "SplitChannelD",
        "args": {"keys": "pred_inverted", "output_postfixes": ["background", "class1", "class2"]}
    }
]
"metrics": [
    {
        "name": "Accuracy",  # ignite metric
        "args": {"output_keys": ["pred", "label"]}
    },
    {
        "name": "MeanDice",  # subclass of ignite metric
        "args": {"output_keys": ["pred_inverted", "label_inverted"]}
    },
    {
        "name": "MeanDice",  # metric on class 1
        "args": {"output_keys": ["pred_inverted_class1", "label_inverted_class1"]}
    },
    {
        "name": "MeanDice",  # metric on class 2
        "args": {"output_keys": ["pred_inverted_class2", "label_inverted_class2"]}
    }
]
```
What do you think is the best way to support it?

Thanks.

vfdev-5: Thanks for the details, Nic ! That's definitely requires `output_transform` to handle each case. Let me think about that.

> We already received much feedback that our workflow programs have too many lambda functions and hard to understand. I am trying to reduce the output_transform= lambda x: ... in the code.

In this case, maybe, it is possible to introduce one function to deal with those required for metrics like:
```python
def from_dict(*keys):
    def wrapper(output: dict):
        return tuple(output[k] for k in keys)
    return wrapper

acc = Accuracy(output_transform=from_dict("pred_inverted", "label_inverted"))
```
what do you think ?
Nic-Ma: OK, thanks for the solution.
Let me try to add this API in MONAI utilities.

Thanks.


Patch:
diff --git a/ignite/metrics/metric.py b/ignite/metrics/metric.py
index 65c2afd9..ca7ce38e 100644
--- a/ignite/metrics/metric.py
+++ b/ignite/metrics/metric.py
@@ -2,7 +2,8 @@ import warnings
 from abc import ABCMeta, abstractmethod
 from collections.abc import Mapping
 from functools import wraps
-from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple, Union
+from numbers import Number
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Sequence, Tuple, Union
 
 import torch
 
@@ -279,8 +280,17 @@ class Metric(metaclass=ABCMeta):
         """Helper method to update metric's computation. It is automatically attached to the
         `engine` with :meth:`~ignite.metrics.metric.Metric.attach`.
 
+        Note:
+            ``engine.state.output`` is used to compute metric values.
+            The majority of implemented metrics accepts the following formats for ``engine.state.output``:
+            ``(y_pred, y)`` or ``{'y_pred': y_pred, 'y': y}``. ``y_pred`` and ``y`` can be torch tensors or
+            list of tensors/numbers if applicable.
+
         Args:
             engine: the engine to which the metric must be attached
+
+        .. versionchanged:: 0.4.5
+            ``y_pred`` and ``y`` can be torch tensors or list of tensors/numbers
         """
 
         output = self._output_transform(engine.state.output)
@@ -296,7 +306,20 @@ class Metric(metaclass=ABCMeta):
                     f"it should contain {self.required_output_keys} keys, but given {list(output.keys())}"
                 )
             output = tuple(output[k] for k in self.required_output_keys)
-        self.update(output)
+
+        if isinstance(output, Sequence) and all([_is_list_of_tensors_or_numbers(o) for o in output]):
+            if not (len(output) == 2 and len(output[0]) == len(output[1])):
+                raise ValueError(
+                    f"Output should have 2 items of the same length, "
+                    f"got {len(output)} and {len(output[0])}, {len(output[1])}"
+                )
+            for o1, o2 in zip(output[0], output[1]):
+                # o1 and o2 are list of tensors or numbers
+                tensor_o1 = _to_batched_tensor(o1)
+                tensor_o2 = _to_batched_tensor(o2, device=tensor_o1.device)
+                self.update((tensor_o1, tensor_o2))
+        else:
+            self.update(output)
 
     def completed(self, engine: Engine, name: str) -> None:
         """Helper method to compute metric's value and put into the engine. It is automatically attached to the
@@ -579,3 +602,13 @@ def reinit__is_reduced(func: Callable) -> Callable:
 
     setattr(wrapper, "_decorated", True)
     return wrapper
+
+
+def _is_list_of_tensors_or_numbers(x: Sequence[Union[torch.Tensor, float]]) -> bool:
+    return isinstance(x, Sequence) and all([isinstance(t, (torch.Tensor, Number)) for t in x])
+
+
+def _to_batched_tensor(x: Union[torch.Tensor, float], device: Optional[torch.device] = None) -> torch.Tensor:
+    if isinstance(x, torch.Tensor):
+        return x.unsqueeze(dim=0)
+    return torch.tensor([x,], device=device)


Test patch:
diff --git a/tests/ignite/metrics/test_accumulation.py b/tests/ignite/metrics/test_accumulation.py
index 1034c3cf..89e5dcec 100644
--- a/tests/ignite/metrics/test_accumulation.py
+++ b/tests/ignite/metrics/test_accumulation.py
@@ -320,11 +320,27 @@ def _test_distrib_geom_average(device):
             _test(idist.device())
 
 
+def _dist_mean(y_true):
+    y_true = idist.all_reduce(y_true) / idist.get_world_size()
+    if len(y_true.shape) > 2:
+        y_true = y_true.reshape(-1, y_true.shape[-1])
+    return y_true.mean(dim=0).cpu().numpy()
+
+
+def _dist_geom_mean(y_true):
+    log_y_true = torch.log(y_true)
+    log_y_true = idist.all_reduce(log_y_true)
+    if len(log_y_true.shape) > 2:
+        log_y_true = log_y_true.reshape(-1, log_y_true.shape[-1])
+    np_t = log_y_true.cpu().numpy()
+    return np.exp(np.mean(np_t, axis=0) / idist.get_world_size())
+
+
 def _test_distrib_integration(device):
-    def _test(metric_cls, true_result_fn, metric_device, tol=1e-5):
+    def _test(metric_cls, shape, true_result_fn, metric_device, tol=1e-5):
 
         size = 100
-        custom_variable = 10.0 + 5.0 * torch.rand(size, 12, dtype=torch.float64)
+        custom_variable = 10.0 + 5.0 * torch.rand(size, *shape, dtype=torch.float64)
         custom_variable = custom_variable.to(device)
 
         def update_fn(engine, batch):
@@ -336,10 +352,11 @@ def _test_distrib_integration(device):
         custom_var_mean.attach(engine, "agg_custom_var")
 
         state = engine.run([0] * size)
+        true_val = true_result_fn(custom_variable)
+        assert len(true_val) == shape[-1]
+
         np.testing.assert_almost_equal(
-            state.metrics["agg_custom_var"].cpu().numpy(),
-            true_result_fn(custom_variable),
-            decimal=int(np.log10(1.0 / tol)),
+            state.metrics["agg_custom_var"].cpu().numpy(), true_val, decimal=int(np.log10(1.0 / tol)),
         )
 
         size = 100
@@ -357,22 +374,13 @@ def _test_distrib_integration(device):
         state = engine.run([0] * size)
         assert state.metrics["agg_custom_var"] == pytest.approx(true_result_fn(custom_variable), abs=tol)
 
-    def _mean(y_true):
-        y_true = idist.all_reduce(y_true)
-        return y_true.mean(dim=0).cpu().numpy() / idist.get_world_size()
-
-    def _geom_mean(y_true):
-        log_y_true = torch.log(y_true)
-        log_y_true = idist.all_reduce(log_y_true)
-        np_t = log_y_true.cpu().numpy()
-        return np.exp(np.mean(np_t, axis=0) / idist.get_world_size())
-
     metric_devices = ["cpu"]
     if device.type != "xla":
         metric_devices.append(idist.device())
     for metric_device in metric_devices:
-        _test(Average, _mean, metric_device)
-        _test(GeometricAverage, _geom_mean, metric_device, tol=1e-4)
+        _test(Average, (12,), _dist_mean, metric_device)
+        _test(Average, (4, 12), _dist_mean, metric_device)
+        _test(GeometricAverage, (12,), _dist_geom_mean, metric_device, tol=1e-4)
 
 
 def _test_distrib_accumulator_device(device):
diff --git a/tests/ignite/metrics/test_accuracy.py b/tests/ignite/metrics/test_accuracy.py
index c997d2f6..ba38afb8 100644
--- a/tests/ignite/metrics/test_accuracy.py
+++ b/tests/ignite/metrics/test_accuracy.py
@@ -5,6 +5,7 @@ import torch
 from sklearn.metrics import accuracy_score
 
 import ignite.distributed as idist
+from ignite.engine import Engine
 from ignite.exceptions import NotComputableError
 from ignite.metrics import Accuracy
 
@@ -374,8 +375,6 @@ def _test_distrib_multilabel_input_NHW(device):
 
 def _test_distrib_integration_multiclass(device):
 
-    from ignite.engine import Engine
-
     rank = idist.get_rank()
     torch.manual_seed(12)
 
@@ -427,8 +426,6 @@ def _test_distrib_integration_multiclass(device):
 
 def _test_distrib_integration_multilabel(device):
 
-    from ignite.engine import Engine
-
     rank = idist.get_rank()
     torch.manual_seed(12)
 
@@ -500,6 +497,57 @@ def _test_distrib_accumulator_device(device):
         ), f"{type(acc._num_correct.device)}:{acc._num_correct.device} vs {type(metric_device)}:{metric_device}"
 
 
+def _test_distrib_integration_list_of_tensors_or_numbers(device):
+
+    rank = idist.get_rank()
+    torch.manual_seed(12)
+
+    def _test(n_epochs, metric_device):
+        metric_device = torch.device(metric_device)
+        n_iters = 80
+        s = 16
+        n_classes = 10
+
+        offset = n_iters * s
+        y_true = torch.randint(0, n_classes, size=(offset * idist.get_world_size(),)).to(device)
+        y_preds = torch.rand(offset * idist.get_world_size(), n_classes).to(device)
+
+        def update(_, i):
+            return (
+                [v for v in y_preds[i * s + rank * offset : (i + 1) * s + rank * offset, :]],
+                [v.item() for v in y_true[i * s + rank * offset : (i + 1) * s + rank * offset]],
+            )
+
+        engine = Engine(update)
+
+        acc = Accuracy(device=metric_device)
+        acc.attach(engine, "acc")
+
+        data = list(range(n_iters))
+        engine.run(data=data, max_epochs=n_epochs)
+
+        assert (
+            acc._num_correct.device == metric_device
+        ), f"{type(acc._num_correct.device)}:{acc._num_correct.device} vs {type(metric_device)}:{metric_device}"
+
+        assert "acc" in engine.state.metrics
+        res = engine.state.metrics["acc"]
+        if isinstance(res, torch.Tensor):
+            res = res.cpu().numpy()
+
+        true_res = accuracy_score(y_true.cpu().numpy(), torch.argmax(y_preds, dim=1).cpu().numpy())
+
+        assert pytest.approx(res) == true_res
+
+    metric_devices = ["cpu"]
+    if device.type != "xla":
+        metric_devices.append(idist.device())
+    for metric_device in metric_devices:
+        for _ in range(2):
+            _test(n_epochs=1, metric_device=metric_device)
+            _test(n_epochs=2, metric_device=metric_device)
+
+
 @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")
@@ -510,6 +558,7 @@ def test_distrib_nccl_gpu(distributed_context_single_node_nccl):
     _test_distrib_integration_multiclass(device)
     _test_distrib_integration_multilabel(device)
     _test_distrib_accumulator_device(device)
+    _test_distrib_integration_list_of_tensors_or_numbers(device)
 
 
 @pytest.mark.distributed
@@ -521,6 +570,7 @@ def test_distrib_gloo_cpu_or_gpu(distributed_context_single_node_gloo):
     _test_distrib_integration_multiclass(device)
     _test_distrib_integration_multilabel(device)
     _test_distrib_accumulator_device(device)
+    _test_distrib_integration_list_of_tensors_or_numbers(device)
 
 
 @pytest.mark.distributed
@@ -535,6 +585,7 @@ def test_distrib_hvd(gloo_hvd_executor):
     gloo_hvd_executor(_test_distrib_integration_multiclass, (device,), np=nproc, do_init=True)
     gloo_hvd_executor(_test_distrib_integration_multilabel, (device,), np=nproc, do_init=True)
     gloo_hvd_executor(_test_distrib_accumulator_device, (device,), np=nproc, do_init=True)
+    gloo_hvd_executor(_test_distrib_integration_list_of_tensors_or_numbers, (device,), np=nproc, do_init=True)
 
 
 @pytest.mark.tpu
@@ -546,6 +597,7 @@ def test_distrib_single_device_xla():
     _test_distrib_integration_multiclass(device)
     _test_distrib_integration_multilabel(device)
     _test_distrib_accumulator_device(device)
+    _test_distrib_integration_list_of_tensors_or_numbers(device)
 
 
 def _test_distrib_xla_nprocs(index):
@@ -554,6 +606,7 @@ def _test_distrib_xla_nprocs(index):
     _test_distrib_integration_multiclass(device)
     _test_distrib_integration_multilabel(device)
     _test_distrib_accumulator_device(device)
+    _test_distrib_integration_list_of_tensors_or_numbers(device)
 
 
 @pytest.mark.tpu
@@ -574,6 +627,7 @@ def test_multinode_distrib_gloo_cpu_or_gpu(distributed_context_multi_node_gloo):
     _test_distrib_integration_multiclass(device)
     _test_distrib_integration_multilabel(device)
     _test_distrib_accumulator_device(device)
+    _test_distrib_integration_list_of_tensors_or_numbers(device)
 
 
 @pytest.mark.multinode_distributed
@@ -586,3 +640,4 @@ def test_multinode_distrib_nccl_gpu(distributed_context_multi_node_nccl):
     _test_distrib_integration_multiclass(device)
     _test_distrib_integration_multilabel(device)
     _test_distrib_accumulator_device(device)
+    _test_distrib_integration_list_of_tensors_or_numbers(device)
diff --git a/tests/ignite/metrics/test_metric.py b/tests/ignite/metrics/test_metric.py
index 8664d6c2..51da80ed 100644
--- a/tests/ignite/metrics/test_metric.py
+++ b/tests/ignite/metrics/test_metric.py
@@ -934,3 +934,75 @@ def test_override_required_output_keys():
         (torch.rand(4, 10), torch.randint(0, 3, size=(4,))),
     ]
     evaluator.run(data)
+
+
+@pytest.mark.parametrize("shapes", [[(10,), ()], [(5, 32, 32), (5, 32, 32)],])
+def test_list_of_tensors_and_numbers(shapes):
+    def check_fn(output):
+        assert len(output) == 2
+        assert isinstance(output[0], torch.Tensor)
+        assert isinstance(output[1], torch.Tensor)
+        assert output[0].shape == (1,) + shapes[0]
+        assert output[1].shape == (1,) + shapes[1]
+
+    def get_data(gt_as_scalar=False):
+        return [
+            (
+                [torch.rand(shapes[0]) for _ in range(3 + i)],  # predictions
+                [
+                    torch.rand(shapes[1]).item() if gt_as_scalar else torch.rand(shapes[1]) for _ in range(3 + i)
+                ],  # ground truth
+            )
+            for i in range(5)
+        ]
+
+    class MyMetric(Metric):
+        def __init__(self, check_fn):
+            super(MyMetric, self).__init__()
+            self.check_fn = check_fn
+
+        def reset(self):
+            pass
+
+        def compute(self):
+            pass
+
+        def update(self, output):
+            self.check_fn(output)
+
+    engine = Engine(lambda e, b: b)
+    m = MyMetric(check_fn)
+    m.attach(engine, "m")
+
+    data = get_data()
+    engine.run(data)
+
+    if len(shapes[1]) == 0:
+        data = get_data(gt_as_scalar=True)
+        engine.run(data)
+
+
+def test_list_of_tensors_and_numbers_unsupported_output():
+    class MyMetric(Metric):
+        def reset(self):
+            pass
+
+        def compute(self):
+            pass
+
+        def update(self, output):
+            pass
+
+    engine = Engine(lambda e, b: ([0, 1, 2], [0, 1, 2], [0, 1, 2]))
+    m = MyMetric()
+    m.attach(engine, "m")
+
+    with pytest.raises(ValueError, match=r"Output should have 2 items of the same length"):
+        engine.run([0] * 10)
+
+    engine = Engine(lambda e, b: ([0, 1, 2], [0, 1, 2, 4]))
+    m = MyMetric()
+    m.attach(engine, "m")
+
+    with pytest.raises(ValueError, match=r"Output should have 2 items of the same length"):
+        engine.run([0] * 10)