为什么 `np.sum(range(N))` 非常慢？

Question

我看到了一个video关于python中循环的速度，其中解释说执行sum(range(N))比手动循环通过range并将变量添加在一起要快得多，因为前者在C中运行由于内置-in 正在使用的函数，而在后者中，求和是在（慢）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: # condition is checked in python
        s += i
        i += 1 # both additions are done in python
    return s

def sum1(N = N):
    s = 0
    for i in range(N): # increment in C
        s += i # addition in python
    return s

def sum2(N = N):
    return sum(range(N)) # everything in C

def sum3(N = N):
    return sum(list(range(N)))

def sum4(N = N):
    return np.sum(range(N)) # very slow np.array conversion

def sum5(N = N):
    # much faster np.array conversion
    return np.sum(np.fromiter(range(N),dtype = int))

def sum5v2_(N = N):
    # much faster np.array conversion
    return np.sum(np.fromiter(range(N),dtype = np.int_))

def sum6(N = N):
    # possibly slow conversion to Py_long from np.int
    return sum(np.arange(N))

def sum7(N = N):
    # list returns a list of np.int-s
    return sum(list(np.arange(N)))

def sum7v2(N = N):
    # tolist conversion to python int seems faster than the implicit conversion
    # in sum(list()) (tolist returns a list of python int-s)
    return sum(np.arange(N).tolist())

def sum8(N = N):
    return np.sum(np.arange(N)) # everything in numpy (fortran libblas?)

def sum9(N = N):
    return np.arange(N).sum() # remove dispatch overhead

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's source code mentions to use fromiter to convert from generators
    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)}")

# Example output:
#
# 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

Answer 1

np.sum(range(N)) 很慢，主要是因为当前的 Numpy 实现没有使用关于生成器 range(N) 提供的值的确切类型/内容的足够信息。一般问题的核心本质上是由于 Python 和大整数的动态类型，尽管 Numpy 可以优化这种特定情况。

首先，range(N) 返回一个动态类型的 Python 对象，它是一种（特殊类型的）Python 生成器。此生成器提供的对象也是动态类型的。它实际上是一个纯 Python 整数。

问题是 Numpy 是用静态类型语言 C 编写的，因此它不能有效地处理动态类型的纯 Python 对象。 Numpy 的策略是尽可能将此类对象转换为 C 类型。在这种情况下，一个大问题是生成器提供的整数理论上可能很大：Numpy 不知道这些值是否会溢出np.int32 甚至np.int64 类型。因此，Numpy 首先检测要使用的好类型，然后使用该类型计算结果。

这个翻译过程可能非常昂贵，这里似乎不需要，因为所有值都由range(10_000_000) 提供。但是，range(5_000_000_000) 返回相同的对象类型，纯 Python 整数溢出 np.int32 和 Numpy 需要自动检测这种情况才不会返回错误结果。问题也是可以正确识别输入类型（我的机器上的np.int32），这并不意味着输出结果是正确的，因为在计算总和期间可能会出现溢出。可悲的是，我的机器就是这种情况。

Numpy 开发人员决定弃用这种用法，并在文档中说明应该使用 np.fromiter。 np.fromiter 有一个dtype 必需参数，让用户定义要使用的好类型。

在实践中检查这种行为的一种方法是简单地使用创建一个临时列表：

tmp = list(range(10_000_000))

# Numpy implicitly convert the list in a Numpy array but 
# still automatically detect the input type to use
np.sum(tmp)

更快的实现如下：

tmp = list(range(10_000_000))

# The array is explicitly converted using a well-defined type and 
# thus there is no need to perform an automatic detection 
# (note that the result is still wrong since it does not fit in a np.int32)
tmp2 = np.array(tmp, dtype=np.int32)
result = np.sum(tmp2)

第一种情况在我的机器上需要 476 毫秒，而第二种情况需要 289 毫秒。请注意，np.sum 只需要 4 毫秒。因此，大部分时间都花在了将纯 Python 整数对象转换为内部 int32 类型（更具体地说，是管理纯 Python 整数）。 list(range(10_000_000)) 也很昂贵，因为它需要 205 毫秒。这又是由于纯 Python 整数的开销（即分配、释放、引用计数、可变大小整数的增量、内存间接和动态类型引起的条件）以及生成器的开销。

