Repository: Lightning-AI/torchmetrics

Issue:
MetricWrapper for Target Binarization
## 🚀 Feature

Add a `TargetBinarizationWrapper` that cast continuous labels to binary labels given a threshold.

### Motivation

Evaluating two metrics that require different label formats (e.g., one binary, the other continuous) is cumbersome since it requires setting up two different evaluation stacks where one is fed with binarized label data and the other is fed the original continuous data. This leads to code duplication. Also, persisting binarized labels into the dataset in scenarios where a metric requires different input than what is given in the ground-truth data diminishes code clarity w.r.t. the evaluation process.

### Pitch

A metric wrapper that casts target data to binary targets during the `.update()` and `.forward()` methods. Can be applied to either a single `Metric`, or a whole `MetricCollection`.

### Alternatives

- using a `MultiTaskWrapper` is possible, but has two caveats: (1) metrics with a different signature than `update(pred, target)` are not supported, and (2) it requires the user to implement the thresholding logic by themselves before feeding it into the `MultiTaskWrapper`
- use a more generic target processing wrapper that would allow supplying, e.g., a custom `lambda` that is applied to targets; more flexible, but also requires the user to implement their own logic. I think binarization is a common enough problem in `torchmetrics` (since its metrics make a binary vs. non-binary distinction) to warrant its own wrapper.

### Additional Information

Consider the following example of the desired behaviour:
```python

import torch
from torchmetrics.wrappers import BinarizedTargetWrapper # <-- This does not exist
from torchmetrics.collections import MetricCollection
from torchmetrics.retrieval import RetrievalNormalizedDCG, RetrievalMRR 

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])

metrics = MetricCollection({
    "RetrievalNormalizedDCG": RetrievalNormalizedDCG(),
    "RetrievalMRR": BinarizedTargetWrapper(RetrievalMRR(), threshold=0)), # <-- Enable this kind of metric composition which is not possible currently
})

metrics.update(preds, targets, indexes=topics)
metrics.compute()
```

If simple binarization as in the example is a desired solution, I have all the code needed for a pull request ready and can take on this issue.


Discussion:
Hi! thanks for your contribution!, great first issue!
Hi @lgienapp, thanks for opening this issue.
I am not against adding this feature, but maybe we should consider adding:
* an general `MetricInputTransformer` class where user can provide custom functions for transforming the input
* then `BinarizedTargetWrapper` can be included as just a subclass `MetricInputTransformer` with pre-selected transforms

how does that sound?
Establishing a general class sounds good. Just for clarification: the general `MetricInputTransformer`  would still subclass the `WrapperMetric` to inherit all the "reset-sync" code from there (e.g., live in `torchmetrics.wrappers.transformations`)? Or be its own thing (e.g., live in `torchmetrics.transformations.abstract`), possibly duplicating the sync code from the wrapper base class?


In either case, I would propose an implementation like this (subclassing wrappers here), assuming that only positional params like `preds` and `targets` would be interesting to modify (and thus ignoring kwargs e.g. `indices`):

```python
class MetricInputTransformer(WrapperMetric):
  
    def __init__(self, wrapped_metric: Union[Metric, MetricCollection], **kwargs: Any):
        super().__init__(**kwargs)
        self.wrapped_metric = wrapped_metric

    def transform(self, *args) -> Tuple[torch.Tensor]:
        raise NotImplementedError

    def update(self, *args, **kwargs: Any) -> None:
        self.wrapped_metric.update(*self.transform(*args), **kwargs)

    def compute(self) -> Any:
        return self.wrapped_metric.compute()

    def forward(self, *args, **kwargs: Any) -> Any:
        self.wrapped_metric.forward(*self.transform(*args), **kwargs)
```



Resolving patch:
diff --git a/src/torchmetrics/wrappers/__init__.py b/src/torchmetrics/wrappers/__init__.py
--- a/src/torchmetrics/wrappers/__init__.py
+++ b/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",
diff --git a/src/torchmetrics/wrappers/transformations.py b/src/torchmetrics/wrappers/transformations.py
new file mode 100644
--- /dev/null
+++ b/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)


