使用Numpy广播以向量化计算欧几里得距离。

16

我有大小为2x4和3x4的矩阵。我想要计算行之间的欧几里得距离,并得到一个大小为2x3的矩阵。这是一个使用for循环计算a矩阵中每一行向量与b矩阵中所有行向量欧几里得距离的代码。如何在不使用for循环的情况下完成相同的操作?

 import numpy as np
a = np.array([[1,1,1,1],[2,2,2,2]])
b = np.array([[1,2,3,4],[1,1,1,1],[1,2,1,9]])
dists = np.zeros((2, 3))
for i in range(2):
      dists[i] = np.sqrt(np.sum(np.square(a[i] - b), axis=1))
5个回答

43

这里是原始输入变量:

A = np.array([[1,1,1,1],[2,2,2,2]])
B = np.array([[1,2,3,4],[1,1,1,1],[1,2,1,9]])
A
# array([[1, 1, 1, 1],
#        [2, 2, 2, 2]])
B
# array([[1, 2, 3, 4],
#        [1, 1, 1, 1],
#        [1, 2, 1, 9]])

A是一个2x4的数组。 B是一个3x4的数组。

我们想要通过一个完全向量化的操作来计算欧几里得距离矩阵运算,其中dist[i,j]包含A中第i个实例和B中第j个实例之间的距离。因此,在这个例子中,dist是2x3的。

距离

enter image description here

可以用numpy来表述

dist = np.sqrt(np.sum(np.square(A-B))) # DOES NOT WORK
# Traceback (most recent call last):
#   File "<stdin>", line 1, in <module>
# ValueError: operands could not be broadcast together with shapes (2,4) (3,4)

然而,如上所示,问题在于逐元素减法操作 A-B 涉及不兼容的数组大小,特别是第一维中的 2 和 3。
A has dimensions 2 x 4
B has dimensions 3 x 4

为了进行按元素减法,我们必须填充A或B中的一个以满足numpy的广播规则。我选择用额外的一维来填充A,使其变成2 x 1 x 4,这样可以使数组的维度对齐以进行广播。有关numpy广播的更多信息,请参见scipy手册中的教程本教程的最后一个示例。
您可以使用np.newaxis值或np.reshape命令执行填充。我两种方法都会展示:
# First approach is to add the extra dimension to A with np.newaxis
A[:,np.newaxis,:] has dimensions 2 x 1 x 4
B has dimensions                     3 x 4

# Second approach is to reshape A with np.reshape
np.reshape(A, (2,1,4)) has dimensions 2 x 1 x 4
B has dimensions                          3 x 4

如您所见,使用任何一种方法都可以使维度对齐。我将使用第一种方法,即 np.newaxis。因此,现在可以使用以下代码创建一个 2x3x4 的数组 A-B:

diff = A[:,np.newaxis,:] - B
# Alternative approach:
# diff = np.reshape(A, (2,1,4)) - B
diff.shape
# (2, 3, 4)

现在我们可以将这个差异表达式放入“dist”方程语句中,以获得最终结果:
dist = np.sqrt(np.sum(np.square(A[:,np.newaxis,:] - B), axis=2))
dist
# array([[ 3.74165739,  0.        ,  8.06225775],
#        [ 2.44948974,  2.        ,  7.14142843]])

请注意,sum 是在 axis=2 上进行的,这意味着对 2x3x4 数组的第三个轴求和(其中轴 id 从 0 开始)。
如果您的数组很小,则上述命令将正常工作。但是,如果您有大型数组,则可能会遇到内存问题。请注意,在上面的示例中,numpy 在内部创建了一个 2x3x4 数组来执行广播。如果我们将 A 的维度通用化为 a x z,将 B 的维度通用化为 b x z,则 numpy 将在广播时在内部创建一个 a x b x z 数组。
我们可以通过进行一些数学操作来避免创建此中间数组。因为您正在计算欧几里得距离作为平方差的总和,所以我们可以利用平方差可重写的数学事实。

enter image description here

