有什么有效的方法可以使用 PyTorch 计算协方差矩阵吗？

Question

我想从向量 a 和 b 计算协方差矩阵，例如 k[i][j] = exp( -(a[i]-b[j])**2 )。

在numpy中，我可以这样写，

import numpy as np

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

在 PyTorch 中，我可以编写简单的代码，但它比 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 )

我应该如何使用 PyTorch 编写高效的代码？

Answer 1

您可以使用广播：

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

Answer 2

我会使用重塑和乘法 ndims 数组的产品：

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

编辑

此方法也适用于 numpy:

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

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