Repository: Lightning-AI/torchmetrics

Issue:
Implement pre and postfix for Classwise Wrapper
## 🚀 Feature

A `prefix` and `postfix` parameter to change the metric names of `ClasswiseWrapper` class.

### Motivation

I came up with this feature while implementing multiple classification metrics for my Lightning Module which I then want to log via `log_dict` to Weight and Biases:

In `__init__` method of my lightning module:
```py
def __init__(self, ...):
    ...
    macro_metrics = MetricCollection(
            {
                "accuracy": MulticlassAccuracy(num_classes=M, average="macro"),
                "precision": MulticlassPrecision(num_classes=M, average="macro"),
                "recall": MulticlassRecall(num_classes=M, average="macro"),
                "f1": MulticlassF1Score(num_classes=M, average="macro"),
                "jaccard": MulticlassJaccardIndex(num_classes=M, average="macro"),
            },
            prefix="macro-",
        )

        micro_metrics = MetricCollection(
            {
                "accuracy": MulticlassAccuracy(num_classes=M, average="micro"),
                "precision": MulticlassPrecision(num_classes=M, average="micro"),
                "recall": MulticlassRecall(num_classes=M, average="micro"),
                "f1": MulticlassF1Score(num_classes=M, average="micro"),
                "jaccard": MulticlassJaccardIndex(num_classes=M, average="micro"),
            },
            prefix="micro-",
        )

        classwise_metrics = MetricCollection(
            {
                "accuracy": ClasswiseWrapper(MulticlassAccuracy(num_classes=M, average=None), labels=self.labels),
                "precision": ClasswiseWrapper(MulticlassPrecision(num_classes=M, average=None), labels=self.labels),
                "recall": ClasswiseWrapper(MulticlassRecall(num_classes=M, average=None), labels=self.labels),
                "f1": ClasswiseWrapper(MulticlassF1Score(num_classes=M, average=None), labels=self.labels),
                "jaccard": ClasswiseWrapper(MulticlassJaccardIndex(num_classes=M, average=None), labels=self.labels),
            }
        )
        self.metrics = MetricCollection([macro_metrics, micro_metrics, classwise_metrics], prefix="val/")
```

In validation step:

```py
def validation_step(self, ...):
    ...
    self.metrics.update(preds, target)
    self.log_dict(self.metrics, on_epoch=True)
```

As of right now this would result in Metrics like: "val/multiclassaccuracy_0". However, for aesthetics and sorting/grouping reasons I would like to have metrics named: "val/accuracy/class_0" or similar. I know I could manually log each item but this would result in such big boilerplate that the reason to use torchmetrics would be gone.

### Pitch

Currently the name of the wrapper is determined from the class name of the metric:

```py
>>> import torch
>>> from torchmetrics import ClasswiseWrapper
>>> from torchmetrics.classification import MulticlassAccuracy
>>> metric = ClasswiseWrapper(
...    MulticlassAccuracy(num_classes=3, average=None),
...    labels=["horse", "fish", "dog"]
... )
>>> preds = torch.randn(10, 3).softmax(dim=-1)
>>> target = torch.randint(3, (10,))
>>> metric(preds, target)  
{'multiclassaccuracy_horse': tensor(0.3333),
'multiclassaccuracy_fish': tensor(0.6667),
'multiclassaccuracy_dog': tensor(0.)}
```

With a prefix (and/or a postfix) it would be possible to customize the output name of the metric, which would be relevant for e.g. logging via `log_dict` of a Lightning Module:

```py
>>> import torch
>>> from torchmetrics import ClasswiseWrapper
>>> from torchmetrics.classification import MulticlassAccuracy
>>> metric = ClasswiseWrapper(
...    MulticlassAccuracy(num_classes=3, average=None),
...    labels=["horse", "fish", "dog"],
...    prefix="accuracy-"
... )
>>> preds = torch.randn(10, 3).softmax(dim=-1)
>>> target = torch.randint(3, (10,))
>>> metric(preds, target)  
{'accuracy-horse': tensor(0.3333),
'accuracy-fish': tensor(0.6667),
'accuracy-dog': tensor(0.)}
```

### Alternatives

1. A `name` parameter which would replace `this.metric.__class__.__name__` if set. This would be an easier implementation but wouldn't give the user the amount of customization options he may need.
2. Use the dict key as a name when passed to a `MetricCollection`. This could be an additional feature. Needs to be discussed what should have prio: the `prefix`/`postfix` parameter or the keyname in the `MetricCollection` dict.

### Additional context

I already opened a PR for this: #1866

---

### Important Note

I think that Lightning can't handle `ClasswiseWrapper` as of right now when logging via `self.log_dict(self.metrics, on_epoch=True)` since it raises an error in the argument validation process: `ValueError: The ".compute()" return of the metric logged as 'val/accuracy' must be a tensor. Found {'multiclassaccuracy_0': tensor(0.6000), 'multiclassaccuracy_1':  tensor(0.6000)}` The reason for this is probably that the `ClasswiseWrapper` counts as a single item when the `MetricCollection.items()` function is called:

