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))

这是最符合Python风格的实现方式。这段单行代码易读且使用了内置生成器。让我们用 $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 循环重写，但采用不那么Python风格的方式：

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

有趣的是，这种非惯用写法反而提升了性能。为什么？让我们使用 Python 反汇编模块 dis 来探究。

这是原始版本的反汇编代码（已添加注释）：

# 加载变量
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

# 调用求和函数
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)

# 跳转至循环开始
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
函数 'basel_np' 执行耗时 0.460317 秒。
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

操作步骤	单步耗时 (毫秒)	累计耗时 (毫秒)
创建全1数组	97	97
创建范围数组	79	176
平方运算	98	274
除法运算	112	387
最终求和	58	444

仔细分析这些数据：创建全1数组和范围数组的步骤并不涉及繁重的计算，但耗时却与平方、除法等"更复杂"的步骤相当。出乎意料的是，对最终数组求和竟然是最快的步骤！这表明性能瓶颈不在 CPU 计算，而在内存访问。让我们观察当 $N$ 从 $1 0^{7}$ 到 $1 0^{9}$ 变化时的性能表现。

在此图表中，最上方线条的边缘代表总耗时，相邻线条之间的区域代表每个步骤的耗时。

我们可以观察到：

总耗时呈非线性增长，尽管算法时间复杂度是 $O (n)$
除法步骤耗时最多，在 $\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}$ 时：

操作步骤	单步耗时 (毫秒)	累计耗时 (毫秒)
创建全1标量	0.00	0
创建等差序列	68.56	70.48
计算序列平方	105.14	180.74
执行标量除法	133.08	271.30
最终求和	71.08	310.41

从现在起，我将把广播版本称为"NumPy解决方案"。

虽然取得了显著改进，但我们仍然观察到那些延迟峰值。接下来让我们尝试解决这个问题。

内存优化思路

为了优化函数的内存使用，我们可以采用分块处理的方式：将整个范围分成小块，对每个小块执行相同操作，最后将结果累加。这样只需在内存中保留一个数据块，既能利用 NumPy 的向量化计算优势，又能提升缓存利用率。

让我们修改函数以支持处理指定范围：

# [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 浮点数的特性，计算结果实际上会略微更精确。现在让我们像之前那样保持分块大小不变，观察性能扩展情况。

首先注意纵轴单位已变为秒级而非分钟级。运行时间呈线性增长，与算法 $O (n)$ 的时间复杂度相符。可以推测在 chunk_size=20000 时，所有数据都能装入缓存，因此不会出现因缓存溢出导致的运行时间突增。

现在固定 $N = 1 0^{9}$ ，在 $[5 \times 1 0^{2}, 1 0^{6}]$ 区间内调整 chunk_size：

从图中可见，当 chunk_size 达到约 51000 之前性能持续提升，之后出现的延迟峰值推测是由缓存未命中引起的。

速算分析

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

我的 M1 Pro 芯片具有以下规格：

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

这表明函数可用的最大 $L 1 + L 2$ 缓存空间约为 $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 模块会将函数对象序列化，并生成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 左右也是最优的，这说明操作系统对单个核心的 L2 缓存使用施加了一些限制。

结果摘要

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	内存不足
`basel_less_pythonic`	5.99	59.25	592 (预估)
`basel` (惯用写法)	7.64	68.28	682 (预估)
`basel_np`	0.618	44.27	内存不足

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

语言	$N = 1 0^{9}$ (毫秒)
Rust (rc, –release, 多核 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

显然，对于较小的 $N$ ，Rust 的并行化方法（通过操作系统线程）远比 Python 基于进程的并行更高效。但随着 $N$ 增大，两者之间的差异会缩小。

总结

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

from math import pi

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

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

在考虑性能时，你需要格外小心。在测量并确认性能确实成为问题之前，不要先入为主地假设它是个问题。如果确实存在性能问题，也请先看看是否有更聪明的解决方案，不要急着输入import multiprocessing，用蛮力去解决一个本来很简单的问题。

另外，如果你的应用真的对性能如此敏感，或许你本就不该选择Python来编写。Python生来就是一款擅长快速原型开发而非执行速度的工具。但正如你所见，要在Python中获得出色性能也并非不可能。

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

更新日志

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

←

巴塞尔问题（Hello, World!）

交互式MNIST探索器

→