Python 性能优化：不再随意

2023年8月1日

我的上一篇文章（老实说，主要是为了测试这个网站的主题）提供了一些代码片段，用于计算 $N$ 项逆平方和。我用我最喜欢的四种语言——Python、C、Rust 和 Haskell——编写了这些代码，但当我运行 Python 代码时，它的速度慢得令人尴尬。与 Rust 顺序执行所需的 $\approx 950$ 毫秒相比，Python 竟然花了 70 秒！因此，在这篇文章中，我们将尝试让 Python 的表现更加合理。

原始解决方案

def basel(N: int) -> float:
    return sum(x**(-2) for x in range(1,N))

这是最 Pythonic 的实现方式。一行易于阅读的代码，使用了内置的生成器。让我们用 $N = 1 0^{8}$ 来计时（而不是像上一篇文章那样用 $1 0^{9}$ ，因为我不想等那么久）：

import time

# 计时函数
def time_function(func):
    def wrapper(*args, **kwargs):
        start_time = time.perf_counter()
        result = func(*args, **kwargs)
        end_time = time.perf_counter()
        execution_time = end_time - start_time
        print(f"函数 '{func.__name__}' 执行时间为 {execution_time:.6f} 秒。")
        return result
    return wrapper

>>> f = time_function(basel)
>>> f(100000000) # 10^8
函数 'basel' 执行时间为 6.641589 秒。
1.644934057834575

让我们尝试用 for 循环重写它，虽然不那么 Pythonic：

def basel_less_pythonic(N: int) -> float:
    s = 0.0
    for x in range(1, N):
        s += x**(-2)
    return s

>>> f(100000000) # 10^8
函数 'basel_less_pythonic' 执行时间为 5.908466 秒。
1.644934057834575

有趣的是，不那么符合惯用语的写法实际上提升了性能。为什么？让我们使用 dis，Python 的反汇编模块来研究一下。

这是原始版本的反汇编代码，我添加了一些有用的注释：

# 加载变量
00 LOAD_GLOBAL              0 (sum)
02 LOAD_CONST               1 (<code object <genexpr> at 0x104f42e40>)
04 LOAD_CONST               2 ('basel.<locals>.<genexpr>')

# 创建函数
06 MAKE_FUNCTION            0

# 创建 `range` 对象
08 LOAD_GLOBAL              1 (range)
10 LOAD_CONST               3 (1)
12 LOAD_FAST                0 (N)
14 CALL_FUNCTION            2

# 转换为迭代器
16 GET_ITER

# 调用生成器
18 CALL_FUNCTION            1

# 调用 sum 函数
20 CALL_FUNCTION            1

# 返回
22 RETURN_VALUE

f <code object <genexpr> at 0x104f42e40>:
# 本质上是一个带有 yield 的 `for` 循环
00 GEN_START                0
02 LOAD_FAST                0 (.0)
04 FOR_ITER                 7 (to 20)
06 STORE_FAST               1 (x)
08 LOAD_FAST                1 (x)
10 LOAD_CONST               0 (-2)
12 BINARY_POWER
14 YIELD_VALUE
16 POP_TOP
18 JUMP_ABSOLUTE            2 (to 4)
20 LOAD_CONST               1 (None)
22 RETURN_VALUE

这是 for 循环版本的反汇编代码：

# 加载变量
00 LOAD_CONST               1 (0.0)
02 STORE_FAST               1 (s)

# 创建 `range` 对象
04 LOAD_GLOBAL              0 (range)
06 LOAD_CONST               2 (1)
08 LOAD_FAST                0 (N)
10 CALL_FUNCTION            2
12 GET_ITER

# 声明 `for` 循环
14 FOR_ITER                 8 (to 32)
16 STORE_FAST               2 (x)

# 幂运算、加法、存储
18 LOAD_FAST                1 (s)
20 LOAD_FAST                2 (x)
22 LOAD_CONST               3 (-2)
24 BINARY_POWER
26 INPLACE_ADD
28 STORE_FAST               1 (s)

# 跳回 `for` 循环顶部
30 JUMP_ABSOLUTE            7 (to 14)

# 返回 `s`
32 LOAD_FAST                1 (s)
34 RETURN_VALUE

在这个 特定情况下，使用生成器反而拖慢了程序。正如我们所看到的，生成器循环的代码与第二个版本完全相同。第一个版本只是额外处理了生成器对象，这需要（较慢的）CALL_FUNCTION 指令。需要注意的是，在实践中，生成器通常由于其惰性特性而更快。只是在这个案例中，额外的开销并不值得。