Test patch:
diff --git a/tests/unittests/wrappers/test_transformations.py b/tests/unittests/wrappers/test_transformations.py
new file mode 100644
--- /dev/null
+++ b/tests/unittests/wrappers/test_transformations.py
@@ -0,0 +1,234 @@
+# 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 ClassVar
+
+import pytest
+from torch import Tensor
+from torchmetrics.aggregation import MeanMetric
+from torchmetrics.classification import BinaryAccuracy
+from torchmetrics.retrieval import RetrievalMAP
+from torchmetrics.wrappers import BinaryTargetTransformer, LambdaInputTransformer, MetricInputTransformer
+
+from unittests._helpers import seed_all
+
+seed_all(42)
+
+
+class TestMetricInputTransformer:
+    """Test suite for MetricInputTransformer."""
+
+    def test_no_base_metric(self) -> None:
+        """Tests that TypeError is raised when no wrapped_metric is passed."""
+        with pytest.raises(TypeError, match=r"Expected wrapped metric to be an instance of .*"):
+            MetricInputTransformer([])
+
+
+class TestLambdaInputTransformer:
+    """Test suite for LambdaInputTransformer."""
+
+    _test_signature: ClassVar = (
+        "cls",
+        "transform_pred",
+        "transform_target",
+        "preds",
+        "preds_transformed",
+        "targets",
+        "targets_transformed",
+    )
+    _test_cases: ClassVar = [
+        # No change to input data (identity transform)
+        (
+            BinaryAccuracy,
+            None,
+            None,
+            [0.9, 0.8, 0.7, 0.6, 0.5, 0.6, 0.7, 0.8, 0.5, 0.4],
+            [0.9, 0.8, 0.7, 0.6, 0.5, 0.6, 0.7, 0.8, 0.5, 0.4],
+            [1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0],
+            [1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0],
+        ),
+        # Change to pred data (invert transform)
+        (
+            BinaryAccuracy,
+            lambda pred: 1 - pred,
+            None,
+            [0.9, 0.8, 0.7, 0.6, 0.5, 0.6, 0.7, 0.8, 0.5, 0.4],
+            [0.1, 0.2, 0.3, 0.4, 0.5, 0.4, 0.3, 0.2, 0.5, 0.6],
+            [1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0],
+            [1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0],
+        ),
+        # Change to target data (invert transform)
+        (
+            BinaryAccuracy,
+            None,
+            lambda target: 1 - target,
+            [0.9, 0.8, 0.7, 0.6, 0.5, 0.6, 0.7, 0.8, 0.5, 0.4],
+            [0.9, 0.8, 0.7, 0.6, 0.5, 0.6, 0.7, 0.8, 0.5, 0.4],
+            [1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0],
+            [0.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 1.0, 1.0, 1.0],
+        ),
+        # Change to all input data (invert transform)
+        (
+            BinaryAccuracy,
+            lambda pred: 1 - pred,
+            lambda target: 1 - target,
+            [0.9, 0.8, 0.7, 0.6, 0.5, 0.6, 0.7, 0.8, 0.5, 0.4],
+            [0.1, 0.2, 0.3, 0.4, 0.5, 0.4, 0.3, 0.2, 0.5, 0.6],
+            [1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0],
+            [0.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 1.0, 1.0, 1.0],
+        ),
+    ]
+
+    @pytest.mark.parametrize(_test_signature, _test_cases)
+    def test_forward(
+        self, cls, transform_pred, transform_target, preds, preds_transformed, targets, targets_transformed
+    ) -> None:
+        """Tests if the binarized forward matches the output of the metric on manually binarized targets."""
+        metric = cls()
+        wrapped_metric = LambdaInputTransformer(
+            metric, transform_pred=transform_pred, transform_target=transform_target
+        )
+        preds = Tensor(preds).float()
+        targets = Tensor(targets).float()
+        preds_transformed = Tensor(preds_transformed).float()
+        targets_transformed = Tensor(targets_transformed).float() if targets_transformed is not None else targets
+
+        args = (preds, targets)
+        transformed_args = (preds_transformed, targets_transformed)
+        assert metric(*transformed_args) == wrapped_metric(*args)
+
+    @pytest.mark.parametrize(_test_signature, _test_cases)
+    def test_update(
+        self, cls, transform_pred, transform_target, preds, preds_transformed, targets, targets_transformed
+    ) -> None:
+        """Tests if the binarized update matches the output of the metric on manually binarized targets."""
+        metric = cls()
+        wrapped_metric = LambdaInputTransformer(
+            metric, transform_pred=transform_pred, transform_target=transform_target
+        )
+        preds = Tensor(preds).float()
+        targets = Tensor(targets).float()
+        preds_transformed = Tensor(preds_transformed).float() if preds_transformed is not None else preds
+        targets_transformed = Tensor(targets_transformed).float() if targets_transformed is not None else targets
+
+        args = (preds, targets)
+        transformed_args = (preds_transformed, targets_transformed)
+        metric.update(*transformed_args)
+        wrapped_metric.update(*args)
+        assert metric.compute() == wrapped_metric.compute()
+
+    def test_no_transform_pred(self) -> None:
+        """Tests that TypeError is raised when a non-callable is passed as `transform_pred`."""
+        with pytest.raises(TypeError, match=r"Expected `transform_pred` to be of type .*"):
+            LambdaInputTransformer(BinaryAccuracy(), transform_pred=[])
+
+    def test_no_transform_target(self) -> None:
+        """Tests that TypeError is raised when a non-callable is passed as `transform_target`."""
+        with pytest.raises(TypeError, match=r"Expected `transform_target` to be of type .*"):
+            LambdaInputTransformer(BinaryAccuracy(), transform_target=[])
+
+
+class TestBinaryTargetTransformer:
+    """Test class for BinaryTargetTransformer."""
+
+    _test_signature: ClassVar = ("cls", "threshold", "preds", "targets", "targets_binary", "kwargs")
+    _test_cases: ClassVar = [
+        # Metric with targets and with kwargs
+        (
+            RetrievalMAP,
+            0,
+            [0.9, 0.8, 0.7, 0.6, 0.5, 0.6, 0.7, 0.8, 0.5, 0.4],
+            [1.0, 0.0, 0.0, 0.0, 0.0, 2.0, 1.0, 0.0, 0.0, 0.0],
+            [1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0],
+            {"indexes": [0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0]},
+        ),
+        (
+            RetrievalMAP,
+            1,
+            [0.9, 0.8, 0.7, 0.6, 0.5, 0.6, 0.7, 0.8, 0.5, 0.4],
+            [1.0, 0.0, 0.0, 0.0, 0.0, 2.0, 1.0, 0.0, 0.0, 0.0],
+            [0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0],
+            {"indexes": [0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0]},
+        ),
+        (
+            RetrievalMAP,
+            2,
+            [0.9, 0.8, 0.7, 0.6, 0.5, 0.6, 0.7, 0.8, 0.5, 0.4],
+            [1.0, 0.0, 0.0, 0.0, 0.0, 2.0, 1.0, 0.0, 0.0, 0.0],
+            [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+            {"indexes": [0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0]},
+        ),
+        # Metric with targets and without kwargs
+        (
+            BinaryAccuracy,
+            0,
+            [0.9, 0.8, 0.9, 0.6, 0.5, 0.4, 0.4, 0.1, 0.5, 0.4],
+            [1.0, 0.0, 0.0, 0.0, 0.0, 2.0, 1.0, 0.0, 0.0, 0.0],
+            [1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0],
+            {},
+        ),
+        (
+            BinaryAccuracy,
+            1,
+            [0.9, 0.8, 0.9, 0.6, 0.5, 0.4, 0.4, 0.1, 0.5, 0.4],
+            [1.0, 0.0, 0.0, 0.0, 0.0, 2.0, 1.0, 0.0, 0.0, 0.0],
+            [0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0],
+            {},
+        ),
+        (
+            BinaryAccuracy,
+            2,
+            [0.9, 0.8, 0.9, 0.6, 0.5, 0.4, 0.4, 0.1, 0.5, 0.4],
+            [1.0, 0.0, 0.0, 0.0, 0.0, 2.0, 1.0, 0.0, 0.0, 0.0],
+            [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+            {},
+        ),
+        # Metric without targets and without kwargs
+        (MeanMetric, 2, [0.9, 0.8, 0.9, 0.6, 0.5, 0.4, 0.4, 0.1, 0.5, 0.4], None, None, {}),
+    ]
+
+    @pytest.mark.parametrize(_test_signature, _test_cases)
+    def test_forward(self, cls, threshold, preds, targets, targets_binary, kwargs) -> None:
+        """Tests if the binarized forward matches the output of the metric on manually binarized targets."""
+        metric = cls()
+        wrapped_metric = BinaryTargetTransformer(metric, threshold=threshold)
+        preds = Tensor(preds).float()
+        targets = Tensor(targets).float() if targets is not None else None
+        targets_binary = Tensor(targets_binary).float() if targets_binary is not None else None
+
+        args = (preds, targets) if targets is not None else (preds,)
+        wrapped_args = (preds, targets_binary) if targets_binary is not None else (preds,)
+        kwargs = {k: Tensor(v).long() for k, v in kwargs.items()}
+        assert metric(*wrapped_args, **kwargs) == wrapped_metric(*args, **kwargs)
+
+    @pytest.mark.parametrize(_test_signature, _test_cases)
+    def test_update(self, cls, threshold, preds, targets, targets_binary, kwargs) -> None:
+        """Tests if the binarized update matches the output of the metric on manually binarized targets."""
+        metric = cls()
+        wrapped_metric = BinaryTargetTransformer(metric, threshold=threshold)
+        preds = Tensor(preds).float()
+        if targets is not None:
+            targets = Tensor(targets).float()
+            targets_binary = Tensor(targets_binary).float()
+
+        args = (preds, targets) if targets is not None else [preds]
+        wrapped_args = (preds, targets_binary) if targets_binary is not None else [preds]
+        kwargs = {k: Tensor(v).long() for k, v in kwargs.items()}
+        metric.update(*wrapped_args, **kwargs)
+        wrapped_metric.update(*args, **kwargs)
+        assert metric.compute() == wrapped_metric.compute()
+
+    def test_threshold(self) -> None:
+        """Tests that TypeError is raised when invalid threshold is passed."""
+        with pytest.raises(TypeError, match=r"Expected `threshold` to be of type .*"):
+            BinaryTargetTransformer(RetrievalMAP(), "a")