Repository: Lightning-AI/torchmetrics

Issue:
Support for MetricTracker of MultioutputWrapper
## 🚀 Feature

Currently, `MetricTracker` works in the following cases:

- The base metric is of type `Metric`, and its `compute` method returns a single `Tensor`
- The base metric is of type `MetricCollection` and its `compute` method returns a `dict[str, Tensor]`

I would like to add support for two other cases which I came across in my projects:

- The base metric is of type `MultioutputWrapper`. Note that `MultioutputWrapper` is a subclass of `Metric`, but it is not supported by `MetricTracker` because its `compute` method returns a `List[Tensor]` rather than a single `Tensor`.
- The base metric is a `MetricCollection` of `MultioutputWrapper`. In thise case the `compute` method returns a `dict[str, List[Tensor]]`.

### Motivation

- I use `MultioutputWrapper` quite a lot for multi-regression problems. If I want to track its evolution over the epochs, I need to use a `MetricTracker` of `MultioutputWrapper`. 
- In some cases, I want to track multiple metrics (like MSE, MAE) for each of the outputs of my models, so I need to use a `MetricTracker` of `MetricCollection` of `MultioutputWrapper`.

### Pitch

After looking at the code, here is my proposed solution:
- Change `MetricTracker`'s `compute_all` method to support both of the new cases mentioned above.
- Change `MetricTracker`s `best_metric` method to support both of the new cases mentioned above.

### Alternatives

Add a boolean parameter to `MetricTracker` to make it return a `Tensor` of higher dimension rather than a `List[Tensor]`.

I don't really know which solution is best so I'll try to implement both in 2 different branches.


Discussion:
Hi! thanks for your contribution!, great first issue!
After looking at the code and starting to implement a solution, I completely updated the issue's description to make it much more clear. 
Hi @ValerianRey,
Just wanted to check in how much process you had on this issue? it seems you commit a lot of code that references this issue. Is it close to done or do you want me to take a stab at it :)
I added some new tests and the solution that makes them work when using a MetricTracker of MultioutputWrapper or a MetricTracker of MetricCollection of MultioutputWrapper. 

Everything works, but I don't like it this way. The MetricTracker class becomes too complex for what it does. The return type of best_metric becomes even more complex for example: 12 possible return types.

I think this class needs some deep changes before we invest too much development into it (some changes that would not make it backward-compatible).
For now, we could just specify in the documentation  which use-cases of MetricTracker work.

Resolving patch:
diff --git a/src/torchmetrics/wrappers/tracker.py b/src/torchmetrics/wrappers/tracker.py
--- a/src/torchmetrics/wrappers/tracker.py
+++ b/src/torchmetrics/wrappers/tracker.py
@@ -35,6 +35,11 @@ class MetricTracker(ModuleList):
         -``MetricTracker.compute_all()``: get the metric value for all steps
         -``MetricTracker.best_metric()``: returns the best value
 
+    Out of the box, this wrapper class fully supports that the base metric being tracked is a single `Metric`, a
+    `MetricCollection` or another `MetricWrapper` wrapped around a metric. However, multiple layers of nesting, such
+    as using a `Metric` inside a `MetricWrapper` inside a `MetricCollection` is not fully supported, especially the
+    `.best_metric` method that cannot auto compute the best metric and index for such nested structures.
+
     Args:
         metric: instance of a ``torchmetrics.Metric`` or ``torchmetrics.MetricCollection``
             to keep track of at each timestep.
@@ -135,12 +140,13 @@ def compute(self) -> Any:
         self._check_for_increment("compute")
         return self[-1].compute()
 
