PyTorch中如何向量化计算两个3D张量两两的平均绝对差

完全可以通过PyTorch广播机制实现无Python循环的向量化计算，效率远高于for循环写法。

原有代码问题 #

你之前的维度调整逻辑错误：将B变形为(1,10,28,28)后，和3维的A(20,28,28)做广播对齐时，无法生成A、B所有图像的配对组合，最终仅得到对角线位置的计算结果，不符合(20,10)全配对的输出要求。

正确实现方法 #

广播机制要求两个张量从最后一维开始向前对齐，对应维度要么长度相等，要么其中一个长度为1（该维度会自动扩展复制到匹配长度）。要生成所有图像对的差张量，需要先调整两个张量的维度：

• 给A在第1维插入长度为1的维度，形状从(20,28,28)转为(20, 1, 28, 28)
• 给B在第0维插入长度为1的维度，形状从(10,28,28)转为(1, 10, 28, 28)

两个张量做差时会自动广播为(20,10,28,28)的4维张量，其中位置(i,j,:,:)对应A中第i张图和B中第j张图的逐像素差值，之后在最后两个图像维度求平均绝对差，就能得到目标形状为(20,10)的结果矩阵。

代码实现如下，两种写法等价：

# 用None插入维度的简写方式
L1_vec = (A[:, None] - B[None]).abs().mean(dim=(-2, -1))

# 用unsqueeze插入维度的显式写法，可读性更强
L1_vec = (A.unsqueeze(1) - B.unsqueeze(0)).abs().mean(dim=(-2, -1))

结果验证 #

执行以下代码可验证向量化结果和原有for循环结果完全一致（浮点计算精度误差范围内）：

print(torch.allclose(L1, L1_vec))  # 输出为True

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