如何使用PyTorch高效计算k[i][j] = exp( -(a[i]-b[j])² )形式的协方差矩阵？

PyTorch高效实现方案 #

你可以直接利用PyTorch的广播机制实现和NumPy版本完全等价的向量化运算，完全避免Python层的循环，代码如下：

import torch

# 假设a是形状为(N,)的一维张量，b是形状为(M,)的一维张量
r = a[:, None] - b[None, :]  # 广播后形状为(N, M)，每个元素为a[i]-b[j]
k = torch.exp(-r ** 2)

如果需要完全对齐NumPy的subtract.outer写法，PyTorch 1.13及以上版本也支持torch.subtract.outer接口，直接写也可以：

import torch

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

性能说明 #

你之前的双循环版本慢的核心原因是所有运算都在Python层逐元素执行，无法利用PyTorch底层的向量化优化、多CPU并行以及GPU加速能力。上述向量化实现的运算逻辑完全在PyTorch底层执行，性能和NumPy版本基本持平，如果将张量迁移到GPU上运行，速度会远超过NumPy的CPU版本。
如果后续需要处理批量的高维张量计算核矩阵，只要对上述广播逻辑做对应维度调整即可，不需要修改核心运算逻辑。

内容的提问来源于stack exchange，提问作者kkk