sum(np.arange(N)) 很慢，因为sum 是一个处理 Numpy 定义的对象的纯 Python 函数。 CPython 解释器需要调用 Numpy 函数来执行基本的添加操作。此外，Numpy 定义的整数对象仍然是 Python 对象，因此它们会受到引用计数、分配、释放等的影响。更不用说 Numpy 和 CPython 在函数中添加了许多检查，目的是最终将两个原生数字相加。 Numba 等支持 Numpy 的即时编译器可以解决此问题。实际上，我的机器上的 Numba 需要 23 毫秒来计算 np.arange(10_000_000) 的总和（代码仍然用 Python 编写），而 CPython 解释器需要 556 毫秒。

Answer 2

看看我能不能总结一下结果。

sum 可以与任何可迭代对象一起工作，反复询问下一个值并添加它。 range 是一个生成器，它很乐意提供下一个值

# sum_range:          1.4830789409988938

从一个范围内制作一个列表需要时间：

# sum_rangelist:      3.6745876889999636

对预先生成的列表求和实际上比对范围求和更快：

%%timeit x = list(range(N))
    ...: sum(x)

np.sum 旨在对数组求和。它是np.add.reduce 的包装器。

np.sum对np.sum(generator)有弃用警告，推荐使用fromiter或Pythonsum：

# npsum_range:        16.216972655000063

fromiter 是从生成器创建数组的最佳方式。在range 上使用np.array 是遗留代码，将来可能会消失。我认为这是np.array 将接受的唯一generator。

np.array 是一个通用函数，可以处理许多情况，包括嵌套数组和转换为各种 dtype。因此，它必须处理整个输入参数，同时推断形状和 dtype。

# npsum_fromiterrange:3.47655400199983

numpy 数组的迭代比列表慢，因为它必须“拆箱”每个元素。

# sum_arange:         16.656015603000924

同样，从数组中创建列表也很慢；同类型的python级迭代。

# sum_list_arange:    19.500842117000502

arr.tolist()比较快，在编译后的代码中创建一个纯python列表。所以速度类似于从范围列列表。

# sum_arange_tolist:  4.004777374000696

数组的np.sum 是纯numpy 并且相当快。 np.sum(x) x=np.arange(N) 甚至更快（大约 4 倍）

# npsum_arange:       0.2332638230000157

np.sum from range or list 主要是先创建数组的成本：

# array_basic:        16.1631146109994
# array_dtype:        16.550737804000164
# array_iter:         3.9803170430004684

Answer 3

从cpython source code for sum sum 最初似乎尝试了一条假设所有输入都是相同类型的快速路径。如果失败，它将只是迭代：

/* Fast addition by keeping temporary sums in C instead of new Python objects.
   Assumes all inputs are the same type.  If the assumption fails, default
   to the more general routine.
*/

我不完全确定引擎盖下发生了什么，但很可能是 C 类型重复创建/转换为 Python 对象导致了这些减速。值得注意的是sum和range都是用C实现的。

---

下一点并不是问题的真正答案，但我想知道我们是否可以为 python ranges 加速 sum，因为 range 相当 a smart object。

为此，我使用functools.singledispatch 覆盖了专门针对range 类型的built-in sum 函数；然后实现了一个小函数来计算sum of an arithmetic progression。

from functools import singledispatch

def sum_range(range_, /, start=0):
    """Overloaded `sum` for range, compute arithmetic sum"""
    n = len(range_)
    if not n:
        return start
    return int(start + (n * (range_[0] + range_[-1]) / 2))

sum = singledispatch(sum)
sum.register(range, sum_range)

def test():
    """
    >>> sum(range(0, 100))
    4950
    >>> sum(range(0, 10, 2))
    20
    >>> sum(range(0, 9, 2))
    20
    >>> sum(range(0, -10, -1))
    -45
    >>> sum(range(-10, 10))
    -10
    >>> sum(range(-1, -100, -2))
    -2500
    >>> sum(range(0, 10, 100))
    0
    >>> sum(range(0, 0))
    0
    >>> sum(range(0, 100), 50)
    5000
    >>> sum(range(0, 0), 10)
    10
    """

if __name__ == "__main__":
    import doctest
    doctest.testmod()

我不确定这是否完整，但它肯定比循环更快。