-    def compute_all(self) -> Union[Tensor, Dict[str, Tensor]]:
+    def compute_all(self) -> Any:
         """Compute the metric value for all tracked metrics.
 
         Return:
-            Either a single tensor if the tracked base object is a single metric, else if a metric collection is
-            provide a dict of tensors will be returned
+            By default will try stacking the results from all increaments into a single tensor if the tracked base
+            object is a single metric. If a metric collection is provided a dict of stacked tensors will be returned.
+            If the stacking process fails a list of the computed results will be returned.
 
         Raises:
             ValueError:
@@ -149,10 +155,15 @@ def compute_all(self) -> Union[Tensor, Dict[str, Tensor]]:
         self._check_for_increment("compute_all")
         # The i!=0 accounts for the self._base_metric should be ignored
         res = [metric.compute() for i, metric in enumerate(self) if i != 0]
-        if isinstance(self._base_metric, MetricCollection):
-            keys = res[0].keys()
-            return {k: torch.stack([r[k] for r in res], dim=0) for k in keys}
-        return torch.stack(res, dim=0)
+        try:
+            if isinstance(res[0], dict):
+                keys = res[0].keys()
+                return {k: torch.stack([r[k] for r in res], dim=0) for k in keys}
+            if isinstance(res[0], list):
+                return torch.stack([torch.stack(r, dim=0) for r in res], 0)
+            return torch.stack(res, dim=0)
+        except TypeError:  # fallback solution to just return as it is if we cannot succesfully stack
+            return res
 
     def reset(self) -> None:
         """Reset the current metric being tracked."""
@@ -181,7 +192,6 @@ def best_metric(
         Returns:
             Either a single value or a tuple, depends on the value of ``return_step`` and the object being tracked.
 
-
             - If a single metric is being tracked and ``return_step=False`` then a single tensor will be returned
             - If a single metric is being tracked and ``return_step=True`` then a 2-element tuple will be returned,
               where the first value is optimal value and second value is the corresponding optimal step
@@ -193,12 +203,23 @@ def best_metric(
               of the first dict being the optimal values and the values of the second dict being the optimal step
 
             In addtion the value in all cases may be ``None`` if the underlying metric does have a proper defined way
-            of being optimal.
+            of being optimal or in the case where a nested structure of metrics are being tracked.
         """
+        res = self.compute_all()
+        if isinstance(res, list):
+            rank_zero_warn(
+                "Encounted nested structure. You are probably using a metric collection inside a metric collection, or"
+                " a metric wrapper inside a metric collection, which is not supported by `.best_metric()` method."
+                "Returning `None` instead. Please consider "
+            )
+            if return_step:
+                return None, None
+            return None
+
         if isinstance(self._base_metric, Metric):
             fn = torch.max if self.maximize else torch.min
             try:
-                value, idx = fn(self.compute_all(), 0)  # type: ignore[arg-type]
+                value, idx = fn(res, 0)
                 if return_step:
                     return value.item(), idx.item()
                 return value.item()
@@ -214,10 +235,9 @@ def best_metric(
                 return None
 
         else:  # this is a metric collection
-            res = self.compute_all()
             maximize = self.maximize if isinstance(self.maximize, list) else len(res) * [self.maximize]
             value, idx = {}, {}
-            for i, (k, v) in enumerate(res.items()):  # type: ignore[union-attr]
+            for i, (k, v) in enumerate(res.items()):
                 try:
                     fn = torch.max if maximize[i] else torch.min
                     out = fn(v, 0)
@@ -238,4 +258,4 @@ def best_metric(
     def _check_for_increment(self, method: str) -> None:
         """Check that a metric that can be updated/used for computations has been intialized."""
         if not self._increment_called:
-            raise ValueError(f"`{method}` cannot be called before `.increment()` has been called")
+            raise ValueError(f"`{method}` cannot be called before `.increment()` has been called.")


Test patch:
diff --git a/tests/unittests/wrappers/test_tracker.py b/tests/unittests/wrappers/test_tracker.py
--- a/tests/unittests/wrappers/test_tracker.py
+++ b/tests/unittests/wrappers/test_tracker.py
@@ -15,7 +15,7 @@
 import pytest
 import torch
 
-from torchmetrics import MeanAbsoluteError, MeanSquaredError, MetricCollection
+from torchmetrics import MeanAbsoluteError, MeanSquaredError, MetricCollection, MultioutputWrapper
 from torchmetrics.classification import (
     MulticlassAccuracy,
     MulticlassConfusionMatrix,
@@ -184,3 +184,33 @@ def test_best_metric_for_not_well_defined_metric_collection(base_metric):
         else:
             assert best is None
             assert idx is None
+
+
+@pytest.mark.parametrize(
+    ("input_to_tracker", "assert_type"),
+    [
+        (MultioutputWrapper(MeanSquaredError(), num_outputs=2), torch.Tensor),
+        (  # nested version
+            MetricCollection(
+                {
+                    "mse": MultioutputWrapper(MeanSquaredError(), num_outputs=2),
+                    "mae": MultioutputWrapper(MeanAbsoluteError(), num_outputs=2),
+                }
+            ),
+            list,
+        ),
+    ],
+)
+def test_metric_tracker_and_collection_multioutput(input_to_tracker, assert_type):
+    """Check that MetricTracker support wrapper inputs and nested structures."""
+    tracker = MetricTracker(input_to_tracker)
+    for _ in range(5):
+        tracker.increment()
+        for _ in range(5):
+            preds, target = torch.randn(100, 2), torch.randn(100, 2)
+            tracker.update(preds, target)
+    all_res = tracker.compute_all()
+    assert isinstance(all_res, assert_type)
+    best_metric, which_epoch = tracker.best_metric(return_step=True)
+    assert best_metric is None
+    assert which_epoch is None