为什么NumPy有时会比NumPy +纯Python循环慢？

Question

这是基于2018-10年度提出的这个问题。

考虑下面的代码。三种简单的函数可以计算NumPy 3D数组（1000×1000×1000）中的非零元素。

import numpy as np

def f_1(arr):
    return np.sum(arr > 0)

def f_2(arr):
    ans = 0
    for val in range(arr.shape[0]):
        ans += np.sum(arr[val, :, :] > 0)
    return ans

def f_3(arr):
    return np.count_nonzero(arr)

if __name__ == '__main__':

    data = np.random.randint(0, 10, (1_000, 1_000, 1_000))
    print(f_1(data))
    print(f_2(data))
    print(f_3(data))

我的机器上的执行阶段（Python 3.7.?、Windows 10，NumPy 1.16。？）：

%timeit f_1(data)
1.73 s ± 21.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit f_2(data)
1.4 s ± 1.36 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit f_3(data)
2.38 s ± 956 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

因此，f_2()比f_1()和更快地工作f_3()。但是，data较小的情况并非如此。问题是-为什么呢？是NumPy，Python还是其他？

Answer 1

这是由于内存访问和缓存。这些功能每个都在做两件事，以第一个代码为例：

np.sum(arr > 0)

它首先进行比较以查找arr大于零（或非零，因为arr包含非负整数）的位置。这将创建一个与形状相同的中间数组arr。然后，将这个数组求和。

坦白吧？好吧，当使用np.sum(arr > 0)这个数组时。当它足够大而不能容纳在高速缓存中时，性能将降低，因为当处理器开始执行时，大多数数组元素的总和将从内存中逐出并需要重新加载。

由于f_2遍历第一维，因此它正在处理较小的子数组。完成相同的复制和求和，但是这次中间数组适合内存。它是在不离开内存的情况下创建，使用和销毁的。这要快得多。

现在，您会认为这f_3将是最快的（使用内置方法和全部方法），但是查看源代码表明它使用以下操作：

a_bool = a.astype(np.bool_, copy=False)
return a_bool.sum(axis=axis, dtype=np.intp

a_bool 这只是查找非零条目并创建大型中间数组的另一种方法。

结论

经验法则就是这样，并且经常是错误的。如果您想要更快的代码，请对其进行概要分析，然后查看问题出在哪里（在此处进行出色的工作）。

Python做得很好。在进行优化的情况下，它可能比快numpy。不要害怕将普通的旧python代码或数据类型与numpy结合使用。

如果您经常发现自己手动编写循环以提高性能，则可能需要看看numexpr-它会自动执行其中的一些操作。我自己并没有使用太多，但是如果中间数组使您的程序变慢，它应该可以提供很好的加速。

Answer 2

这完全取决于数据在内存中的布局方式以及代码如何访问它。本质上，数据是从内存中以块的形式获取的，然后将其缓存。如果算法设法使用缓存中的块中的数据，则无需再次从内存中读取数据。这可以节省大量时间，尤其是当缓存远小于您要处理的数据时。

考虑这些变化，这些变化仅在我们要迭代的轴上有所不同：

def f_2_0(arr):
    ans = 0
    for val in range(arr.shape[0]):
        ans += np.sum(arr[val, :, :] > 0)
    return ans

def f_2_1(arr):
    ans = 0
    for val in range(arr.shape[1]):
        ans += np.sum(arr[:, val, :] > 0)
    return ans

def f_2_2(arr):
    ans = 0
    for val in range(arr.shape[2]):
        ans += np.sum(arr[:, :, val] > 0)
    return ans

结果在我的笔记本电脑上：

%timeit f_1(data)
2.31 s ± 47.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit f_2_0(data)
1.88 s ± 60 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit f_2_1(data)
2.65 s ± 142 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit f_2_2(data)
12.8 s ± 650 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

您f_2_1几乎可以看到f_1，这使我认为numpy没有使用最佳访问模式（所使用的模式f_2_0）。关于缓存如何精确影响时序的说明在另一个答案中。