Fix ComputationalMetric.forward() to not only consider the last batch

CHANGELOG.md

@@ -47,6 +47,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Fix edge case of AUROC with `average=weighted` on GPU ([#606](https://github.com/PyTorchLightning/metrics/pull/606))
 
 
+- Fixed `forward` in compositional metrics ([#645](https://github.com/PyTorchLightning/metrics/pull/645))
+
 
 ## [0.6.0] - 2021-10-28
 

tests/bases/test_composition.py

@@ -24,7 +24,7 @@
 class DummyMetric(Metric):
  def __init__(self, val_to_return):
  super().__init__()
- self._num_updates = 0
+ self.add_state("_num_updates", tensor(0), dist_reduce_fx="sum")
  self._val_to_return = val_to_return
  self._update_called = True
 
@@ -34,10 +34,6 @@ def update(self, *args, **kwargs) -> None:
  def compute(self):
  return tensor(self._val_to_return)
 
- def reset(self):
- self._num_updates = 0
- return super().reset()
-
 
 @pytest.mark.parametrize(
  ["second_operand", "expected_result"],
@@ -544,7 +540,7 @@ def test_metrics_getitem(value, idx, expected_result):
 
 
 def test_compositional_metrics_update():
-
+ """test update method for compositional metrics."""
  compos = DummyMetric(5) + DummyMetric(4)
 
  assert isinstance(compos, CompositionalMetric)
@@ -557,3 +553,26 @@ def test_compositional_metrics_update():
 
  assert compos.metric_a._num_updates == 3
  assert compos.metric_b._num_updates == 3
+
+
+@pytest.mark.parametrize("compute_on_step", [True, False])
+@pytest.mark.parametrize("metric_b", [4, DummyMetric(4)])
+def test_compositional_metrics_forward(compute_on_step, metric_b):
+ """test forward method of compositional metrics."""
+ metric_a = DummyMetric(5)
+ metric_a.compute_on_step = compute_on_step
+ compos = metric_a + metric_b
+
+ assert isinstance(compos, CompositionalMetric)
+ for _ in range(3):
+ val = compos()
+ assert val == 9 if compute_on_step else val is None
+
+ assert isinstance(compos.metric_a, DummyMetric)
+ assert compos.metric_a._num_updates == 3
+
+ if isinstance(metric_b, DummyMetric):
+ assert isinstance(compos.metric_b, DummyMetric)
+ assert compos.metric_b._num_updates == 3
+
+ compos.reset()

torchmetrics/metric.py

@@ -756,6 +756,35 @@ def compute(self) -> Any:
 
  return self.op(val_a, val_b)
 
+ @torch.jit.unused
+ def forward(self, *args: Any, **kwargs: Any) -> Any:
+
+ val_a = (
+ self.metric_a(*args, **self.metric_a._filter_kwargs(**kwargs))
+ if isinstance(self.metric_a, Metric)
+ else self.metric_a
+ )
+ val_b = (
+ self.metric_b(*args, **self.metric_b._filter_kwargs(**kwargs))
+ if isinstance(self.metric_b, Metric)
+ else self.metric_b
+ )
+
+ if val_a is None:
+ # compute_on_step of metric_a is False
+ return None
+
+ if val_b is None:
+ if isinstance(self.metric_b, Metric):
+ # compute_on_step of metric_b is False
+ return None
+
+ # Unary op
+ return self.op(val_a)
+
+ # Binary op
+ return self.op(val_a, val_b)
+
  def reset(self) -> None:
  if isinstance(self.metric_a, Metric):
  self.metric_a.reset()
@@ -774,3 +803,6 @@ def __repr__(self) -> str:
  repr_str = self.__class__.__name__ + _op_metrics
 
  return repr_str
+
+ def _wrap_compute(self, compute: Callable) -> Callable:
+ return compute