Is there any efficient way to calculate covariance matrix using PyTorch?-CodePudding

I want to calculate covariance matrix from vectors a and b, like k[i][j] = exp( -(a[i]-b[j])**2 ).

In numpy, I can write as follows,

import numpy as np

r = np.subtract.outer(a, b)
k = np.exp(-r*r)

In PyTorch, I can write naive code, but it's slower than numpy.

import torch

for i in range(len(a)):
    for j in range(len(b)):
        k[i][j] = torch.exp( -(a[i]-b[j])**2 )

How should I write efficient code using PyTorch?

CodePudding user response：

You can use broadcasting:

r = a[:, None] - b[None, :]
k = torch.exp(-r**2)

CodePudding user response：

I would use the goods of reshape and multipliing ndims arrays:

k = torch.exp(- (a.reshape(-1,1)*b.reshape(1,-1))**2)

EDIT

Also this method is valid with numpy:

k = np.exp(- (a.reshape(-1,1)*b.reshape(1,-1))**2)