pytorch 快速计算两个tensor的欧式距离

给定两个tensor： A 和 B。A的维度为 [m，h],  B的维度为 [n，h]。要求每两行之间计算他们的欧氏距离，返回一个维度为[m，n]的tensor。要求不能使用循环，也不能扩充复制一个tensor。

解：将欧式距离的公式展开，根号下面是。a是tensor A中某一行的平方和，b是tensor B中某一行的平方和，ab是两个向量的内积。

所以代码为：

def EuclideanDistances(a,b):
    sq_a = a**2
    sum_sq_a = torch.sum(sq_a,dim=1).unsqueeze(1)  # m->[m, 1]
    sq_b = b**2
    sum_sq_b = torch.sum(sq_b,dim=1).unsqueeze(0)  # n->[1, n]
    bt = b.t()
    return torch.sqrt(sum_sq_a+sum_sq_b-2*a.mm(bt))


a = torch.rand(3,5)
b = torch.rand(4,5)
print(EuclideanDistances(a,b))

 

解释：其中问题关键就在于[m,1 ]+[1, n] 变成了一个[m,  n]

 

例如：