spearman_corrcoef

[AhmedHussKhalifa]

Hi,
I am trying to calculate the spearman coefficient during training on the GPU, and it is slowing down. Is there any advice to speed it up?

torchmetrics.functional import spearman_corrcoef

[mhorlacher]

You could try my custom implementation for PCC and change it to SCC.

import torch
import torchmetrics

def pearson_corrcoef(x, y, dim=-1):
    x = x - torch.unsqueeze(torch.mean(x, dim), dim)
    y = y - torch.unsqueeze(torch.mean(y, dim), dim)
    return torch.sum(x * y, dim) / torch.sqrt(torch.sum(x ** 2, dim) * torch.sum(y ** 2, dim))

class PearsonCorrCoeff(torchmetrics.MeanMetric):
    def __init__(self, dim=-1, reduction=torch.mean, *args, **kwargs):
        super(PearsonCorrCoeff,  self).__init__(*args, **kwargs)

        self.dim = dim
        self.reduction = reduction

    def update(self, y: torch.Tensor, y_pred: torch.Tensor):
        assert y.shape == y_pred.shape

        pcc = pearson_corrcoef(y, y_pred, dim=self.dim)
        pcc = torch.nan_to_num(pcc, 0.0) # replace nan's with 0 (this might underestimate the pcc)

        reduced_pcc = self.reduction(pcc)

        # update (i.e. take mean)
        super().update(reduced_pcc)