注意中间项涉及按元素相乘求和。这个乘法求和更为常见的称呼是点乘。因为A和B都是矩阵,所以这个操作实际上是矩阵乘法。我们可以将上述式子重写为:

enter image description here

我们可以编写以下NumPy代码:

threeSums = np.sum(np.square(A)[:,np.newaxis,:], axis=2) - 2 * A.dot(B.T) + np.sum(np.square(B), axis=1)
dist = np.sqrt(threeSums)
dist
# array([[ 3.74165739,  0.        ,  8.06225775],
#        [ 2.44948974,  2.        ,  7.14142843]])

请注意,上面的答案与之前的实现完全相同。这里的优点是我们不需要为广播创建中间的2x3x4数组。
为了完整起见,让我们再次检查threeSums中每个加数的维度是否允许广播。
np.sum(np.square(A)[:,np.newaxis,:], axis=2) has dimensions 2 x 1
2 * A.dot(B.T) has dimensions                               2 x 3
np.sum(np.square(B), axis=1) has dimensions                 1 x 3

因此,正如预期的那样,最终的dist数组具有2x3的尺寸。

使用点积代替逐元素乘法之和的方法也在this tutorial中讨论过。


4
这个答案非常有用,尤其是克服广播问题的部分。谢谢@stackoverflowuser2010。 - Allen Qin
1
很棒的答案!但是我有一个问题,因为似乎你仍然需要在求和中广播第一个和最后一个数组。这仍然是可取的吗? - penny

25

最近我在做深度学习(斯坦福cs231n,第一次作业)时遇到了同样的问题,但当我使用

 np.sqrt((np.square(a[:,np.newaxis]-b).sum(axis=2)))

有一个错误发生了

MemoryError

这意味着我的内存不足了(事实上,在中间创建了一个500*5000*1024的数组。它太大了!)

为了防止这个错误,我们可以使用一个公式进行简化:

代码:

import numpy as np
aSumSquare = np.sum(np.square(a),axis=1);
bSumSquare = np.sum(np.square(b),axis=1);
mul = np.dot(a,b.T);
dists = np.sqrt(aSumSquare[:,np.newaxis]+bSumSquare-2*mul)

2
添加某些内容;引用自官方文档。然而,有些情况下广播是不好的,因为它会导致内存使用效率低下,从而减慢计算速度。 - Han Qiu
我正在处理同一个问题,但使用这个无循环实现得到的结果与我使用一次和两次循环解决方案得到的结果不匹配。 - The Governor

20

只需在正确的位置使用np.newaxis

 np.sqrt((np.square(a[:,np.newaxis]-b).sum(axis=2)))

11
能否解释一下Simply using np.newaxis at the right place是如何起作用的?如果您能从a是2x4,b是3x4这个事实开始讲起,那就太好了。 - stackoverflowuser2010

4

这个功能已经包含在scipy的空间模块中,我建议使用它,因为它将在底层进行向量化和高度优化。但是,正如其他答案所示,你也可以自己实现。

import numpy as np
a = np.array([[1,1,1,1],[2,2,2,2]])
b = np.array([[1,2,3,4],[1,1,1,1],[1,2,1,9]])
np.sqrt((np.square(a[:,np.newaxis]-b).sum(axis=2)))
# array([[ 3.74165739,  0.        ,  8.06225775],
#       [ 2.44948974,  2.        ,  7.14142843]])
from scipy.spatial.distance import cdist
cdist(a,b)
# array([[ 3.74165739,  0.        ,  8.06225775],
#       [ 2.44948974,  2.        ,  7.14142843]])

2
使用 numpy.linalg.norm 函数并结合广播功能可以很好地实现矢量范数计算。通过指定整数值作为 axis 参数,可以使用向量范数,默认为欧几里得范数。请注意保留 HTML 标签。
import numpy as np

a = np.array([[1,1,1,1],[2,2,2,2]])
b = np.array([[1,2,3,4],[1,1,1,1],[1,2,1,9]])
np.linalg.norm(a[:, np.newaxis] - b, axis = 2)

# array([[ 3.74165739,  0.        ,  8.06225775],
#       [ 2.44948974,  2.        ,  7.14142843]])

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接