Question

我有两个numpy数组，X和Y，分别为(n,d)和(m,d)形状。假设我们想要计算X的每一行与Y的每一行之间的欧几里德距离，并将结果存储在数组Z中，形状为(n,m)。我有两个实现。第一个实现使用两个for循环，如下所示：

for i in range(n):
      for j in range(m):
        Z[i,j] = np.sqrt(np.sum(np.square(X[i] - Y[j])))

第二个实现只使用一个矢量化循环：

for i in range(n):
      Z[i] = np.sqrt(np.sum(np.square(X[i]-Y), axis=1))

当我在特定X和Y数据上运行这些代码时，第一个实现需要将近30秒，而第二个实现需要将近60秒。我希望第二个实现更快，因为它使用矢量化。它运行缓慢的原因是什么？我知道我们可以通过完全矢量化代码来获得更快的实现，但我不明白为什么第二个代码（部分矢量化）比非矢量化版本慢。

以下是完整的代码：

n,m,d = 5000,500,3000
X = np.random.rand(n,d)
Y = np.random.rand(m,d)
Z = np.zeros((n,m))

tic = time.time()
for i in range(n):
      for j in range(m):
        Z[i,j] = np.sqrt(np.sum(np.square(X[i] - Y[j])))
print('Elapsed time 1: ', time.time()-tic)

tic = time.time()
for i in range(n):
      Z[i] = np.sqrt(np.sum(np.square(X[i]-Y), axis=1))
print('Elapsed time 2: ', time.time()-tic)


tic = time.time()
train_squared = np.square(X).sum(axis=1).reshape((1,n))
test_squared = np.square(Y).sum(axis=1).reshape((m,1))
test_train = -2*np.matmul(Y, X.T)
dists = np.sqrt(test_train + train_squared + test_squared)
print('Elapsed time 3: ', time.time()-tic)

这是输出：

Elapsed time 1:  35.659096002578735
Elapsed time 2:  65.57051086425781
Elapsed time 3:  0.3912069797515869

Answer 1

我拆开了你的方程并将其缩小到MVCE：

for i in range(n):
    for j in range(m):
        Y[j].copy()

for i in range(n):
    Y.copy()

这里的copy()只是为了模拟X的减法。减法本身应该很便宜。

这是我计算机上的结果：

第一个花了10毫秒。
第二个花了13s！

我正在复制完全相同数量的数据。使用您的选择n=5000, m=500, d=3000，此代码正在复制 60千兆字节的数据。

说实话，我在13秒内并不感到惊讶。它已经超过4GB / s，基本上是我的CPU和RAM之间的最大带宽（例如memcpy）。

真正令人惊讶的是，第一次测试只能在0.01秒内复制60GB，转换为6TB / s！

我很确定这是因为数据实际上并没有离开CPU。它只是在CPU和L1缓存之间来回反复：3000个双精度数字的数组很容易放入32KiB L1缓存中。

因此，我推断出你的第二个算法并不像人们想象的那么大的主要原因是因为每次迭代处理一大块500×3000个元素对CPU缓存非常不友好：你基本上将整个缓存驱逐到RAM中！相比之下，你的第一个算法确实在某种程度上利用了缓存，因为3000元素在计算sum时仍然在缓存中，因此不会像大量数据在CPU和RAM之间移动。（一旦得到总和，3000元素数组将被丢弃＆＃34;这意味着它可能只会被覆盖在缓存中并且永远不会使其回到实际的RAM中。）

当然，做矩阵乘法的速度要快得多，因为你的问题基本上是以下形式：

C[i, j] = ∑[k] f(A[i, k], B[j, k])

如果您将f(x, y)替换为x * y，您可以看到它只是矩阵乘法的一种变体。操作f在这里并不是非常重要 - 重要的是索引在这个等式中的行为，它决定了数组在内存中的存储方式。矩阵乘法算法的本质在于能够通过阻塞来处理这种数组访问，因此原则上整个算法即使对于用户定义的f也不会发生显着变化。不幸的是，在 practice 中，很少有库支持用户定义的操作，因此您可以像使用那样使用技巧(X - Y)**2 = X**2 - 2 X Y + Y**2。但它完成了工作：D

为什么向量化的numpy代码比for循环慢？

1 个答案: