如何用单循环计算两组张量两两欧氏距离并填充对应矩阵

你现有代码的问题出在torch.sum没有指定求和维度：train形状为(num_train, 特征维度)，test[i]形状为(特征维度,)，二者相减得到的张量形状是(num_train, 特征维度)，直接调用无维度参数的torch.sum会把整个张量所有元素累加为单个标量，自然会出现整列填充同一个值的问题。

仅需给torch.sum指定dim=1，沿特征维度求和即可，同时展平张量的写法也可以简化，避免索引取shape出错：

# 展平：保留第0维的batch维度，后面所有维度展平为一维
train = x_train.flatten(start_dim=1)
test = x_test.flatten(start_dim=1)
num_train = train.shape[0]
num_test = test.shape[0]
# 初始化距离矩阵 形状为(num_train, num_test)
dists = torch.zeros(num_train, num_test, device=train.device)

for i in range(num_test):
    # dim=1 沿特征维度求和，得到形状为(num_train,)的张量，对应第i列所有行的距离
    dists[:,i] = torch.sqrt(torch.sum(torch.square(train - test[i]), dim=1))

如果要避免循环、提升计算速度，可以用广播机制实现完全向量化的距离计算：

train = x_train.flatten(1)
test = x_test.flatten(1)
# 利用欧氏距离展开公式 (a-b)² = a² + b² - 2ab
dist_sq = torch.sum(train**2, dim=1).unsqueeze(1) + torch.sum(test**2, dim=1).unsqueeze(0) - 2 * train @ test.T
dists = torch.sqrt(dist_sq)

该实现得到的dists形状同样为(num_train, num_test)，其中dists[m][n]就是train第m个样本和test第n个样本的欧氏距离，完全匹配你需要的填充规则。

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