矩阵之间的行向量两两距离
给定数据矩阵
A
∈
R
n
×
d
A\in R^{n\times d}
A∈Rn×d和矩阵
B
∈
R
m
×
d
B\in R^{m\times d}
B∈Rm×d ,A和B中每一行都是一个数据点,现在要去求A中所有元素和B中所有元素之间的欧氏距离。即计算矩阵
D
=
(
d
i
j
=
∣
∣
a
i
,
:
−
b
j
,
:
∣
∣
)
∈
R
n
×
m
D =(d_{ij} = ||a_{i,:} - b_{j, :}||) \in R^{n\times m}
D=(dij=∣∣ai,:−bj,:∣∣)∈Rn×m。
直接去做
n
×
m
n\times m
n×m的两层循环当然可以得到结果,但是既然数据都已经放在矩阵里了,它自然可以“并行”地去计算,好利用一些现有的矩阵计算库(比如numpy
、torch
)等的优化。
具体来讲,把欧氏距离的运算展开
∣
∣
a
i
,
:
−
b
j
,
:
∣
∣
2
=
(
a
i
,
:
−
b
j
,
:
)
2
=
a
i
,
:
2
+
b
j
,
:
2
−
2
a
i
,
:
⋅
b
j
,
:
\begin{aligned} ||a_{i,:}-b_{j,:}||^2 & =(a_{i,:}-b_{j,:})^2 \\ & = a_{i,:}^2 + b_{j,:}^2 - 2a_{i,:}\cdot b_{j,:} \end{aligned}
∣∣ai,:−bj,:∣∣2=(ai,:−bj,:)2=ai,:2+bj,:2−2ai,:⋅bj,:
观察可以知道,式子中的三项分别等于A中第i行的平方和、B中第j行的平方和、A中第i行和B中第j行的内积,这样的话,可以分别对A和B做逐元素平方求和,而向量两两做内积,得到的结果其实就是A与B转置相乘。矩阵相乘作为一个经典操作那是被优化得好好的,比起自己去做两层循环方便得多。
整个过程就成了 R o w S u m ( A . ∗ A ) T + R o w S u m ( B . ∗ B ) − 2 A ⋅ B T RowSum(A .* A)^T + RowSum(B .* B) - 2A\cdot B^T RowSum(A.∗A)T+RowSum(B.∗B)−2A⋅BT 。
这里有个问题,A平方和之后转置得到的是 n × 1 n\times 1 n×1 的向量,B平方和得到的是 1 × m 1\times m 1×m的向量,要把它们都扩展成 n × m n\times m n×m的向量才好相加。幸运的是,现在常见的矩阵运算库基本都是有“广播”功能的。当向量 v ∈ R n × 1 v \in R^{n\times 1} v∈Rn×1和向量 u ∈ R 1 × m u\in R^{1\times m} u∈R1×m 相加时,它们分别会被自动扩展成为矩阵 V ∈ R n × m V\in R^{n\times m} V∈Rn×m和 U ∈ R n × m U\in R^{n\times m} U∈Rn×m ,显然可以想到V就是把v复制了m列,U是把u复制了n行。通过自动广播,维度的处理就完全不是问题了。
具体代码如下
# numpy version
A2 = (A ** 2).sum(axis=1).reshape((-1, 1)) # 转为n*1
B2 = (B ** 2).sum(axis=1) # 这里不用reshape了,(n,)和(1,n)的向量在计算过程中效果是一样的
D = A2 + B2 - 2 * A.dot(B.transpose())
# D = np.clamp(D, a_min=0) # 有时候会出现精度问题导致结果小于0 ,我习惯于在这里进行一次校正
D = D.sqrt() # 开方得到欧氏距离
# ...
# torch version
A2 = (A ** 2).sum(dim=1).reshape((-1, 1))
B2 = (B ** 2).sum(dim=1)
D = A2 + B2 - 2 * A.mm(B.t())
# ...
实际场景经常会遇到一个点矩阵 A ∈ R a × d A\in R^{a\times d} A∈Ra×d 需要求所有数据点之间的两两距离,此时就可以用上面的方式去计算A和A自身的两两距离。
批量计算点集与点集之间的两两距离
今天在做实验时遇到这样一个场景,
X
∈
R
B
×
n
×
d
X\in R^{B\times n\times d}
X∈RB×n×d 和
Y
∈
R
B
×
m
×
d
Y\in R^{B\times m\times d}
Y∈RB×m×d,其实这也是玩深度学习时常会出现到的情况,就是数据集被划分为一个一个的Batch,这里我想对每一个batch b分别计算
X
[
b
,
:
,
:
]
X[b,:, :]
X[b,:,:] 和
Y
[
b
,
:
,
:
]
Y[b,:, :]
Y[b,:,:] 这两个矩阵中的两两距离,放在常用的计算库(如torch
、paddle
等)中,一样是可以转换为矩阵运算的,只是需要做一点点的reshape而已。
# paddle version
X2 = (X ** 2).sum(axis=2).reshape((batch_size, 1, -1))
Y2 = (Y ** 2).sum(axis=2).reshape((batch_size, -1, 1)) # 由于多了第一维,这里的reshape不能偷懒不写了
D = X2 + Y2 - 2 * X.bmm(Y.transpose((0, 2, 1)))
之所以能这么做,得益于很多库实现了bmm
运算,即对每个batch分别进行矩阵乘法(这个函数应用于张量的时代,在numpy里还没有)。
这一需求其实在pytorch中有一个特定的函数torch.cdist
完成了。然而很不幸的百度paddle框架中并没有类似函数,所以只能自己实现了一版。当然torch内部似乎是用了contiguous函数去做了些矩阵结构的调整(从GitHub issue中看到),具体实现细节倒没有去看。