为什么矢量化 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 个双精度数字的数组很容易放入 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