[bugfix] MetricCollection should return metrics with prefix on items(), keys()

tests/bases/test_collections.py

@@ -168,6 +168,21 @@ def test_metric_collection_prefix_postfix_args(prefix, postfix):
     for name in names:
         assert f"new_prefix_{name}" in out, 'prefix argument not working as intended with clone method'
 
+    for k, _ in new_metric_collection.items():
+        assert 'new_prefix_' in k
+
+    for k in new_metric_collection.keys():
+        assert 'new_prefix_' in k
+
+    for k, _ in new_metric_collection.items(keep_base=True):
+        assert 'new_prefix_' not in k
+
+    for k in new_metric_collection.keys(keep_base=True):
+        assert 'new_prefix_' not in k
+
+    repr = f'MetricCollection(\n  prefix=new_prefix_,\n  postfix={postfix},\n  (DummyMetricSum): DummyMetricSum()\n  (DummyMetricDiff): DummyMetricDiff()\n)'
+    assert new_metric_collection.__repr__() == repr
+
     new_metric_collection = new_metric_collection.clone(postfix='_new_postfix')
     out = new_metric_collection(5)
     names = [n[:-len(postfix)] if postfix is not None else n for n in names]  # strip away old postfix

torchmetrics/collections.py

@@ -13,9 +13,10 @@
 # limitations under the License.
 
 from copy import deepcopy
-from typing import Any, Dict, Optional, Sequence, Union
+from typing import Any, Dict, Iterable, Optional, Sequence, Tuple, Union
 
 from torch import nn
+from torch._jit_internal import _copy_to_script_wrapper
 
 from torchmetrics.metric import Metric
 from torchmetrics.utilities import rank_zero_warn
@@ -144,24 +145,24 @@ def forward(self, *args, **kwargs) -> Dict[str, Any]:  # pylint: disable=E0202
         be passed to every metric in the collection, while keyword arguments (kwargs)
         will be filtered based on the signature of the individual metric.
         """
-        return {self._set_name(k): m(*args, **m._filter_kwargs(**kwargs)) for k, m in self.items()}
+        return {k: m(*args, **m._filter_kwargs(**kwargs)) for k, m in self.items()}
 
     def update(self, *args, **kwargs):  # pylint: disable=E0202
         """
         Iteratively call update 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.
         """
-        for _, m in self.items():
+        for _, m in self.items(keep_base=True):
             m_kwargs = m._filter_kwargs(**kwargs)
             m.update(*args, **m_kwargs)
 
     def compute(self) -> Dict[str, Any]:
-        return {self._set_name(k): m.compute() for k, m in self.items()}
+        return {k: m.compute() for k, m in self.items()}
 
     def reset(self) -> None:
         """ Iteratively call reset for each metric """
-        for _, m in self.items():
+        for _, m in self.items(keep_base=True):
             m.reset()
 
     def clone(self, prefix: Optional[str] = None, postfix: Optional[str] = None) -> 'MetricCollection':
@@ -182,16 +183,40 @@ 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():
+        for _, m in self.items(keep_base=True):
             m.persistent(mode)
 
     def _set_name(self, base: str) -> str:
         name = base if self.prefix is None else self.prefix + base
         name = name if self.postfix is None else name + self.postfix
         return name
 
+    @_copy_to_script_wrapper
+    def keys(self, keep_base: bool = False):
+        r"""Return an iterable of the ModuleDict key.
+        Args:
+            keep_base: Whether to add prefix/postfix on the items collection.
+        """
+        if keep_base:
+            return self._modules.keys()
+        return [self._set_name(k) for k in self._modules.keys()]
+
+    @_copy_to_script_wrapper
+    def items(self, keep_base: bool = False) -> Iterable[Tuple[str, nn.Module]]:
+        r"""Return an iterable of the ModuleDict key/value pairs.
+        Args:
+            keep_base: Whether to add prefix/postfix on the items collection.
+        """
+        if keep_base:
+            return self._modules.items()
+        return [(self._set_name(k), v) for k, v in self._modules.items()]
+
     @staticmethod
     def _check_arg(arg: Optional[str], name: str) -> Optional[str]:
         if arg is None or isinstance(arg, str):
             return arg
         raise ValueError(f'Expected input `{name}` to be a string, but got {type(arg)}')
+
+    def extra_repr(self) -> Optional[str]:
+        if self.prefix or self.postfix:
+            return f"prefix={self.prefix},\npostfix={self.postfix},"