Refactor/remove double forward

CHANGELOG.md

@@ -22,7 +22,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Reimplemented the `signal_distortion_ratio` metric, which removed the absolute requirement of `fast-bss-eval` ([#964](https://github.com/PyTorchLightning/metrics/pull/964))
 
 
--
+- Changed `forward` method to only calling `update` a single time instead of twice ([#984](https://github.com/PyTorchLightning/metrics/pull/984))
 
 
 ### Deprecated

docs/source/pages/implement.rst

@@ -64,18 +64,13 @@ The cache is first emptied on the next call to ``update``.
 
 ``forward`` serves the dual purpose of both returning the metric on the current data and updating the internal
 metric state for accumulating over multiple batches. The ``forward()`` method achieves this by combining calls
-to ``update`` and ``compute`` in the following way:
-
-1. Calls ``update()`` to update the global metric state (for accumulation over multiple batches)
-2. Caches the global state.
-3. Calls ``reset()`` to clear global metric state.
-4. Calls ``update()`` to update local metric state.
-5. Calls ``compute()`` to calculate metric for current batch.
-6. Restores the global state.
-
-This procedure has the consequence of calling the user defined ``update`` **twice** during a single
-forward call (one to update global statistics and one for getting the batch statistics).
+to ``update``, ``compute`` and ``reset`` in the following way:
 
+1. Caches the global state.
+2. Calls ``reset`` the metric to its default state
+3. Calls ``update`` to update the state with local batch statistics
+4. Calls ``compute`` to calculate the metric for the current batch
+5. Reduce the global state and batch state into a single state that becomes the new global state
 
 ---------
 

tests/helpers/testers.py

@@ -569,7 +569,7 @@ class DummyMetric(Metric):
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
-        self.add_state("x", tensor(0.0), dist_reduce_fx=None)
+        self.add_state("x", tensor(0.0), dist_reduce_fx="sum")
 
     def update(self):
         pass
@@ -583,7 +583,7 @@ class DummyListMetric(Metric):
 
     def __init__(self):
         super().__init__()
-        self.add_state("x", [], dist_reduce_fx=None)
+        self.add_state("x", [], dist_reduce_fx="cat")
 
     def update(self):
         pass

torchmetrics/classification/roc.py

@@ -114,8 +114,8 @@ def __init__(
         self.num_classes = num_classes
         self.pos_label = pos_label
 
-        self.add_state("preds", default=[], dist_reduce_fx=None)
-        self.add_state("target", default=[], dist_reduce_fx=None)
+        self.add_state("preds", default=[], dist_reduce_fx="cat")
+        self.add_state("target", default=[], dist_reduce_fx="cat")
 
         rank_zero_warn(
             "Metric `ROC` will save all targets and predictions in buffer."

torchmetrics/classification/stat_scores.py

@@ -151,7 +151,7 @@ def __init__(
             raise ValueError(f"The `ignore_index` {ignore_index} is not valid for inputs with {num_classes} classes")
 
         default: Callable = lambda: []
-        reduce_fn: Optional[str] = None
+        reduce_fn: Optional[str] = "cat"
         if mdmc_reduce != "samplewise" and reduce != "samples":
             if reduce == "micro":
                 zeros_shape = []

torchmetrics/image/psnr.py

@@ -83,8 +83,8 @@ def __init__(
             self.add_state("sum_squared_error", default=tensor(0.0), dist_reduce_fx="sum")
             self.add_state("total", default=tensor(0), dist_reduce_fx="sum")
         else:
-            self.add_state("sum_squared_error", default=[])
-            self.add_state("total", default=[])
+            self.add_state("sum_squared_error", default=[], dist_reduce_fx="cat")
+            self.add_state("total", default=[], dist_reduce_fx="cat")
 
         if data_range is None:
             if dim is not None:

torchmetrics/metric.py

@@ -197,41 +197,38 @@ def add_state(
 
     @torch.jit.unused
     def forward(self, *args: Any, **kwargs: Any) -> Any:
-        """Automatically calls ``update()``.
+        """``forward`` serves the dual purpose of both computing the metric on the current batch of inputs but also
+        add the batch statistics to the overall accumululating metric state.
 
-        Returns the metric value over inputs if ``compute_on_step`` is True.
+        Input arguments are the exact same as corresponding ``update`` method. The returned output is the exact same as
+        the output of ``compute``.
         """
-        # add current step
+        # check if states are already synced
         if self._is_synced:
             raise TorchMetricsUserError(
-                "The Metric shouldn't be synced when performing ``update``. "
+                "The Metric shouldn't be synced when performing ``forward``. "
                 "HINT: Did you forget to call ``unsync`` ?."
             )
+        # store global state and reset to default
+        global_state = {attr: getattr(self, attr) for attr in self._defaults.keys()}
+        self.reset()
 
-        # global accumulation
-        self.update(*args, **kwargs)
-
-        self._to_sync = self.dist_sync_on_step  # type: ignore
-        # skip restore cache operation from compute as cache is stored below.
+        # local syncronization settings
+        self._to_sync = self.dist_sync_on_step
         self._should_unsync = False
-        # skip computing on cpu for the batch
         _temp_compute_on_cpu = self.compute_on_cpu
         self.compute_on_cpu = False
-
-        # save context before switch
-        cache = {attr: getattr(self, attr) for attr in self._defaults}
-
-        # call reset, update, compute, on single batch
         self._enable_grad = True  # allow grads for batch computation
-        self.reset()
+
+        # calculate batch state and compute batch value
         self.update(*args, **kwargs)
         self._forward_cache = self.compute()
 
+        # reduce batch and global state
+        self._reduce_state(global_state)
+
         # restore context
-        for attr, val in cache.items():
-            setattr(self, attr, val)
         self._is_synced = False
-
         self._should_unsync = True
         self._to_sync = True
         self._computed = None
@@ -240,6 +237,33 @@ def forward(self, *args: Any, **kwargs: Any) -> Any:
 
         return self._forward_cache
 
+    def _reduce_state(self, state_to_reduce: Dict[str, Any]) -> None:
+        for attr in self._defaults.keys():
+            current_state = getattr(self, attr)
+            incoming_state = state_to_reduce[attr]
+            reduce_fn = self._reductions[attr]
+            if reduce_fn == dim_zero_sum:
+                reduced = current_state + incoming_state
+            elif reduce_fn == dim_zero_mean:
+                reduced = (current_state + incoming_state) / 2.0
+            elif reduce_fn == dim_zero_max:
+                reduced = torch.max(current_state, incoming_state)
+            elif reduce_fn == dim_zero_min:
+                reduced = torch.min(current_state, incoming_state)
+            elif reduce_fn == dim_zero_cat:  # or (reduce_fn is None and isinstance(current_state, list)):
+                reduced = incoming_state + current_state
+            elif reduce_fn is None and isinstance(current_state, Tensor):
+                if incoming_state.ndim > 0:  # TODO: figure out why this works
+                    reduced = torch.cat([current_state.expand(1), incoming_state], dim=0)
+                else:
+                    reduced = torch.stack([current_state, incoming_state], dim=0)
+            elif reduce_fn is None and isinstance(current_state, list):
+                reduced = _flatten([current_state, incoming_state])
+            else:
+                reduced = reduce_fn(torch.stack([current_state, incoming_state]))  # type: ignore
+
+            setattr(self, attr, reduced)
+
     def _sync_dist(self, dist_sync_fn: Callable = gather_all_tensors, process_group: Optional[Any] = None) -> None:
         input_dict = {attr: getattr(self, attr) for attr in self._reductions}