Pytorch中计算余弦相似度、欧式距离、范数 (捋清pairwise distance, norm, 详解cdist)

设a,b分别为两个tensor

import torch
import torch.nn as nn
a = torch.tensor([1,2],dtype=float)
b = torch.tensor([5,7],dtype=float)

余弦相似度

余弦相似度非常简单

cos_sim = nn.CosineSimilarity(dim=0, eps=1e-6)
sim = cos_sim(a,b)
print(sim) # tensor(0.9878, dtype=torch.float64)

欧式距离

这个欧氏距离实现起来也很方便，不过用Pytorch有很多种实现方式，顺便帮大家捋清楚torch.PairwiseDistance、torch.linalg.vector_norm、torch.cdist的异同

torch.PairwiseDistance

设$x=a-b$
那么这个pairwise distance输出的结果就是$x$的范数： ${\Vert x\Vert }_p$
$${\Vert x\Vert }_p={\left(\sum _{i=1}^n{\left|x_i\right|}^p\right)}^{1/p}$$
说人话就是$x$里面元素各$p$次方加和再开$p$次方，所以算欧式距离那就是$p=2$的时候：

pdist = nn.PairwiseDistance(p=2)
output = pdist(a,b)
print(output) # tensor(6.4031, dtype=torch.float64)

torch.linalg.vector_norm

这个函数的话只是计算某个向量的norm，就是上面 ${\Vert x\Vert }_p$的式子。
那我们手动求一下$x=a-b$这部分就行了：

from torch import linalg as LA
x = a-b
output = LA.vector_norm(x, ord=2)
print(output) # tensor(6.4031, dtype=torch.float64)

可以看到这个和torch.PairwiseDistance求出来的值是一毛一样的

torch.cdist

这个也可以像PairwiseDistance那样直接用于计算p-norm distance，但它复杂一些。
我们先来理解一下torch.cdist
它的官方文档里是这样说的 “Computes batched the p-norm distance between each pair of the two collections of row vectors.”¹

注意一下我画出来的部分，下面我会仔细讲
设它的输入是m1,m2
m1的shape是(B,P,M)
m2的shape是(B,R,M)
那么这个B，就是batch size, 对应的是官方说明里所说的batched。
这个M，就是官方说明里row vectors的长度，这个长度是需要m1,m2是一样的，不然怎么做p-norm的计算呢？
P和R只是m1, m2的每个batch中，各自有的row vectors的数量，这个可以不一样的，因为对于每一个batch，m1里的每个row vector都会和m2里的每个row vector进行距离的计算。也就是说你m1里有P个row vectors, m2里有R个row vectors的话，就会进行P*R次计算。
因此，cdist计算后的输出形状会是(B,P,R)。

举个例子：

m1 = torch.tensor([[[1.5,2.5],[1.0,0.5],[3.5,5.5]],[[6.5,7.0],[2.0,3.0],[5.0,8.0]]]) 
# shape为2,3,2

m2 = torch.tensor([[[4.5,0.5],[9.0,10.0]],[[4.0,6.0],[2.5,3.2]]]) 
# shape为2,2,2

output = torch.cdist(m1,m2,p=2)
print(output)
'''
tensor([[[ 3.6056, 10.6066],
         [ 3.5000, 12.4197],
         [ 5.0990,  7.1063]],

        [[ 2.6926,  5.5172],
         [ 3.6056,  0.5385],
         [ 2.2361,  5.4120]]])
'''

我们来看一下这是怎么算出来的
对于第一个batch，我们有(P,M)和(R,M)大小的两个tensor

[[1.5,2.5],[1.0,0.5],[3.5,5.5]] #shape 3,2 来自m1的第一个batch
[[4.5,0.5],[9.0,10.0]] #shape 2,2 来自m2的第一个batch

这里的每一个row vector, 也就是像[1.5,2.5]这样的东西，会进行P*R次的组合。
比如[1.5,2.5]会分别和[4.5,0.5]、[9.0,10.0]都计算一次欧氏距离。我们来手动用torch.PairwiseDistance算一下：

g = torch.tensor([1.5,2.5])
h = torch.tensor([4.5,0.5])
i = torch.tensor([9.0,10.0])
print(pdist(g,h)) # tensor(3.6056)
print(pdist(g,i)) # tensor(10.6066)

这两个值不就正好对应刚才cdist计算出的第一个batch的第一行那两列的元素吗？
这样，就理解了为什么cdist计算后的输出形状会是(B,P,R)。

ok，理解了cdist，我们用它来计算一下开头的a,b的欧氏距离（实际上这个属于"可以但没有必要"的一个计算hhh，因为用cdist算这个好像不太必要哈）
cdist需要你的输入至少是2D，如果你这样的话：

output = torch.cdist(a, b, p=2)

会报错：“cdist only supports at least 2D tensors, X1 got: 1D”

那就先unsqueeze一下：

output = torch.cdist(a.unsqueeze(0), b.unsqueeze(0), p=2)
print(output) # tensor([[6.4031]], dtype=torch.float64)

可以注意到现在的输出的shape是(1,1),因为假如我们默认batch是1，那你a和b的shape相当于是(1,1,2)和(1,1,2), P和R的值分别是1和1，所以输出的shape也就是(1,1)了。

---

1. https://pytorch.org/docs/stable/generated/torch.cdist.html?highlight=cdist#torch.cdist ↩︎