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

我有两个 numpy 数组，X and Y，有形状(n,d) and (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 and 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

我拆解了你的方程并将其简化为这个MVCE https://stackoverflow.com/help/mcve:

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

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

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

这是我电脑上的结果：

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

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

说实话，我对这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) with x * y，你可以看到它只是矩阵乘法的一种变体。操作f这里并不是非常重要 - 重要的是索引在这个方程中的表现，它决定了数组如何存储在内存中。矩阵乘法算法的本质在于能够通过blocking，因此原则上，即使对于用户定义的算法，整体算法也不会发生显着变化f。不幸的是，在practice支持用户定义操作的库很少，所以你必须使用这个技巧(X - Y)**2 = X**2 - 2 X Y + Y**2正如你所做的那样。但它完成了工作：D