无论如何，两个版本之间的性能差异（约 1 秒）与 Python 和 C/Rust 之间的整体差异相比是微不足道的。即使使用更快的版本，Rust 代码也比 Python 快约 65 倍。

为什么？主要是因为 Python 是一种松散类型的解释型语言。这意味着 Python 解释器（CPython）必须将 Python 代码即时翻译成计算机易于执行的形式。它通过将 Python 代码快速编译成一种通用的汇编形式（称为 Python 字节码）来实现这一点，我们刚刚已经看到了。

然后，解释器会执行这些字节码。这个编译步骤是 Python 性能受损的第一个原因。由于像 Rust 这样的语言只需要编译一次程序，因此这个时间不会计入程序的 运行时。但性能受损的最大原因是字节码的质量较差，它是与架构无关的，而不是原生优化的。这只是解释型语言本质的一部分，因为它们无法花费太多时间进行高质量的编译。

那么，如何编写快速的 Python 代码呢？其实你做不到。但是，你可以通过调用高度优化的库（如 Numpy）来调用快速的 C 代码。这些库包含预编译的、向量化的 C 函数，让你可以绕过整个 Python 解释器。

让我们使用 Numpy

import numpy as np

def basel_np(N: int) -> float:
    # [1, 1, ..., 1]
    ones = np.ones(N - 1)
    # [1, 2, ..., N]
    r = np.arange(1, N)
    # [1, 1/2, ..., 1/N]
    div = ones / r
    # [1, 1/4, ..., 1/N^2]
    inv_squares = np.square(div)
    # ~ pi^2/6
    return float(np.sum(inv_squares))

>>> f(100000000) # 10^8
Function 'basel_np' executed in 0.460317 seconds.
1.6449340568482196

哇，这次运行速度提升了约 13 倍！在这种情况下查看字节码不会告诉我们太多信息，因为它只会包含一些对 numpy 的 CALL_FUNCTION 调用，而 numpy 才是实际执行工作的部分。让我们看看哪一行代码耗时最多：

def basel_np(N: int) -> tuple[float, list[float]]:
    times = []

    # 计时单个步骤
    start = time.perf_counter()
    ones = np.ones(N - 1)
    end = time.perf_counter()
    step_time = end - start
    times.append(step_time)

    # 剩余的计时代码省略
    r = np.arange(1, N)
    div = ones / r
    square = np.square(div)
    ret = np.sum(square)

    return ret, times

操作	操作耗时 (ms)	累计耗时 (ms)
创建 ones	97	97
创建 range	79	176
平方 range	98	274
除法 ones/squares	112	387
最终求和	58	444

让我们思考一下这些数字。创建 ones/range 的步骤并不涉及繁重的计算工作。然而，它们花费的时间几乎与“更复杂”的步骤（如平方和除法）相同。令人惊讶的是，对最终数组求和是最快的步骤！这表明这里的瓶颈不是 CPU，而是内存访问。让我们看看性能如何随着 $N$ 从 $1 0^{7}$ 到 $1 0^{9}$ 的变化而变化。

在这张图中，最上面线的边缘代表总耗时，而连续线之间的空间代表每个步骤的耗时。

我们可以看到：

总时间呈非线性增长，尽管算法是 $O (n)$ 的。
divide 步骤耗时最多，在 $\approx 3 \times 1 0^{8}$ 和 $\approx 7 \times 1 0^{8}$ 时出现急剧增加。

这是有道理的，因为与其他操作不同，np.divide 需要同时访问两个大数组，并将结果写入一个新数组。这意味着大量的主内存访问，甚至可能从 SSD 读取数据，而 SSD 的速度非常慢。

广播

实际上，Numpy 为这类问题内置了一种优化，称为广播，它可以在向量-向量操作中，将较小尺寸的向量“投影”到较大尺寸。

def basel_np_broadcast(N) -> float:
    ones = 1
    r = np.arange(1, N)
    div = ones / r
    square = np.square(div)
    # 就像 [1, 1, ..., 1] / [1, 4, ..., N^2]
    return float(np.sum(square))

这让我们节省了大量宝贵的缓存空间。让我们运行与之前相同的指标。当 $N = 1 0^{8}$ 时：

操作	操作时间 (ms)	累计时间 (ms)
创建 ones	0.00	0
创建 range	68.56	70.48
对 range 进行平方	105.14	180.74
进行 ones/squares 除法	133.08	271.30
最终求和	71.08	310.41

