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

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 毫秒。
• 第二个用了 13 秒!

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

老实说，那 13 秒我一点都不惊讶。这已经超过 4GB/s，基本上是我的 CPU 和 RAM 之间的最大带宽(例如 memcpy)。

真正令人惊讶的是，第一个测试成功地在仅 0.01 秒内复制了 60GB，换算为 6TB/s!

我很确定这是因为数据实际上根本没有离开 CPU。它只是在 CPU 和 L1 缓存之间来回跳动:一个 3000 个 double 字的数组很容易放入 32KiB L1 缓存中。

因此，我推断您的第二个算法不如人们天真地期望的那么好，主要原因是因为处理了一大块 500×3000。每次迭代的元素对 CPU 缓存非常不友好:你基本上将整个缓存驱逐到 RAM 中!相比之下，您的第一个算法确实在某种程度上利用了缓存，因为 3000到 sum 时，元素仍将在缓存中得到计算，因此在 CPU 和 RAM 之间移动的数据几乎没有那么多。 (一旦你得到总和，3000 元素数组就会被“丢弃”，这意味着它可能只会在缓存中被覆盖而永远不会回到实际的 RAM 中。)

---

自然地，进行矩阵乘法要快得多，因为您的问题本质上是以下形式:

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

如果你替换f(x, y)与 x * y ，你可以看到它只是矩阵乘法的一个变体。操作f在这里并不是特别重要——重要的是索引在这个方程中的表现，它决定了数组在内存中的存储方式。矩阵乘法算法的本质在于能够通过阻塞来应对这种数组访问，所以原则上即使对于用户定义的f，整体算法也不会发生显着变化。 .不幸的是，在实践中，支持用户定义操作的库非常少，因此您使用了技巧 (X - Y)**2 = X**2 - 2 X Y + Y**2正如你所做的那样。但它完成了工作 :D