PyTorch无循环计算成对欧氏距离：张量形状错误求助

问题描述 #

我正在自学MIT深度学习与计算机视觉课程（EECS 498-007 / 598-005）的作业1，该作业与斯坦福CS 231n相关内容类似。需求是实现一个函数，计算成对欧氏距离：输入为维度[N,x,x]的xtrain和[M,x,x]的xtest，输出[N,M]的距离矩阵，作业提示需用两次广播求和与一次矩阵乘法实现。

我尝试基于广播实现该运算，但遇到张量形状问题，当前代码如下：

def euc_no_loop(x,y):
  #hint: two broadcast sums
  xsq = torch.sum(x**2,axis=1)
  print(xsq.shape)
  ysq = torch.sum(y**2,axis=1)
  print(ysq.shape)
  #and one matrix multiply
  mixprod = -2 * x.view(x.shape[0],-1).matmul(y.view(y.shape[0],-1).T)
  print(mixprod.shape)
  euc_dist = torch.sqrt(xsq + mixprod + ysq.unsqueeze(1).T)
  return euc_dist

输入示例：

x = torch.randn(5,3,3)
y = torch.randn(3,3,3)

运行后各张量形状：

• xsq: [5,3]
• ysq: [3,3]
• mixprod: [5,3]
  最终输出维度为[3,5,3]，不符合[N,M]要求。参考过NumPy实现，但torch.matmul灵活性不如numpy.dot，不清楚问题所在，请求帮助排查。

问题分析 #

核心错误出在求和维度和张量维度调整上：

• 计算xsq和ysq时，仅在axis=1求和，未覆盖所有空间维度，导致得到的是[N, x]和[M, x]的张量，而非每个样本的总平方和（应为[N]和[M]）。
• 后续的广播拼接因维度不匹配，导致最终输出形状错误。

修正后的代码 #

import torch

def euc_no_loop(x, y):
    # 对所有空间维度求和，得到每个样本的平方和，形状[N]和[M]
    xsq = torch.sum(x ** 2, dim=(1, 2))
    ysq = torch.sum(y ** 2, dim=(1, 2))
    
    # 将张量展平为[N, D]和[M, D]，D为x*x，矩阵乘法得到[N, M]的交叉项
    x_flat = x.view(x.shape[0], -1)
    y_flat = y.view(y.shape[0], -1)
    mixprod = -2 * x_flat.matmul(y_flat.T)
    
    # 利用广播机制完成求和：xsq[:, None]转为[N,1]，与ysq[M]、mixprod[N,M]广播为[N,M]
    euc_dist = torch.sqrt(xsq[:, None] + mixprod + ysq)
    return euc_dist

验证结果 #

运行输入示例：

x = torch.randn(5,3,3)
y = torch.randn(3,3,3)
dist = euc_no_loop(x, y)
print(dist.shape)  # 输出 torch.Size([5, 3])，符合[N,M]要求

关键细节说明 #

• 求和维度修正：dim=(1,2)指定对输入张量的后两个维度求和，将[N,x,x]压缩为[N]，每个元素对应一个样本所有元素的平方和。
• 广播适配：xsq[:, None]将一维张量转为[N,1]，PyTorch会自动将其与ysq（[M]）、mixprod（[N,M]）广播为相同的[N,M]形状，确保每对样本的距离计算正确。
• 矩阵乘法逻辑：展平后的x_flat是[N,D]，y_flat.T是[D,M]，矩阵乘法结果正好是[N,M]的交叉项矩阵，完全匹配成对距离的计算需求。

内容的提问来源于stack exchange，提问作者Martin ægidius