从现在开始，我将把广播版本称为“Numpy 解决方案”。

尽管有了显著的改进，但我们仍然看到那些延迟峰值。让我们尝试解决这个问题。

关于内存的思考

为了改进函数的内存使用，我们可以逐块处理，每次对范围内的小块数据进行相同的操作，最后将所有结果相加。这样，我们一次只需在内存中保留一个数据块，有望在仍受益于 Numpy 向量化代码的同时，提高缓存利用率。

让我们修改函数以处理一个 $N$ 的范围：

# [N1, N2]，包含两端
def basel_np_range(N1: int, N2: int) -> float:
    # 计时代码省略
    ones = 1
    r = np.arange(N1, N2 + 1)
    div = ones / r
    squares = np.square(div)
    return float(np.sum(squares))

并编写一个函数来汇总所有块的结果。

def basel_chunks(N: int, chunk_size: int) -> float:
    # 计时代码省略
    s = 0.0
    num_chunks = N // chunk_size
    for i in range(num_chunks):
        s += basel_np_range(i * chunk_size + 1, (i + 1) * chunk_size)
    return s

当 $N = 1 0^{8}$ 时：

函数 'basel_chunks' 在 0.108557 秒内执行完毕。
1.6449340568482258

不错！它的运行速度比 Numpy 解决方案快了约 3 倍。

此外，由于 IEEE 754 浮点数的特性，答案实际上会稍微更准确。让我们像之前一样看看性能如何扩展，保持块大小不变。

首先，看看 y 轴。我们现在的工作单位是秒，而不是分钟。我们还看到运行时间线性增加，与算法的 $O (n)$ 时间复杂度一致。我们可以推测，当 chunk_size 为 20000 时，所有数据都能放入缓存中，因此不会出现因缓存未命中而导致的运行时峰值。

让我们尝试在 $[5 \times 1 0^{2}, 1 0^{6}]$ 范围内改变 chunk_size，同时保持 $N$ 为 $1 0^{9}$ 不变。

从图中可以看出，性能提升一直持续到 chunk_size 约为 51000，之后由于缓存未命中，出现了延迟峰值。

一些粗略计算

每个 float64 占用 64 位，即 8 字节。我们处理的是 3 个长度为 $N$ 的 float64 数组，因此一次调用的内存使用量为 $51000 (8) (3) = 1224000 \approx 1.2$ MB。

我的 M1 Pro 具有以下规格：

内存类型	大小
L1	128KB（数据缓存，每核心）
L2	24MB（6 个性能核心共享）
L3	24MB
主内存	16GB

这表明函数可用的最大 $L 1 + L 2$ 缓存空间约为 $\approx 1.2$ MB，如果超过这个值，我们就需要等待 $L 3$ 缓存。

多进程处理

现在，让我们尝试让计算机将更多的数组放入缓存中。一种可能性是尝试在所有核心上运行函数，这样我们就可以为我们的函数使用更多的 $L 2$ 缓存，所以让我们试试看。

首先，让我们修改 basel_chunks 以接受一个 $N$ 的范围作为输入。

# (N1, N2]
def basel_chunks_range(N1: int, N2: int, chunk_size: int):
    # 计时代码省略
    s = 0.0
    num_chunks = (N2 - N1) // chunk_size
    for i in range(num_chunks):
        s += basel_np_range(N1 + i * chunk_size + 1, N1 + (i + 1) * chunk_size)
    return s

现在进行实际的多进程处理：

from multiprocessing import Pool

def basel_multicore(N: int, chunk_size: int):
    # 我的机器上有 10 个核心
    NUM_CORES = 10
    N_PER_CORE = N // NUM_CORES
    Ns = [
        (i * N_PER_CORE, (i + 1) * N_PER_CORE, chunk_size)
        for i in range(NUM_CORES)
    ]
    # 并行处理 10 个批次
    with Pool(NUM_CORES) as p:
        result = p.starmap(basel_chunks_range, Ns)
    return sum(result)

在底层，Python 的 multiprocessing 模块会将函数对象 pickle 化，并生成 9 个新的 python3.10 实例，所有这些实例都会以 $N = 1 0^{9}$ 和 num_chunks = 50000 执行该函数。

让我们比较一下性能。

函数 'basel_multicore' 执行时间为 1.156558 秒。
1.6449340658482263
函数 'basel_chunks' 执行时间为 1.027350 秒。
1.6449340658482263