```
│ /home/tobias/Repositories/eSPA-Python/.venv/lib/python3.10/site-packages/lightning/pytorch/train │
│ er/connectors/logger_connector/result.py:474 in metrics                                          │
│                                                                                                  │
│   471 │   │                                                                                      │
│   472 │   │   for _, result_metric in self.valid_items():                                        │
│   473 │   │   │   # extract forward_cache or computed from the _ResultMetric                     │
│ ❱ 474 │   │   │   value = self._get_cache(result_metric, on_step)                                │
│   475 │   │   │   if not isinstance(value, Tensor):                                              │
│   476 │   │   │   │   continue                                                                   │
│   477                                                                                            │

│ /home/tobias/Repositories/eSPA-Python/.venv/lib/python3.10/site-packages/lightning/pytorch/train │
│ er/connectors/logger_connector/result.py:445 in _get_cache                                       │
│                                                                                                  │
│   442 │   │                                                                                      │
│   443 │   │   if cache is not None:                                                              │
│   444 │   │   │   if not isinstance(cache, Tensor):                                              │
│ ❱ 445 │   │   │   │   raise ValueError(                                                          │
│   446 │   │   │   │   │   f"The `.compute()` return of the metric logged as {result_metric.met   │
│   447 │   │   │   │   │   f" Found {cache}"                                                      │
│   448 │   │   │   │   )                                                                          │
```


Discussion:

Resolving patch:
diff --git a/src/torchmetrics/wrappers/classwise.py b/src/torchmetrics/wrappers/classwise.py
--- a/src/torchmetrics/wrappers/classwise.py
+++ b/src/torchmetrics/wrappers/classwise.py
@@ -33,8 +33,12 @@ class ClasswiseWrapper(Metric):
         metric: base metric that should be wrapped. It is assumed that the metric outputs a single
             tensor that is split along the first dimension.
         labels: list of strings indicating the different classes.
+        prefix: string that is prepended to the metric names.
+        postfix: string that is appended to the metric names.
+
+    Example::
+        Basic example where the ouput of a metric is unwrapped into a dictionary with the class index as keys:
 
-    Example:
         >>> import torch
         >>> _ = torch.manual_seed(42)
         >>> from torchmetrics.wrappers import ClasswiseWrapper
@@ -47,7 +51,29 @@ class ClasswiseWrapper(Metric):
         'multiclassaccuracy_1': tensor(0.7500),
         'multiclassaccuracy_2': tensor(0.)}
 
-    Example (labels as list of strings):
+    Example::
+        Using custom name via prefix and postfix:
+
+        >>> import torch
+        >>> _ = torch.manual_seed(42)
+        >>> from torchmetrics.wrappers import ClasswiseWrapper
+        >>> from torchmetrics.classification import MulticlassAccuracy
+        >>> metric_pre = ClasswiseWrapper(MulticlassAccuracy(num_classes=3, average=None), prefix="acc-")
+        >>> metric_post = ClasswiseWrapper(MulticlassAccuracy(num_classes=3, average=None), postfix="-acc")
+        >>> preds = torch.randn(10, 3).softmax(dim=-1)
+        >>> target = torch.randint(3, (10,))
+        >>> metric_pre(preds, target)  # doctest: +NORMALIZE_WHITESPACE
+        {'acc-0': tensor(0.5000),
+         'acc-1': tensor(0.7500),
+         'acc-2': tensor(0.)}
+        >>> metric_post(preds, target)  # doctest: +NORMALIZE_WHITESPACE
+        {'0-acc': tensor(0.5000),
+         '1-acc': tensor(0.7500),
+         '2-acc': tensor(0.)}
+
+    Example::
+        Providing labels as a list of strings:
+
         >>> from torchmetrics.wrappers import ClasswiseWrapper
         >>> from torchmetrics.classification import MulticlassAccuracy
         >>> metric = ClasswiseWrapper(
@@ -61,7 +87,10 @@ class ClasswiseWrapper(Metric):
         'multiclassaccuracy_fish': tensor(0.6667),
         'multiclassaccuracy_dog': tensor(0.)}
 
-    Example (in metric collection):
+    Example::
+        Classwise can also be used in combination with :class:`~torchmetrics.MetricCollection`. In this case, everything
+        will be flattened into a single dictionary:
+
         >>> from torchmetrics import MetricCollection
         >>> from torchmetrics.wrappers import ClasswiseWrapper
         >>> from torchmetrics.classification import MulticlassAccuracy, MulticlassRecall
