Numpy：与两次循环迭代相比，单循环矢量化代码速度较慢

Question

以下代码遍历两个数组的每个元素以计算成对欧氏距离。

def compute_distances_two_loops(X, Y):
    num_test = X.shape[0]
    num_train = Y.shape[0]
    dists = np.zeros((num_test, num_train))
    for i in range(num_test):
        for j in range(num_train):
            dists[i][j] = np.sqrt(np.sum((X[i] - Y[j])**2))
    return dists

以下代码具有相同的目的，但只有一个循环。

def compute_distances_one_loop(X, Y):
    num_test = X.shape[0]
    num_train = Y.shape[0]
    dists = np.zeros((num_test, num_train))
    for i in range(num_test):
        dists[i, :] = np.sqrt(np.sum((Y - X[i])**2, axis=1))
    return dists

以下是两者的时间比较。

two_loop_time = time_function(compute_distances_two_loops, X, Y)
print ('Two loop version took %f seconds' % two_loop_time)

>> Two loop version took 20.3 seconds

one_loop_time = time_function(compute_distances_one_loop, X, Y)
print ('One loop version took %f seconds' % one_loop_time)

>> One loop version took 80.9 seconds

X 和 Y 都是 numpy.ndarray，形状为 -

X: (500, 3000) Y: (5000, 3000)

凭直觉结果不正确，单循环至少应运行具有相同的速度。我在这里错过了什么？

PS：结果不是来自单个运行。我运行代码多次，在不同的机器上，结果都差不多。

Answer 1

原因是循环体内数组的大小。在两个循环中，变体适用于两个包含 3000 个元素的数组。这至少很容易适合 cpu 的二级缓存，它比主内存快得多，但它也足够大，计算距离比 python 循环迭代开销慢得多。

第二种情况循环体作用于 5000 * 3000 个元素。这太多了，以至于在每个计算步骤中数据都需要进入主内存（首先将 Y-X[i] 减法到一个临时数组中，将临时数组平方成另一个临时数组，然后将其读回以求和）。对于所涉及的简单操作，主内存比 cpu 慢得多，因此尽管删除了循环，但它花费的时间要长得多。您可以通过使用写入预分配临时数组的就地操作来加快速度，但对于这些数组大小，它仍然比两个循环变体慢。

请注意 scipy.spatial.distance.cdist(X, Y) 可能是最快的，因为它根本不需要任何临时文件

Numpy：与两次循环迭代相比，单循环矢量化代码速度较慢

Numpy: Single loop vectorized code slow compared to two loop iteration

python

performance

time

numpy

nearest-neighbor