Bugfix for using metric collection and aggregation metric (#1896)

CHANGELOG.md

@@ -29,6 +29,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- Fixes corner case when using `MetricCollection` together with aggregation metrics ([#1896](https://github.com/Lightning-AI/torchmetrics/pull/1896))
+
+
 - Fixed the use of `max_fpr` in `AUROC` metric when only one class is present ([#1895](https://github.com/Lightning-AI/torchmetrics/pull/1895))
 
 

src/torchmetrics/aggregation.py

@@ -39,14 +39,14 @@ class BaseAggregator(Metric):
  - ``'ignore'``: all `nan` values are silently removed
  - a float: if a float is provided will impude any `nan` values with this value
 
+ state_name: name of the metric state
  kwargs: Additional keyword arguments, see :ref:`Metric kwargs` for more info.
 
  Raises:
  ValueError:
  If ``nan_strategy`` is not one of ``error``, ``warn``, ``ignore`` or a float
  """
 
- value: Tensor
  is_differentiable = None
  higher_is_better = None
  full_state_update: bool = False
@@ -56,6 +56,7 @@ def __init__(
  fn: Union[Callable, str],
  default_value: Union[Tensor, List],
  nan_strategy: Union[str, float] = "error",
+ state_name: str = "value",
  **kwargs: Any,
  ) -> None:
  super().__init__(**kwargs)
@@ -67,7 +68,8 @@ def __init__(
  )
 
  self.nan_strategy = nan_strategy
- self.add_state("value", default=default_value, dist_reduce_fx=fn)
+ self.add_state(state_name, default=default_value, dist_reduce_fx=fn)
+ self.state_name = state_name
 
  def _cast_and_nan_check_input(
  self, x: Union[float, Tensor], weight: Optional[Union[float, Tensor]] = None
@@ -105,7 +107,7 @@ def update(self, value: Union[float, Tensor]) -> None:
 
  def compute(self) -> Tensor:
  """Compute the aggregated value."""
- return self.value
+ return getattr(self, self.state_name)
 
 
 class MaxMetric(BaseAggregator):
@@ -144,6 +146,7 @@ class MaxMetric(BaseAggregator):
  """
 
  full_state_update: bool = True
+ max_value: Tensor
 
  def __init__(
  self,
@@ -154,6 +157,7 @@ def __init__(
  "max",
  -torch.tensor(float("inf")),
  nan_strategy,
+ state_name="max_value",
  **kwargs,
  )
 
@@ -166,7 +170,7 @@ def update(self, value: Union[float, Tensor]) -> None:
  """
  value, _ = self._cast_and_nan_check_input(value)
  if value.numel(): # make sure tensor not empty
- self.value = torch.max(self.value, torch.max(value))
+ self.max_value = torch.max(self.max_value, torch.max(value))
 
  def plot(
  self, val: Optional[Union[Tensor, Sequence[Tensor]]] = None, ax: Optional[_AX_TYPE] = None
@@ -244,6 +248,7 @@ class MinMetric(BaseAggregator):
  """
 
  full_state_update: bool = True
+ min_value: Tensor
 
  def __init__(
  self,
@@ -254,6 +259,7 @@ def __init__(
  "min",
  torch.tensor(float("inf")),
  nan_strategy,
+ state_name="min_value",
  **kwargs,
  )
 
@@ -266,7 +272,7 @@ def update(self, value: Union[float, Tensor]) -> None:
  """
  value, _ = self._cast_and_nan_check_input(value)
  if value.numel(): # make sure tensor not empty
- self.value = torch.min(self.value, torch.min(value))
+ self.min_value = torch.min(self.min_value, torch.min(value))
 
  def plot(
  self, val: Optional[Union[Tensor, Sequence[Tensor]]] = None, ax: Optional[_AX_TYPE] = None
@@ -343,6 +349,8 @@ class SumMetric(BaseAggregator):
  tensor(6.)
  """
 
+ sum_value: Tensor
+
  def __init__(
  self,
  nan_strategy: Union[str, float] = "warn",
@@ -352,6 +360,7 @@ def __init__(
  "sum",
  torch.tensor(0.0),
  nan_strategy,
+ state_name="sum_value",
  **kwargs,
  )
 
@@ -364,7 +373,7 @@ def update(self, value: Union[float, Tensor]) -> None:
  """
  value, _ = self._cast_and_nan_check_input(value)
  if value.numel():
- self.value += value.sum()
+ self.sum_value += value.sum()
 
  def plot(
  self, val: Optional[Union[Tensor, Sequence[Tensor]]] = None, ax: Optional[_AX_TYPE] = None
@@ -442,6 +451,8 @@ class CatMetric(BaseAggregator):
  tensor([1., 2., 3.])
  """
 
+ value: Tensor
+
  def __init__(
  self,
  nan_strategy: Union[str, float] = "warn",
@@ -503,6 +514,8 @@ class MeanMetric(BaseAggregator):
  tensor(2.)
  """
 
+ mean_value: Tensor
+
  def __init__(
  self,
  nan_strategy: Union[str, float] = "warn",
@@ -512,6 +525,7 @@ def __init__(
  "sum",
  torch.tensor(0.0),
  nan_strategy,
+ state_name="mean_value",
  **kwargs,
  )
  self.add_state("weight", default=torch.tensor(0.0), dist_reduce_fx="sum")
@@ -537,12 +551,12 @@ def update(self, value: Union[float, Tensor], weight: Union[float, Tensor] = 1.0
 
  if value.numel() == 0:
  return
- self.value += (value * weight).sum()
+ self.mean_value += (value * weight).sum()
  self.weight += weight.sum()
 
  def compute(self) -> Tensor:
  """Compute the aggregated value."""
- return self.value / self.weight
+ return self.mean_value / self.weight
 
  def plot(
  self, val: Optional[Union[Tensor, Sequence[Tensor]]] = None, ax: Optional[_AX_TYPE] = None

tests/integrations/test_lightning.py

@@ -419,22 +419,22 @@ def configure_optimizers(self):
  return torch.optim.SGD(self.layer.parameters(), lr=0.1)
 
  model = BoringModel()
- assert model.metric.value.dtype == torch.float32
+ assert model.metric.sum_value.dtype == torch.float32
  model = model.half()
- assert model.metric.value.dtype == torch.float32
+ assert model.metric.sum_value.dtype == torch.float32
 
  model = BoringModel()
- assert model.metric.value.dtype == torch.float32
+ assert model.metric.sum_value.dtype == torch.float32
  model = model.double()
- assert model.metric.value.dtype == torch.float32
+ assert model.metric.sum_value.dtype == torch.float32
 
  model = BoringModel(metric_dtype=torch.float16)
- assert model.metric.value.dtype == torch.float16
+ assert model.metric.sum_value.dtype == torch.float16
  model = model.float()
- assert model.metric.value.dtype == torch.float16
+ assert model.metric.sum_value.dtype == torch.float16
 
  model = BoringModel()
- assert model.metric.value.dtype == torch.float32
+ assert model.metric.sum_value.dtype == torch.float32
 
  model = model.type(torch.half)
- assert model.metric.value.dtype == torch.float32
+ assert model.metric.sum_value.dtype == torch.float32

tests/unittests/bases/test_aggregation.py

@@ -2,6 +2,7 @@
 import pytest
 import torch
 from torchmetrics.aggregation import CatMetric, MaxMetric, MeanMetric, MinMetric, SumMetric
+from torchmetrics.collections import MetricCollection
 
 from unittests import BATCH_SIZE, NUM_BATCHES
 from unittests.helpers.testers import MetricTester
@@ -166,6 +167,22 @@ def test_mean_metric_broadcasting(weights, expected):
  assert avg(values, weights) == expected
 
 
+def test_aggregation_in_collection_with_compute_groups():
+ """Check that aggregation metrics work in MetricCollection with compute_groups=True."""
+ m = MetricCollection(MinMetric(), MaxMetric(), SumMetric(), MeanMetric(), compute_groups=True)
+ assert len(m.compute_groups) == 4, "Expected 4 compute groups"
+ m.update(1)
+ assert len(m.compute_groups) == 4, "Expected 4 compute groups"
+ m.update(2)
+ assert len(m.compute_groups) == 4, "Expected 4 compute groups"
+
+ res = m.compute()
+ assert res["MinMetric"] == 1
+ assert res["MaxMetric"] == 2
+ assert res["SumMetric"] == 3
+ assert res["MeanMetric"] == 1.5
+
+
 @pytest.mark.skipif(not hasattr(torch, "broadcast_to"), reason="PyTorch <1.8 does not have broadcast_to")
 @pytest.mark.parametrize("nan_strategy", ["ignore", "warn"])
 def test_mean_metric_broadcast(nan_strategy):