哎呀，它的表现更差了。这可能是因为多进程处理的初始工作成本较高，因此只有在函数完成所需时间相对较长时才会有好处。

将 $N$ 增加到 $1 0^{10}$ ：

函数 'basel_multicore' 执行时间为 2.314828 秒。
1.6449340667482235
函数 'basel_chunks' 执行时间为 10.221904 秒。
1.6449340667482235

这下好了！它快了大约 5 倍。让我们试试 $1 0^{11}$ 。

函数 'basel_multicore' 执行时间为 13.844876 秒。
multicore 1.6449340668379773
函数 'basel_chunks' 执行时间为 102.480372 秒。
chunks 1.6449340668379773

快了大约 8 倍！随着 $N$ 的增加，单核与 10 核的性能比应该接近 10:1。

x 轴是对数刻度，这就是为什么线性增加的运行时间看起来像指数增长。

我们可以看到，使用多进程处理会有大约 1 秒的开销，因此只有在分块运行时间较高时才会有性能优势。这大约发生在 $N = 1.3 \times 1 0^{9}$ 时。然而，之后会有巨大的差异。

通过实验（此处未展示），我发现 chunk_size 大约为 50000 时，对于多进程代码也是最优的，这表明操作系统对单个核心的 $L 2$ 使用施加了一些限制。

结果总结

Python 函数	$N = 1 0^{8}$ (秒)	$N = 1 0^{9}$ (秒)	$N = 1 0^{10}$ (秒)
`basel_multicore`	1.04	1.17	2.33
`basel_chunks`	0.09	0.91	9.18
`basel_np_broadcast`	0.27	9.68	NaN (内存不足)
`basel_less_pythonic`	5.99	59.25	592 (估计值)
`basel` (惯用写法)	7.64	68.28	682 (估计值)
`basel_np`	0.618	44.27	NaN (内存不足)

与上一篇文章中的语言相比如何？

语言	$N = 1 0^{9}$ (毫秒)
Rust (rc, –release, 多核 w/ rayon)	112.65
Python3.10 (`basel_chunks`)	913
Rust (rc, –release)	937.94
C (clang, -O3)	995.38
Python3.10 (`basel_multicore`)	1170
Python3.10 (`basel_np_broadcast`)	9680
Haskell (ghc, -O3)	13454
Python3.10 (`basel_np`)	44270
Python3.10 (`basel_less_pythonic`)	59250
Python3.10 (`basel`)	68280

非常出色！分块的 Numpy 代码是最快的顺序解决方案！而且请记住，Python 在进行计算时，整个解释器都在后台运行。

将 Rust 与 Python 多核进行比较：

语言	$N = 1 0^{8}$ (毫秒)	$N = 1 0^{9}$ (毫秒)	$N = 1 0^{10}$ (毫秒)	$N = 1 0^{11}$ (毫秒)
Rust	12.3	111.2	1083	10970
Python	1040	1173	2330	12629

显然，Rust 的并行化方法（通过操作系统线程）在 $N$ 较小时比 Python 基于进程的并行化效率高得多。但随着 $N$ 的增加，两者之间的差距逐渐缩小。

结论

如果还不清楚的话，整篇文章其实只是一个练习，目的是了解一些 Python 底层的运作机制。请不要试图通过过度优化你的 Python 代码来给你的经理留下深刻印象。几乎总有更好的解决方案，比如：

from math import pi

def basel_smart() -> float:
    return pi*pi / 6

瞧！比所有其他函数都快了几个数量级，而且无限简单。

在考虑性能时，你需要小心。不要假设性能是个问题，除非你测量过并且知道它确实是。如果是这样，在你急着输入 import multiprocessing 并用蛮力解决一个本来很简单的问题之前，先看看是否有更聪明的方法。

此外，如果你的应用程序真的对性能要求如此之高，你可能不应该用 Python 来编写它。Python 被设计为一种擅长原型的工具，而不是执行速度。但正如你所看到的，要获得出色的性能也并非不可能。

无论如何，我希望你喜欢我在这个网站上的第一篇真正的文章。如果你找到了一个更快的 Python “解决方案”来解决这个“问题”，请给我发邮件！

更新日志

日期	描述	感谢
2023年8月2日	新增广播部分，Python成为最快解决方案	u/mcmcmcmcmcmcmcmcmc_ & u/Allanon001
2023年8月2日	将 `time.time()` 改为 `time.perf_counter()`	u/james_pic

←

巴塞尔问题（你好，世界！）

交互式MNIST探索器

→