Repository: Lightning-AI/torchmetrics

Issue:
Nested MetricCollection
## 🚀 Feature

Allowing MetricCollection to contain other MetricCollections for even less boilerplate and coding effort.

### Motivation

When using both micro and macro metrics I either need to call two different collections or provide the collection with a dictionary of names, which could easily be created by appending a post-/prefix to the metric name.

### Pitch

I envision something like:
```python
metrics = MetricCollection([
    MetricCollection([
        StatScores(reduce='macro'),
        F1Score(reduce='macro')
    ], postfix='_macro'),
    MetricCollection([
        StatScores(reduce='micro'),
        F1Score(reduce='micro')
    ], postfix='_micro'),
], prefix='valmetrics/')
```

I want a nested collection to be treated like any other metric. Ideally nested collections would also make use of `compute_groups`.

### Alternatives

The same functionality can already be achieved by:

```python
metrics = MetricCollection({
    'statscores_macro': StatScores(reduce='macro'),
    'f1_macro': F1Score(reduce='macro'),
    'statscores_micro': StatScores(reduce='micro'),
    'f1_micro': F1Score(reduce='micro')
}, prefix='valmetrics/')
```

But I think accepting nested collections would make for a more convenient interface.
(It would probably be the easiest solution to just generate this version if presented with nested collections.)


Discussion:
Hi! thanks for your contribution!, great first issue!
cc: @Borda and @justusschock 
Seems like a good idea to me, but it will not work with the compute group feature we just implemented.

Resolving patch:
diff --git a/torchmetrics/collections.py b/torchmetrics/collections.py
--- a/torchmetrics/collections.py
+++ b/torchmetrics/collections.py
@@ -54,6 +54,10 @@ class name as key for the output dict.
             this argument is ``True`` which enables this feature. Set this argument to `False` for disabling
             this behaviour. Can also be set to a list of list of metrics for setting the compute groups yourself.
 
+    .. note::
+        Metric collections can be nested at initilization (see last example) but the output of the collection will
+        still be a single flattened dictionary combining the prefix and postfix arguments from the nested collection.
+
     Raises:
         ValueError:
             If one of the elements of ``metrics`` is not an instance of ``pl.metrics.Metric``.
@@ -104,6 +108,23 @@ class name as key for the output dict.
         ... )
         >>> pprint(metrics(preds, target))
         {'Accuracy': tensor(0.1250), 'MeanSquaredError': tensor(2.3750), 'Precision': tensor(0.0667)}
+
+    Example (nested metric collections):
+        >>> metrics = MetricCollection([
+        ...     MetricCollection([
+        ...         Accuracy(num_classes=3, average='macro'),
+        ...         Precision(num_classes=3, average='macro')
+        ...     ], postfix='_macro'),
+        ...     MetricCollection([
+        ...         Accuracy(num_classes=3, average='micro'),
+        ...         Precision(num_classes=3, average='micro')
+        ...     ], postfix='_micro'),
+        ... ], prefix='valmetrics/')
+        >>> pprint(metrics(preds, target))  # doctest: +NORMALIZE_WHITESPACE
+        {'valmetrics/Accuracy_macro': tensor(0.1111),
+        'valmetrics/Accuracy_micro': tensor(0.1250),
+        'valmetrics/Precision_macro': tensor(0.0667),
+        'valmetrics/Precision_micro': tensor(0.1250)}
     """
 
     _groups: Dict[int, List[str]]
@@ -195,6 +216,10 @@ def _merge_compute_groups(self) -> None:
     @staticmethod
     def _equal_metric_states(metric1: Metric, metric2: Metric) -> bool:
         """Check if the metric state of two metrics are the same."""
+        # empty state
+        if len(metric1._defaults) == 0 or len(metric2._defaults) == 0:
+            return False
+
         if metric1._defaults.keys() != metric2._defaults.keys():
             return False
 
@@ -280,19 +305,31 @@ def add_metrics(
             # Make sure that metrics are added in deterministic order
             for name in sorted(metrics.keys()):
                 metric = metrics[name]
-                if not isinstance(metric, Metric):
+                if not isinstance(metric, (Metric, MetricCollection)):
                     raise ValueError(
-                        f"Value {metric} belonging to key {name} is not an instance of `pl.metrics.Metric`"
+                        f"Value {metric} belonging to key {name} is not an instance of"
+                        " `torchmetrics.Metric` or `torchmetrics.MetricCollection`"
                     )
-                self[name] = metric
+                if isinstance(metric, Metric):
+                    self[name] = metric
+                else:
+                    for k, v in metric.items(keep_base=False):
+                        self[f"{name}_{k}"] = v
         elif isinstance(metrics, Sequence):
             for metric in metrics:
-                if not isinstance(metric, Metric):
-                    raise ValueError(f"Input {metric} to `MetricCollection` is not a instance of `pl.metrics.Metric`")
-                name = metric.__class__.__name__
-                if name in self:
-                    raise ValueError(f"Encountered two metrics both named {name}")
-                self[name] = metric
+                if not isinstance(metric, (Metric, MetricCollection)):
+                    raise ValueError(
+                        f"Input {metric} to `MetricCollection` is not a instance of"
+                        " `torchmetrics.Metric` or `torchmetrics.MetricCollection`"
+                    )
+                if isinstance(metric, Metric):
+                    name = metric.__class__.__name__
+                    if name in self:
+                        raise ValueError(f"Encountered two metrics both named {name}")
+                    self[name] = metric
+                else:
+                    for k, v in metric.items(keep_base=False):
+                        self[k] = v
         else:
             raise ValueError("Unknown input to MetricCollection.")
 


Test patch:
diff --git a/tests/bases/test_collections.py b/tests/bases/test_collections.py
--- a/tests/bases/test_collections.py
+++ b/tests/bases/test_collections.py
@@ -401,3 +401,36 @@ def test_error_on_wrong_specified_compute_groups():
         MetricCollection(
             ConfusionMatrix(3), Recall(3), Precision(3), compute_groups=[["ConfusionMatrix"], ["Recall", "Accuracy"]]
         )
+
+
+@pytest.mark.parametrize(
+    "input_collections",
+    [
+        [
+            MetricCollection(
+                [Accuracy(num_classes=3, average="macro"), Precision(num_classes=3, average="macro")], prefix="macro_"
+            ),
+            MetricCollection(
+                [Accuracy(num_classes=3, average="micro"), Precision(num_classes=3, average="micro")], prefix="micro_"
+            ),
+        ],
+        {
+            "macro": MetricCollection(
+                [Accuracy(num_classes=3, average="macro"), Precision(num_classes=3, average="macro")]
+            ),
+            "micro": MetricCollection(
+                [Accuracy(num_classes=3, average="micro"), Precision(num_classes=3, average="micro")]
+            ),
+        },
+    ],
+)
+def test_nested_collections(input_collections):
+    """Test that nested collections gets flattened to a single collection."""
+    metrics = MetricCollection(input_collections, prefix="valmetrics/")
+    preds = torch.randn(10, 3).softmax(dim=-1)
+    target = torch.randint(3, (10,))
+    val = metrics(preds, target)
+    assert "valmetrics/macro_Accuracy" in val
+    assert "valmetrics/macro_Precision" in val
+    assert "valmetrics/micro_Accuracy" in val
+    assert "valmetrics/micro_Precision" in val