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

12

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

考虑下列代码。三个简单函数用于计算NumPy 3D数组(1000 x 1000 x 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还是其他原因?


你正在声明一个8Gb的数组,这在许多台式机或笔记本电脑上都无法容纳,因此你隐含地测试了缓存行为。类型为np.int64的3D numpy数组(1000×1000×1000立方体)有10^9个元素,每个元素占用8个字节。因此,这将归结为一个8Gb的整数数组是否适合/被缓存等在你特定的机器上。由于它无法适应许多台式机或笔记本电脑,因此如果您不想测试缓存,请改用尺寸为100×100×100的数组,这将是8Mb。 - smci
此外,为了使其确定性,需要在 np.random.randint 调用之前设置随机种子。 - smci
@Barker,你似乎忽略了秒和毫秒之间的区别。 - plugwash
3个回答

7
这是由于内存访问和缓存所致。每个函数都在做两件事情,以第一个代码示例为例:
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 结合使用。

如果您发现自己经常手动编写 for 循环以获得更好的性能,您可能需要查看 numexpr - 它会自动执行其中的一些操作。我自己没有用过它,但如果中间数组是减慢程序速度的原因,它应该提供一个很好的加速。


很好的发现。确实通过删除> 0测试,numpy是最快的。 - BlackBear
谢谢您的回复。您能解释一下您所说的缓存是什么吗?我问这个问题是因为data大约有3.7GB,data > 0大约有0.9GB,data[0, :, :]大约有3.8MB。 - user10325516
@Poolka,这是由CPU缓存引起的。我不知道具体的数字,但通常它有几MB的内存(这里是我谷歌搜索到的一个页面,其中有一些更好的数字:https://www.makeuseof.com/tag/what-is-cpu-cache/)。 - user2699
@Poolka,Numpy在这里处理内存布局方面做得很正确,问题中给出的每个选项都以相同的方式访问内存,因此这不是问题所在。 - user2699
关于 f_3:实际上在这里不会创建大型中间数组 - 当 axis=None 时(如 np.count_nonzero(arr) 的情况),你突出显示的 Python 代码片段不会被执行。编译后的代码只是遍历数组并计算值的数量(虽然它必须首先设置 nditer,但开销相对较小)。 - Alex Riley
显示剩余2条评论

4

一切都取决于数据在内存中的布局以及代码如何访问它。实质上,数据是以块的形式从内存中获取的,然后被缓存;如果算法成功地使用了缓存中的块中的数据,则不需要再次读取内存。这可以节省大量时间,尤其是当缓存比您处理的数据小得多时。

考虑以下变体,它们只在我们迭代的轴上有所不同:

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所使用的那种)。有关缓存如何影响计时的详细解释,请参见其他答案

谢谢你的回答。我看到你编辑了答案并参考了另一个答案。然而,你注意到的那个对数组轴的强依赖性,是否可以称之为NumPy的问题? - user10325516
@Poolka 这是一个numpy问题,因为data > 0创建了一个_copy_,而不是_view_(与索引相反)。虽然不是迭代的问题:如果您运行相同的测试并删除> 0部分,则会发现numpy最快(尽管不是很快,我的结果是500ms、537ms、945ms和10.8s)。 - BlackBear

2

完全删除临时数组

正如@user2699在他的回答中所提到的,分配和写入一个不适合缓存的大型数组可能会大大减慢进程。为了展示这种行为,我编写了两个使用Numba(JIT-编译器)的小函数。

在编译语言(C、Fortran等)中,通常避免使用临时数组。在解释性Python中(未使用Cython或Numba),您经常希望在更大的数据块上调用编译后的函数(向量化),因为解释性代码中的循环非常缓慢。但是这也可能有一些缺点(例如临时数组、糟糕的缓存使用)

没有临时数组分配的函数

@nb.njit(fastmath=True,parallel=False)
def f_4(arr):
    sum=0
    for i in nb.prange(arr.shape[0]):
        for j in range(arr.shape[1]):
            for k in range(arr.shape[2]):
                if arr[i,j,k]>0:
                    sum+=1
    return sum

使用临时数组

请注意,如果您打开并行化parallel=True,编译器不仅会尝试并行化代码,还会打开其他优化,如循环融合。

最初的回答:

请注意,如果您将parallel=True设置为真,编译器不仅会尝试并行化代码,还会打开其他优化,如循环融合。

@nb.njit(fastmath=True,parallel=False)
def f_5(arr):
    return np.sum(arr>0)

最初的回答
时间安排
%timeit f_1(data)
1.65 s ± 48.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit f_2(data)
1.27 s ± 5.66 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit f_3(data)
1.99 s ± 6.11 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit f_4(data) #parallel=false
216 ms ± 5.45 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit f_4(data) #parallel=true
121 ms ± 4.85 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit f_5(data) #parallel=False
1.12 s ± 19 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit f_5(data) #parallel=true Temp-Array is automatically optimized away
146 ms ± 12.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接