Add MetricInputTransformer

CHANGELOG.md

@@ -12,7 +12,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
--
+- Added `MetricInputTransformer` wrapper ([#2392](https://github.com/Lightning-AI/torchmetrics/pull/2392))
 
 
 ### Changed

docs/source/wrappers/transformations.rst

@@ -0,0 +1,23 @@
+.. customcarditem::
+   :header: Transformations
+   :image: https://pl-flash-data.s3.amazonaws.com/assets/thumbnails/graph_classification.svg
+   :tags: Wrappers
+
+.. include:: ../links.rst
+
+###############
+Transformations
+###############
+
+Transformations allow for modifications to the input a metric receives by wrapping its `pred` and `target` arguments.
+Transformations can be implemented by either subclassing the ``MetricInputTransformer`` base class and overriding the ``.transform_pred()`` and/or ``transform_target()`` functions, or by supplying a lambda function via the ``LambdaInputTransformer``.
+A ``BinaryTargetTransformer`` which casts target labels to 0/1 given a threshold is provided for convenience.
+
+Module Interface
+________________
+
+.. autoclass:: torchmetrics.wrappers.MetricInputTransformer
+
+.. autoclass:: torchmetrics.wrappers.LambdaInputTransformer
+
+.. autoclass:: torchmetrics.wrappers.BinaryTargetTransformer

src/torchmetrics/wrappers/__init__.py

@@ -19,11 +19,19 @@
 from torchmetrics.wrappers.multitask import MultitaskWrapper
 from torchmetrics.wrappers.running import Running
 from torchmetrics.wrappers.tracker import MetricTracker
+from torchmetrics.wrappers.transformations import (
+    BinaryTargetTransformer,
+    LambdaInputTransformer,
+    MetricInputTransformer,
+)
 
 __all__ = [
+    "BinaryTargetTransformer",
     "BootStrapper",
     "ClasswiseWrapper",
     "FeatureShare",
+    "LambdaInputTransformer",
+    "MetricInputTransformer",
     "MinMaxMetric",
     "MultioutputWrapper",
     "MultitaskWrapper",

src/torchmetrics/wrappers/transformations.py

@@ -0,0 +1,175 @@
+# Copyright The Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Any, Callable, Dict, Optional, Tuple, Union
+
+import torch
+
+from torchmetrics.collections import MetricCollection
+from torchmetrics.metric import Metric
+from torchmetrics.wrappers.abstract import WrapperMetric
+
+
+class MetricInputTransformer(WrapperMetric):
+    """Abstract base class for metric input transformations.
+
+    Input transformations are characterized by them applying a transformation to the input data of a metric, and then
+    forwarding all calls to the wrapped metric with modifications applied.
+
+    """
+
+    def __init__(self, wrapped_metric: Union[Metric, MetricCollection], **kwargs: Dict[str, Any]) -> None:
+        super().__init__(**kwargs)
+        if not isinstance(wrapped_metric, (Metric, MetricCollection)):
+            raise TypeError(
+                f"Expected wrapped metric to be an instance of `torchmetrics.Metric` or "
+                f"`torchmetrics.MetricsCollection`but received {wrapped_metric}"
+            )
+        self.wrapped_metric = wrapped_metric
+
+    def transform_pred(self, pred: torch.Tensor) -> torch.Tensor:
+        """Define transform operations on the prediction data.
+
+        Overridden by subclasses. Identity by default.
+
+        """
+        return pred
+
+    def transform_target(self, target: torch.Tensor) -> torch.Tensor:
+        """Define transform operations on the target data.
+
+        Overridden by subclasses. Identity by default.
+
+        """
+        return target
+
+    def _wrap_transform(self, *args: torch.Tensor) -> Tuple[torch.Tensor, ...]:
+        """Wrap transformation functions to dispatch args to their individual transform functions."""
+        if len(args) == 1:
+            return (self.transform_pred(args[0]),)
+        if len(args) == 2:
+            return self.transform_pred(args[0]), self.transform_target(args[1])
+        return self.transform_pred(args[0]), self.transform_target(args[1]), *args[2:]
+
+    def update(self, *args: torch.Tensor, **kwargs: Dict[str, Any]) -> None:
+        """Wrap the update call of the underlying metric."""
+        args = self._wrap_transform(*args)
+        self.wrapped_metric.update(*args, **kwargs)
+
+    def compute(self) -> Any:
+        """Wrap the compute call of the underlying metric."""
+        return self.wrapped_metric.compute()
+
+    def forward(self, *args: torch.Tensor, **kwargs: Dict[str, Any]) -> Any:
+        """Wrap the forward call of the underlying metric."""
+        args = self._wrap_transform(*args)
+        return self.wrapped_metric.forward(*args, **kwargs)
+
+
+class LambdaInputTransformer(MetricInputTransformer):
+    """Wrapper class for transforming a metrics' inputs given a user-defined lambda function.
+
+    Args:
+        wrapped_metric:
+            The underlying `Metric` or `MetricCollection`.
+        transform_pred:
+            The function to apply to the predictions before computing the metric.
+        transform_target:
+            The function to apply to the target before computing the metric.
+
+    Raises:
+        TypeError:
+            If `transform_pred` is not a Callable.
+        TypeError:
+            If `transform_target` is not a Callable.
+
+    Example:
+        >>> import torch
+        >>> from torchmetrics.classification import BinaryAccuracy
+        >>> from torchmetrics.wrappers import LambdaInputTransformer
+        >>>
+        >>> preds = torch.tensor([0.9, 0.8, 0.7, 0.6, 0.5, 0.6, 0.7, 0.8, 0.5, 0.4])
+        >>> targets = torch.tensor([1,0,0,0,0,1,1,0,0,0])
+        >>>
+        >>> metric = LambdaInputTransformer(BinaryAccuracy(), lambda preds: 1 - preds)
+        >>> metric.update(preds, targets)
+        >>> metric.compute()
+        tensor(0.6000)
+
+    """
+
+    def __init__(
+        self,
+        wrapped_metric: Metric,
+        transform_pred: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
+        transform_target: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(wrapped_metric, **kwargs)
+        if transform_pred is not None:
+            if not callable(transform_pred):
+                raise TypeError(f"Expected `transform_pred` to be of type `Callable` but received `{transform_pred}`")
+            self.transform_pred = transform_pred  # type: ignore[assignment,method-assign]
+
+        if transform_target is not None:
+            if not callable(transform_target):
+                raise TypeError(
+                    f"Expected `transform_target` to be of type `Callable` but received `{transform_target}`"
+                )
+            self.transform_target = transform_target  # type: ignore[assignment,method-assign]
+
+
+class BinaryTargetTransformer(MetricInputTransformer):
+    """Wrapper class for computing a metric on binarized targets.
+
+    Useful when the given ground-truth targets are continuous, but the metric requires binary targets.
+
+    Args:
+        wrapped_metric:
+            The underlying `Metric` or `MetricCollection`.
+        threshold:
+            The binarization threshold for the targets. Targets values `t` are cast to binary with `t > threshold`.
+
+    Raises:
+        TypeError:
+            If `threshold` is not an `int` or `float`.
+
+    Example:
+        >>> import torch
+        >>> from torchmetrics.retrieval import RetrievalMRR
+        >>> from torchmetrics.wrappers import BinaryTargetTransformer
+        >>>
+        >>> preds = torch.tensor([0.9, 0.8, 0.7, 0.6, 0.5, 0.6, 0.7, 0.8, 0.5, 0.4])
+        >>> targets = torch.tensor([1,0,0,0,0,2,1,0,0,0])
+        >>> topics = torch.tensor([0,0,0,0,0,1,1,1,1,1])
+        >>>
+        >>> metric = BinaryTargetTransformer(RetrievalMRR())
+        >>> metric.update(preds, targets, indexes=topics)
+        >>> metric.compute()
+        tensor(0.7500)
+
+    """
+
+    def __init__(self, wrapped_metric: Union[Metric, MetricCollection], threshold: float = 0, **kwargs: Any) -> None:
+        super().__init__(wrapped_metric, **kwargs)
+        if not isinstance(threshold, (int, float)):
+            raise TypeError(f"Expected `threshold` to be of type `int` or `float` but received `{threshold}`")
+        self.threshold = threshold
+
+    def transform_target(self, target: torch.Tensor) -> torch.Tensor:
+        """Cast the target tensor to binary values according to the threshold.
+
+        Output assumes same type as input.
+
+        """
+        return target.gt(self.threshold).to(target.dtype)