Repository: Lightning-AI/torchmetrics

Issue:
`log_dict` is missing `prefix` parameter that is mentioned in docs
## 🐛 Bug

[`MetricCollection` documentation](https://pytorch-lightning.readthedocs.io/en/1.2.0/extensions/metrics.html#metriccollection) mentions using `self.log_dict(self.train_metrics, on_step=True, on_epoch=False, prefix='train')`. The `prefix` parameter doesn't seem to be present in `log_dict` function header.


### Expected behavior

`prefix` is most likely usable in this context, so this feature should be implemented. If not - the documentation should be fixed.

### Environment

 - Lightning Version: 1.2.0



Discussion:
@carmocca are you in favor of adding the argument to `log_dict` or fix the documentation?
Fix the documentation. Does `MetricCollection` have a mechanism to set a prefix? If it doesn't, then the argument should be added there
@carmocca I agree this would be a nice addition and the implementation seems easy ...

```diff
-def __init__(self, metrics: Union[List[Metric], Tuple[Metric], Dict[str, Metric]]):
+def __init__(self, metrics: Union[List[Metric], Tuple[Metric], Dict[str, Metric]], prefix: Union[Literal[None],str]=None):
         super().__init__()
         if isinstance(metrics, dict):
+            prefix = prefix if prefix is not None else ""
             # Check all values are metrics
             for name, metric in metrics.items():
                 if not isinstance(metric, Metric):
                     raise ValueError(
                         f"Value {metric} belonging to key {name}"
                         " is not an instance of `pl.metrics.Metric`"
                     )
-                self[name] = metric
+                self[prefix + name] = metric
```

possibly add a warning as well if a list (and not a dict) is passed so you know that it will get ignored.
let me know if you want the PR. (I could also update the docs)

Kindest wishes,
Sebastian
@jspaezp your solution looks very similar to the one I had in mind:
https://github.com/PyTorchLightning/pytorch-lightning/compare/master...SkafteNicki:fix_prefix_arg
Feel free to submit a PR but please wait a couple of days for an update that will come to the metrics package (cannot disclose exactly what right now) :]
Hello @SkafteNicki I like the idea of your implementation and I think it could add as well an optional argument to the `clone` method.

I am thinking something like ...

```python
...
def clone(self, prefix = None):
    cloned = super().clone()
    if prefix is not None:
        cloned.prefix = prefix
    return cloned
...
``` 

Since it would make it a lot easier to make clones of large collections (train and test metrics for instance ...)
let me know when you feel would be a good time to work on it (although i feel like you got it covered, thanks a lot)
I'm just updating to the latest version of pytorch lightning and I ran into this issue also. Do we have an ETA for when it will be released?
Hi! thanks for your contribution!, great first issue!

Resolving patch:
diff --git a/torchmetrics/collections.py b/torchmetrics/collections.py
--- a/torchmetrics/collections.py
+++ b/torchmetrics/collections.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from copy import deepcopy
-from typing import Any, Dict, List, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 from torch import nn
 
@@ -36,6 +36,8 @@ class MetricCollection(nn.ModuleDict):
               dict as key for output dict. Use this format if you want to chain
               together multiple of the same metric with different parameters.
 
+        prefix: a string to append in front of the keys of the output dict
+
     Example (input as list):
         >>> import torch
         >>> from torchmetrics import MetricCollection, Accuracy, Precision, Recall
@@ -58,8 +60,11 @@ class MetricCollection(nn.ModuleDict):
         >>> metrics.persistent()
 
     """
-
-    def __init__(self, metrics: Union[List[Metric], Tuple[Metric], Dict[str, Metric]]):
+    def __init__(
+            self,
+            metrics: Union[List[Metric], Tuple[Metric], Dict[str, Metric]],
+            prefix: Optional[str] = None
+    ):
         super().__init__()
         if isinstance(metrics, dict):
             # Check all values are metrics
@@ -84,13 +89,15 @@ def __init__(self, metrics: Union[List[Metric], Tuple[Metric], Dict[str, Metric]
         else:
             raise ValueError("Unknown input to MetricCollection.")
 
+        self.prefix = self._check_prefix_arg(prefix)
+
     def forward(self, *args, **kwargs) -> Dict[str, Any]:  # pylint: disable=E0202
         """
         Iteratively call forward for each metric. Positional arguments (args) will
         be passed to every metric in the collection, while keyword arguments (kwargs)
         will be filtered based on the signature of the individual metric.
         """
-        return {k: m(*args, **m._filter_kwargs(**kwargs)) for k, m in self.items()}
+        return {self._set_prefix(k): m(*args, **m._filter_kwargs(**kwargs)) for k, m in self.items()}
 
     def update(self, *args, **kwargs):  # pylint: disable=E0202
         """
@@ -103,20 +110,36 @@ def update(self, *args, **kwargs):  # pylint: disable=E0202
             m.update(*args, **m_kwargs)
 
     def compute(self) -> Dict[str, Any]:
-        return {k: m.compute() for k, m in self.items()}
+        return {self._set_prefix(k): m.compute() for k, m in self.items()}
 
-    def reset(self):
+    def reset(self) -> None:
         """ Iteratively call reset for each metric """
         for _, m in self.items():
             m.reset()
 
-    def clone(self):
-        """ Make a copy of the metric collection """
-        return deepcopy(self)
+    def clone(self, prefix: Optional[str] = None) -> 'MetricCollection':
+        """ Make a copy of the metric collection
+        Args:
+            prefix: a string to append in front of the metric keys
+        """
+        mc = deepcopy(self)
+        mc.prefix = self._check_prefix_arg(prefix)
+        return mc
 
-    def persistent(self, mode: bool = True):
+    def persistent(self, mode: bool = True) -> None:
         """Method for post-init to change if metric states should be saved to
         its state_dict
         """
         for _, m in self.items():
             m.persistent(mode)
+
+    def _set_prefix(self, k: str) -> str:
+        return k if self.prefix is None else self.prefix + k
+
+    def _check_prefix_arg(self, prefix: str) -> Optional[str]:
+        if prefix is not None:
+            if isinstance(prefix, str):
+                return prefix
+            else:
+                raise ValueError('Expected input `prefix` to be a string')
+        return None


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
@@ -130,3 +130,28 @@ def test_metric_collection_args_kwargs(tmpdir):
     _ = metric_collection(x=10, y=20)
     assert metric_collection['DummyMetricSum'].x == 10
     assert metric_collection['DummyMetricDiff'].x == -20
+
+
+def test_metric_collection_prefix_arg(tmpdir):
+    """ Test that the prefix arg alters the keywords in the output"""
+    m1 = DummyMetricSum()
+    m2 = DummyMetricDiff()
+    names = ['DummyMetricSum', 'DummyMetricDiff']
+
+    metric_collection = MetricCollection([m1, m2], prefix='prefix_')
+
+    # test forward
+    out = metric_collection(5)
+    for name in names:
+        assert f"prefix_{name}" in out, 'prefix argument not working as intended with forward method'
+
+    # test compute
+    out = metric_collection.compute()
+    for name in names:
+        assert f"prefix_{name}" in out, 'prefix argument not working as intended with compute method'
+
+    # test clone
+    new_metric_collection = metric_collection.clone(prefix='new_prefix_')
+    out = new_metric_collection(5)
+    for name in names:
+        assert f"new_prefix_{name}" in out, 'prefix argument not working as intended with clone method'