@@ -81,21 +110,43 @@ class ClasswiseWrapper(Metric):
          'multiclassrecall_dog': tensor(0.4000)}
     """
 
-    def __init__(self, metric: Metric, labels: Optional[List[str]] = None) -> None:
+    def __init__(
+        self,
+        metric: Metric,
+        labels: Optional[List[str]] = None,
+        prefix: Optional[str] = None,
+        postfix: Optional[str] = None,
+    ) -> None:
         super().__init__()
         if not isinstance(metric, Metric):
             raise ValueError(f"Expected argument `metric` to be an instance of `torchmetrics.Metric` but got {metric}")
+        self.metric = metric
+
         if labels is not None and not (isinstance(labels, list) and all(isinstance(lab, str) for lab in labels)):
             raise ValueError(f"Expected argument `labels` to either be `None` or a list of strings but got {labels}")
-        self.metric = metric
         self.labels = labels
+
+        if prefix is not None and not isinstance(prefix, str):
+            raise ValueError(f"Expected argument `prefix` to either be `None` or a string but got {prefix}")
+        self.prefix = prefix
+
+        if postfix is not None and not isinstance(postfix, str):
+            raise ValueError(f"Expected argument `postfix` to either be `None` or a string but got {postfix}")
+        self.postfix = postfix
+
         self._update_count = 1
 
     def _convert(self, x: Tensor) -> Dict[str, Any]:
-        name = self.metric.__class__.__name__.lower()
+        # Will set the class name as prefix if neither prefix nor postfix is given
+        if not self.prefix and not self.postfix:
+            prefix = f"{self.metric.__class__.__name__.lower()}_"
+            postfix = ""
+        else:
+            prefix = self.prefix or ""
+            postfix = self.postfix or ""
         if self.labels is None:
-            return {f"{name}_{i}": val for i, val in enumerate(x)}
-        return {f"{name}_{lab}": val for lab, val in zip(self.labels, x)}
+            return {f"{prefix}{i}{postfix}": val for i, val in enumerate(x)}
+        return {f"{prefix}{lab}{postfix}": val for lab, val in zip(self.labels, x)}
 
     def forward(self, *args: Any, **kwargs: Any) -> Any:
         """Calculate on batch and accumulate to global state."""


Test patch:
diff --git a/tests/unittests/wrappers/test_classwise.py b/tests/unittests/wrappers/test_classwise.py
--- a/tests/unittests/wrappers/test_classwise.py
+++ b/tests/unittests/wrappers/test_classwise.py
@@ -13,6 +13,12 @@ def test_raises_error_on_wrong_input():
     with pytest.raises(ValueError, match="Expected argument `labels` to either be `None` or a list of strings.*"):
         ClasswiseWrapper(MulticlassAccuracy(num_classes=3), "hest")
 
+    with pytest.raises(ValueError, match="Expected argument `prefix` to either be `None` or a string.*"):
+        ClasswiseWrapper(MulticlassAccuracy(num_classes=3), prefix=1)
+
+    with pytest.raises(ValueError, match="Expected argument `postfix` to either be `None` or a string.*"):
+        ClasswiseWrapper(MulticlassAccuracy(num_classes=3), postfix=1)
+
 
 def test_output_no_labels():
     """Test that wrapper works with no label input."""
@@ -54,6 +60,38 @@ def test_output_with_labels():
             assert val[f"multiclassaccuracy_{lab}"] == val_base[i]
 
 
+def test_output_with_prefix():
+    """Test that wrapper works with prefix."""
+    base = MulticlassAccuracy(num_classes=3, average=None)
+    metric = ClasswiseWrapper(MulticlassAccuracy(num_classes=3, average=None), prefix="pre_")
+    for _ in range(2):
+        preds = torch.randn(20, 3).softmax(dim=-1)
+        target = torch.randint(3, (20,))
+        val = metric(preds, target)
+        val_base = base(preds, target)
+        assert isinstance(val, dict)
+        assert len(val) == 3
+        for i in range(3):
+            assert f"pre_{i}" in val
+            assert val[f"pre_{i}"] == val_base[i]
+
+
+def test_output_with_postfix():
+    """Test that wrapper works with postfix."""
+    base = MulticlassAccuracy(num_classes=3, average=None)
+    metric = ClasswiseWrapper(MulticlassAccuracy(num_classes=3, average=None), postfix="_post")
+    for _ in range(2):
+        preds = torch.randn(20, 3).softmax(dim=-1)
+        target = torch.randint(3, (20,))
+        val = metric(preds, target)
+        val_base = base(preds, target)
+        assert isinstance(val, dict)
+        assert len(val) == 3
+        for i in range(3):
+            assert f"{i}_post" in val
+            assert val[f"{i}_post"] == val_base[i]
+
+
 @pytest.mark.parametrize("prefix", [None, "pre_"])
 @pytest.mark.parametrize("postfix", [None, "_post"])
 def test_using_metriccollection(prefix, postfix):