有没有使用 PyTorch 计算协方差矩阵的有效方法？答案

【问题标题】：Is there any efficient way to calculate covariance matrix using PyTorch?有没有使用 PyTorch 计算协方差矩阵的有效方法？
【发布时间】：2021-11-02 16:28:35
【问题描述】：

我想从向量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 编写高效的代码？

【问题讨论】：

标签： python performance pytorch

【解决方案1】：

你可以使用广播：

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

【讨论】：

这是一个很好的解决方案，但我认为 reshape 稍微快一点！ i.stack.imgur.com/yeh5H.png
@UlisesBussi 这很奇怪，添加单个维度是 O(1) 操作。创建 temp var r 是否有可能产生运行时效果？
不，我已经尝试使用辅助 r 来做到公平......如果你愿意，我可以将此代码添加到我的答案中。并显示结果（我正在审查以检查这不是人为错误）
已解决，您编辑过吗？你以前用过 k=np.exp 吗？这就是 np.exp 似乎比 torch.exp 慢的问题，现在我的时间几乎相同！
@UlisesBussi 我没有编辑。使用timeit 我发现没有明显的时差。

【解决方案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)

【讨论】：

广播对 numpy 也有效。事实上，许多切片/索引技巧在 pytorch 和 numpy 之间共享。