为什么`np.sum(range(N))`运行极慢？Python与Numpy求和性能差异探究
问题描述
我之前看过一个讲解Python循环运行速度的视频，里面提到sum(range(N))比手动遍历range逐次累加的速度快得多，原因是前者调用内置函数在C层运行，而后者的求和操作是在运行效率更低的Python层完成的。我好奇引入numpy后的求和性能表现，测试后发现符合预期的np.sum(np.arange(N))是最快的，但sum(np.arange(N))和np.sum(range(N))的运行速度甚至比朴素for循环还慢。
这种情况出现的原因是什么？
以下是我使用的测试脚本，包含我已知的性能变慢原因注释（大多来自上述视频），以及我在python 3.10.0、numpy 1.21.2环境下得到的运行结果：
更新后的脚本：
import numpy as np from timeit import timeit N = 10_000_000 repetition = 10 def sum0(N = N): s = 0 i = 0 while i < N: # 条件判断在Python层执行 s += i i += 1 # 两次加法操作都在Python层执行 return s def sum1(N = N): s = 0 for i in range(N): # 自增操作在C层执行 s += i # 加法在Python层执行 return s def sum2(N = N): return sum(range(N)) # 所有操作都在C层执行 def sum3(N = N): return sum(list(range(N))) def sum4(N = N): return np.sum(range(N)) # 转np.array的过程非常慢 def sum5(N = N): # 转np.array的过程快很多 return np.sum(np.fromiter(range(N),dtype = int)) def sum5v2_(N = N): # 转np.array的过程快很多 return np.sum(np.fromiter(range(N),dtype = np.int_)) def sum6(N = N): # 从np.int转Py_long的过程可能很慢 return sum(np.arange(N)) def sum7(N = N): # list返回由np.int组成的列表 return sum(list(np.arange(N))) def sum7v2(N = N): # 转python int的操作，用tolist似乎比sum(list())里的隐式转换更快 # tolist返回由python int组成的列表 return sum(np.arange(N).tolist()) def sum8(N = N): return np.sum(np.arange(N)) # 所有操作都在numpy层执行（底层是fortran的libblas？） def sum9(N = N): return np.arange(N).sum() # 省去了分发开销 def array_basic(N = N): return np.array(range(N)) def array_dtype(N = N): return np.array(range(N),dtype = np.int_) def array_iter(N = N): # np.sum的源码提到可以用fromiter来转换生成器 return np.fromiter(range(N),dtype = np.int_) print(f"while loop: {timeit(sum0, number = repetition)}") print(f"for loop: {timeit(sum1, number = repetition)}") print(f"sum_range: {timeit(sum2, number = repetition)}") print(f"sum_rangelist: {timeit(sum3, number = repetition)}") print(f"npsum_range: {timeit(sum4, number = repetition)}") print(f"npsum_iterrange: {timeit(sum5, number = repetition)}") print(f"npsum_iterrangev2: {timeit(sum5, number = repetition)}") print(f"sum_arange: {timeit(sum6, number = repetition)}") print(f"sum_list_arange: {timeit(sum7, number = repetition)}") print(f"sum_arange_tolist: {timeit(sum7v2, number = repetition)}") print(f"npsum_arange: {timeit(sum8, number = repetition)}") print(f"nparangenpsum: {timeit(sum9, number = repetition)}") print(f"array_basic: {timeit(array_basic, number = repetition)}") print(f"array_dtype: {timeit(array_dtype, number = repetition)}") print(f"array_iter: {timeit(array_iter, number = repetition)}") print(f"npsumarangeREP: {timeit(lambda : sum8(N/1000), number = 100000*repetition)}") print(f"npsumarangeREP: {timeit(lambda : sum9(N/1000), number = 100000*repetition)}") # 示例输出： # # while loop: 11.493371912998555 # for loop: 7.385945574002108 # sum_range: 2.4605720699983067 # sum_rangelist: 4.509678105998319 # npsum_range: 11.85120212900074 # npsum_iterrange: 4.464334709002287 # npsum_iterrangev2: 4.498494338993623 # sum_arange: 9.537815956995473 # sum_list_arange: 13.290120724996086 # sum_arange_tolist: 5.231948580003518 # npsum_arange: 0.241889145996538 # nparangenpsum: 0.21876695199898677 # array_basic: 11.736577274998126 # array_dtype: 8.71628468400013 # array_iter: 4.303306431000237 # npsumarangeREP: 21.240833958996518 # npsumarangeREP: 16.690092379001726
原因解答
两种慢的场景核心都是跨类型转换的开销大于底层优化带来的收益，具体拆解如下：
np.sum(range(N))慢的原因np.sum是为numpy数组设计的优化函数，当输入是Python原生的可迭代对象（比如range）时，它会先把整个可迭代对象逐元素转换成numpy数组。默认的np.array(range(N))转换过程需要逐个处理Python int对象，每次都要做类型校验、内存申请，开销远大于纯Python循环的累加，所以整体耗时会超过朴素for循环。你测试里用np.fromiter转换的sum5比sum4快3倍左右，就是因为fromiter是专门为Python生成器/可迭代对象设计的数组转换接口，预分配了内存，减少了中间开销。sum(np.arange(N))慢的原因
Python内置的sum是为Python原生对象设计的，当输入是numpy数组时，它会逐元素遍历数组，每次取出的元素是numpy的int类型，累加前必须先转换成Python原生的int类型，这个类型转换每一步都有Python层的开销，相当于比纯Python for循环多了一步类型转换的成本，所以速度反而更慢。你测试里的sum7v2用tolist()先把整个numpy数组批量转成Python int的列表再求和，速度比sum6快近一倍，就是因为批量转换的开销远低于逐元素隐式转换的开销。
而你测试里最快的np.sum(np.arange(N))，整个流程的数组生成、求和操作都在numpy的C层完成，没有Python层的循环和类型转换开销，底层还用到了向量化、BLAS库优化，所以速度会比纯C实现的sum(range(N))还要快一个数量级。
内容的提问来源于stack exchange，提问作者fbence
