Support multioutput pearson and spearman corrcoef

CHANGELOG.md

@@ -20,6 +20,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added argument `normalize` to `LPIPS` metric ([#1216](https://github.com/Lightning-AI/metrics/pull/1216))
 
 
+- Added support for multioutput in `PearsonCorrCoef` and `SpearmanCorrCoef` ([#1200](https://github.com/Lightning-AI/metrics/pull/1200))
+
+
 ### Changed
 
 - Classification refactor (

src/torchmetrics/functional/regression/pearson.py

@@ -28,6 +28,7 @@ def _pearson_corrcoef_update(
     var_y: Tensor,
     corr_xy: Tensor,
     n_prior: Tensor,
+    n_out: int,
 ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:
     """Updates and returns variables required to compute Pearson Correlation Coefficient.
 
@@ -43,18 +44,24 @@ def _pearson_corrcoef_update(
     """
     # Data checking
     _check_same_shape(preds, target)
-    preds = preds.squeeze()
-    target = target.squeeze()
-    if preds.ndim > 1 or target.ndim > 1:
-        raise ValueError("Expected both predictions and target to be 1 dimensional tensors.")
-
-    n_obs = preds.numel()
-    mx_new = (n_prior * mean_x + preds.mean() * n_obs) / (n_prior + n_obs)
-    my_new = (n_prior * mean_y + target.mean() * n_obs) / (n_prior + n_obs)
+    if preds.ndim > 2 or target.ndim > 2:
+        raise ValueError(
+            "Expected both predictions and target to be either 1 or 2 dimensional tensors,"
+            " but get{target.ndim} and {preds.ndim}."
+        )
+    if (n_out == 1 and preds.ndim != 1) or (n_out > 1 and n_out != preds.shape[-1]):
+        raise ValueError(
+            "Expected argument `num_outputs` to match the second dimension of input, but got {self.n_out}"
+            " and {preds.ndim}."
+        )
+
+    n_obs = preds.shape[0]
+    mx_new = (n_prior * mean_x + preds.mean(0) * n_obs) / (n_prior + n_obs)
+    my_new = (n_prior * mean_y + target.mean(0) * n_obs) / (n_prior + n_obs)
     n_prior += n_obs
-    var_x += ((preds - mx_new) * (preds - mean_x)).sum()
-    var_y += ((target - my_new) * (target - mean_y)).sum()
-    corr_xy += ((preds - mx_new) * (target - mean_y)).sum()
+    var_x += ((preds - mx_new) * (preds - mean_x)).sum(0)
+    var_y += ((target - my_new) * (target - mean_y)).sum(0)
+    corr_xy += ((preds - mx_new) * (target - mean_y)).sum(0)
     mean_x = mx_new
     mean_y = my_new
 
@@ -95,9 +102,19 @@ def pearson_corrcoef(preds: Tensor, target: Tensor) -> Tensor:
         >>> preds = torch.tensor([2.5, 0.0, 2, 8])
         >>> pearson_corrcoef(preds, target)
         tensor(0.9849)
+
+    Example (multi output regression):
+        >>> from torchmetrics.functional import pearson_corrcoef
+        >>> target = torch.tensor([[3, -0.5], [2, 7]])
+        >>> preds = torch.tensor([[2.5, 0.0], [2, 8]])
+        >>> pearson_corrcoef(preds, target)
+        tensor([1., 1.])
     """
-    _temp = torch.zeros(1, dtype=preds.dtype, device=preds.device)
+    d = preds.shape[1] if preds.ndim == 2 else 1
+    _temp = torch.zeros(d, dtype=preds.dtype, device=preds.device)
     mean_x, mean_y, var_x = _temp.clone(), _temp.clone(), _temp.clone()
     var_y, corr_xy, nb = _temp.clone(), _temp.clone(), _temp.clone()
-    _, _, var_x, var_y, corr_xy, nb = _pearson_corrcoef_update(preds, target, mean_x, mean_y, var_x, var_y, corr_xy, nb)
+    _, _, var_x, var_y, corr_xy, nb = _pearson_corrcoef_update(
+        preds, target, mean_x, mean_y, var_x, var_y, corr_xy, nb, n_out=1 if preds.ndim == 1 else preds.shape[-1]
+    )
     return _pearson_corrcoef_compute(var_x, var_y, corr_xy, nb)

src/torchmetrics/functional/regression/spearman.py

@@ -52,7 +52,7 @@ def _rank_data(data: Tensor) -> Tensor:
     return rank
 
 
-def _spearman_corrcoef_update(preds: Tensor, target: Tensor) -> Tuple[Tensor, Tensor]:
+def _spearman_corrcoef_update(preds: Tensor, target: Tensor, n_out: int) -> Tuple[Tensor, Tensor]:
     """Updates and returns variables required to compute Spearman Correlation Coefficient.
 
     Checks for same shape and type of input tensors.
@@ -68,10 +68,17 @@ def _spearman_corrcoef_update(preds: Tensor, target: Tensor) -> Tuple[Tensor, Te
             f" Got preds: {preds.dtype} and target: {target.dtype}."
         )
     _check_same_shape(preds, target)
-    preds = preds.squeeze()
-    target = target.squeeze()
-    if preds.ndim > 1 or target.ndim > 1:
-        raise ValueError("Expected both predictions and target to be 1 dimensional tensors.")
+    if preds.ndim > 2 or target.ndim > 2:
+        raise ValueError(
+            "Expected both predictions and target to be either 1 or 2 dimensional tensors,"
+            " but get{target.ndim} and {preds.ndim}."
+        )
+    if (n_out == 1 and preds.ndim != 1) or (n_out > 1 and n_out != preds.shape[-1]):
+        raise ValueError(
+            "Expected argument `num_outputs` to match the second dimension of input, but got {self.n_out}"
+            " and {preds.ndim}."
+        )
+
     return preds, target
 
 
@@ -86,20 +93,23 @@ def _spearman_corrcoef_compute(preds: Tensor, target: Tensor, eps: float = 1e-6)
     Example:
         >>> target = torch.tensor([3, -0.5, 2, 7])
         >>> preds = torch.tensor([2.5, 0.0, 2, 8])
-        >>> preds, target = _spearman_corrcoef_update(preds, target)
+        >>> preds, target = _spearman_corrcoef_update(preds, target, n_out=1)
         >>> _spearman_corrcoef_compute(preds, target)
         tensor(1.0000)
     """
+    if preds.ndim == 1:
+        preds = _rank_data(preds)
+        target = _rank_data(target)
+    else:
+        preds = torch.stack([_rank_data(p) for p in preds.T]).T
+        target = torch.stack([_rank_data(t) for t in target.T]).T
 
-    preds = _rank_data(preds)
-    target = _rank_data(target)
-
-    preds_diff = preds - preds.mean()
-    target_diff = target - target.mean()
+    preds_diff = preds - preds.mean(0)
+    target_diff = target - target.mean(0)
 
-    cov = (preds_diff * target_diff).mean()
-    preds_std = torch.sqrt((preds_diff * preds_diff).mean())
-    target_std = torch.sqrt((target_diff * target_diff).mean())
+    cov = (preds_diff * target_diff).mean(0)
+    preds_std = torch.sqrt((preds_diff * preds_diff).mean(0))
+    target_std = torch.sqrt((target_diff * target_diff).mean(0))
 
     corrcoef = cov / (preds_std * target_std + eps)
     return torch.clamp(corrcoef, -1.0, 1.0)
@@ -119,13 +129,19 @@ def spearman_corrcoef(preds: Tensor, target: Tensor) -> Tensor:
         preds: estimated scores
         target: ground truth scores
 
-    Example:
+    Example (single output regression):
         >>> from torchmetrics.functional import spearman_corrcoef
         >>> target = torch.tensor([3, -0.5, 2, 7])
         >>> preds = torch.tensor([2.5, 0.0, 2, 8])
         >>> spearman_corrcoef(preds, target)
         tensor(1.0000)
 
+    Example (multi output regression):
+        >>> from torchmetrics.functional import spearman_corrcoef
+        >>> target = torch.tensor([[3, -0.5], [2, 7]])
+        >>> preds = torch.tensor([[2.5, 0.0], [2, 8]])
+        >>> spearman_corrcoef(preds, target)
+        tensor([1.0000, 1.0000])
     """
-    preds, target = _spearman_corrcoef_update(preds, target)
+    preds, target = _spearman_corrcoef_update(preds, target, n_out=1 if preds.ndim == 1 else preds.shape[-1])
     return _spearman_corrcoef_compute(preds, target)

src/torchmetrics/regression/pearson.py

@@ -73,20 +73,29 @@ class PearsonCorrCoef(Metric):
 
     Forward accepts
 
-    - ``preds`` (float tensor): ``(N,)``
-    - ``target``(float tensor): ``(N,)``
+    - ``preds`` (float tensor): either single output tensor with shape ``(N,)`` or multioutput tensor of shape ``(N,d)``
+    - ``target``(float tensor): either single output tensor with shape ``(N,)`` or multioutput tensor of shape ``(N,d)``
 
     Args:
+        num_outputs: Number of outputs in multioutput setting
         kwargs: Additional keyword arguments, see :ref:`Metric kwargs` for more info.
 
-    Example:
+    Example (single output regression):
         >>> from torchmetrics import PearsonCorrCoef
         >>> target = torch.tensor([3, -0.5, 2, 7])
         >>> preds = torch.tensor([2.5, 0.0, 2, 8])
         >>> pearson = PearsonCorrCoef()
         >>> pearson(preds, target)
         tensor(0.9849)
 
+    Example (multi output regression):
+        >>> from torchmetrics import PearsonCorrCoef
+        >>> target = torch.tensor([[3, -0.5], [2, 7]])
+        >>> preds = torch.tensor([[2.5, 0.0], [2, 8]])
+        >>> pearson = PearsonCorrCoef(num_outputs=2)
+        >>> pearson(preds, target)
+        tensor([1., 1.])
+
     """
     is_differentiable = True
     higher_is_better = None  # both -1 and 1 are optimal
@@ -102,16 +111,20 @@ class PearsonCorrCoef(Metric):
 
     def __init__(
         self,
+        num_outputs: int = 1,
         **kwargs: Any,
     ) -> None:
         super().__init__(**kwargs)
+        if not isinstance(num_outputs, int) and num_outputs < 1:
+            raise ValueError("Expected argument `num_outputs` to be an int larger than 0, but got {num_outputs}")
+        self.n_out = num_outputs
 
-        self.add_state("mean_x", default=torch.tensor(0.0), dist_reduce_fx=None)
-        self.add_state("mean_y", default=torch.tensor(0.0), dist_reduce_fx=None)
-        self.add_state("var_x", default=torch.tensor(0.0), dist_reduce_fx=None)
-        self.add_state("var_y", default=torch.tensor(0.0), dist_reduce_fx=None)
-        self.add_state("corr_xy", default=torch.tensor(0.0), dist_reduce_fx=None)
-        self.add_state("n_total", default=torch.tensor(0.0), dist_reduce_fx=None)
+        self.add_state("mean_x", default=torch.zeros(self.n_out), dist_reduce_fx=None)
+        self.add_state("mean_y", default=torch.zeros(self.n_out), dist_reduce_fx=None)
+        self.add_state("var_x", default=torch.zeros(self.n_out), dist_reduce_fx=None)
+        self.add_state("var_y", default=torch.zeros(self.n_out), dist_reduce_fx=None)
+        self.add_state("corr_xy", default=torch.zeros(self.n_out), dist_reduce_fx=None)
+        self.add_state("n_total", default=torch.zeros(self.n_out), dist_reduce_fx=None)
 
     def update(self, preds: Tensor, target: Tensor) -> None:  # type: ignore
         """Update state with predictions and targets.
@@ -120,13 +133,27 @@ def update(self, preds: Tensor, target: Tensor) -> None:  # type: ignore
             preds: Predictions from model
             target: Ground truth values
         """
+        if (self.n_out == 1 and preds.ndim != 1) or (self.n_out > 1 and self.n_out != preds.shape[-1]):
+            raise ValueError(
+                "Expected argument `num_outputs` to match the second dimension of input, but got {self.num_outputs}"
+                " and {preds.ndim}."
+            )
         self.mean_x, self.mean_y, self.var_x, self.var_y, self.corr_xy, self.n_total = _pearson_corrcoef_update(
-            preds, target, self.mean_x, self.mean_y, self.var_x, self.var_y, self.corr_xy, self.n_total
+            preds,
+            target,
+            self.mean_x,
+            self.mean_y,
+            self.var_x,
+            self.var_y,
+            self.corr_xy,
+            self.n_total,
+            self.n_out,
         )
 
     def compute(self) -> Tensor:
         """Computes pearson correlation coefficient over state."""
-        if self.mean_x.numel() > 1:  # multiple devices, need further reduction
+        if (self.n_out == 1 and self.mean_x.numel() > 1) or (self.n_out > 1 and self.mean_x.ndim > 1):
+            # multiple devices, need further reduction
             var_x, var_y, corr_xy, n_total = _final_aggregation(
                 self.mean_x, self.mean_y, self.var_x, self.var_y, self.corr_xy, self.n_total
             )

src/torchmetrics/regression/spearman.py

@@ -32,17 +32,31 @@ class SpearmanCorrCoef(Metric):
     Spearmans correlations coefficient corresponds to the standard pearsons correlation coefficient calculated
     on the rank variables.
 
+    Forward accepts
+
+    - ``preds`` (float tensor): ``(N,d)``
+    - ``target``(float tensor): ``(N,d)``
+
     Args:
+        num_outputs: Number of outputs in multioutput setting
         kwargs: Additional keyword arguments, see :ref:`Metric kwargs` for more info.
 
-    Example:
+    Example (single output regression):
         >>> from torchmetrics import SpearmanCorrCoef
         >>> target = torch.tensor([3, -0.5, 2, 7])
         >>> preds = torch.tensor([2.5, 0.0, 2, 8])
         >>> spearman = SpearmanCorrCoef()
         >>> spearman(preds, target)
         tensor(1.0000)
 
+    Example (multi output regression):
+        >>> from torchmetrics import SpearmanCorrCoef
+        >>> target = torch.tensor([[3, -0.5], [2, 7]])
+        >>> preds = torch.tensor([[2.5, 0.0], [2, 8]])
+        >>> spearman = SpearmanCorrCoef(num_outputs=2)
+        >>> spearman(preds, target)
+        tensor([1.0000, 1.0000])
+
     """
     is_differentiable: bool = False
     higher_is_better: bool = True
@@ -52,13 +66,17 @@ class SpearmanCorrCoef(Metric):
 
     def __init__(
         self,
+        num_outputs: int = 1,
         **kwargs: Any,
     ) -> None:
         super().__init__(**kwargs)
         rank_zero_warn(
             "Metric `SpearmanCorrcoef` will save all targets and predictions in the buffer."
             " For large datasets, this may lead to large memory footprint."
         )
+        if not isinstance(num_outputs, int) and num_outputs < 1:
+            raise ValueError("Expected argument `num_outputs` to be an int larger than 0, but got {num_outputs}")
+        self.n_out = num_outputs
 
         self.add_state("preds", default=[], dist_reduce_fx="cat")
         self.add_state("target", default=[], dist_reduce_fx="cat")
@@ -70,7 +88,7 @@ def update(self, preds: Tensor, target: Tensor) -> None:  # type: ignore
             preds: Predictions from model
             target: Ground truth values
         """
-        preds, target = _spearman_corrcoef_update(preds, target)
+        preds, target = _spearman_corrcoef_update(preds, target, n_out=self.n_out)
         self.preds.append(preds)
         self.target.append(target)