2v2 gpu

2v2_gpu.py

class Time2v2():
    def __init__(self, metric='cosine', scale=True, to=None):
        self.metric = metric
        self.to = to
        self.scale = scale

    def __call__(self, y_true, y_pred):
        from torch.nn import CosineSimilarity
        from numpy.random import permutation

        # define permutations
        assert len(y_true) == len(y_pred)
        ns = len(y_true)
        first = permutation(ns)  # first group of TR
        second = permutation(ns)  # second group of TR
        while (first == second).any():  # check that distinct TRs in pairs
            first[first == second] = np.random.choice((first == second).sum())

        # Prepare tensors
        cos = CosineSimilarity(dim=1, eps=1e-08)
        y_true = torch.tensor(y_true)
        y_pred = torch.tensor(y_pred)
        if torch.cuda.is_available() and self.to in (None, 'cuda'):
            y_true = y_true.to('cuda')
            y_pred = y_pred.to('cuda')

        # scale
        if self.scale:
            y_true -= y_true.mean(0)
            y_pred -= y_pred.mean(0)
            y_true /= y_true.std(0)
            y_pred /= y_pred.std(0)

        # compute cosine distance across dim == 1
        s1 = 1. - cos(y_true[first], y_pred[first])
        s1 += 1. - cos(y_true[second], y_pred[second])
        s2 = 1. - cos(y_true[first], y_pred[second])
        s2 += 1. - cos(y_true[second], y_pred[first])

        acc = (1.*(s1 < s2)).mean(0).cpu().numpy()

        return acc