CPython&Numba | Python程序加速小技巧

本期推文主要分享 Python 程序加速小妙招：Cpython 和 Numba。

 

写在前面的话：由于 MFEA 现已逐步转向 Python，所以在平时的空余时间我会去关注一些 Python 编程的知识，号内近期更新的推文较多是关于 Python。相同的道理，如果以后程序转向 C++的话，号内推文也会向 C++倾斜。

分享的主要内容:

• CPython
• 概念
• 简易使用方法
• Numba
• 概念
• 简易使用方法
• CPython & Numba & 纯 Python & 内置函数，性能对比
•       矩阵乘法
•       矩阵乘法

CPython

CPython 是用 C 语言编写的 Python 解释器，因此得名 “CPython”。它遵循 Python 语言的官方规范，是由 Python 的原始作者 Guido van Rossum 及其后继者们维护的。在 CPython 中，Python 源代码首先被解析成字节码，然后由 CPython 的虚拟机执行。

小伙伴对以上概念可以先不用看，只需记得，可以加速我们的程序就行，接下来看一下如何借助 CPython 加速 Python 程序。

网上有很多有关 CPython 的使用教程，有些操作可能比较高端，难以操作，本次推文只是记录一些简学易用的小妙招，更高端的操作可逛一逛 Github 或者 B 站。

步骤一

将代码保存到一个文件中，例如 compute_square_sum.pyx：

# compute_square_sum.pyx

def compute_square_sum(numbers):
    cdef int total = 0
    cdef int number
    for number in numbers:
        total += number ** 2
    return total

• cdef 是 Cython 的关键字，用于声明 C 语言的变量类型。在这里，total 被声明为 int 类型，并初始化为 0。使用 cdef 声明变量类型可以显著提高执行效率，因为它允许 Cython 将操作编译为更高效的 C 代码。
• 由于 total 是 int 类型，所有对 total 的操作都将在 C 语言级别执行，而不是 Python 层次。Cython 可以在 C 级别处理这些变量，避免了 Python 层次的动态类型检查和操作，这样可以显著提高执行速度。相比于纯 Python 实现，使用 Cython 编写的代码在数值计算密集型任务中通常能获得显著的性能提升。

步骤二

创建 setup.py 文件用于编译 Cython 代码：

from setuptools import setup
from Cython.Build import cythonize

setup(
    ext_modules = cythonize("compute_square_sum.pyx")
)

步骤三

在终端中运行以下命令进行编译：

python setup.py build_ext --inplace

编译成功后，就可以像调用普通 Python 模块一样导入并使用这个函数。就像这样：

from compute_square_sum import compute_square_sum

# 测试
numbers = list(range(1_000_000))
result = compute_square_sum(numbers)

Numba

Numba 是一个用于 Python 的开源 JIT（Just-In-Time）编译器，它能够将 Python 代码中的数值密集型部分即时编译为高效的机器代码，从而大幅提高程序的执行速度。Numba 的设计目标是为科学计算、数据分析和机器学习等领域提供快速、易用的加速工具，而无需离开 Python 语言的舒适环境。

中文文档：https://apachecn.github.io/numba-doc-zh/#/

核心功能

• JIT 编译: Numba 通过 JIT 编译将特定的 Python 函数即时编译为机器代码。这使得 Python 代码在运行时能够接近 C 或 Fortran 代码的性能，尤其在数值计算中。
• GPU 加速: Numba 支持 CUDA 编程，可以将代码编译为 GPU 上运行的并行代码，这对于处理大规模数据集和需要大量并行计算的任务尤为有用。
• 自动并行化: 通过并行选项，Numba 可以将循环和其他计算任务自动并行化，充分利用多核 CPU 的性能。
• 支持 NumPy: Numba 与 NumPy 结合紧密，可以直接加速 NumPy 数组操作。这使得 Numba 成为加速科学计算代码的理想选择。

基本用法

Numba 的使用非常简单。只需在需要加速的 Python 函数上添加装饰器 @jit 或 @njit，Numba 就会在第一次调用该函数时编译其为机器代码。

from numba import jit

@jit
def sum_array(x):
    total = 0
    for i in range(x.shape[0]):
        total += x[i]
    return total

import numpy as np
arr = np.random.rand(1000000)
print(sum_array(arr))

在这个示例中，sum_array 函数使用 Numba 加速后，其执行速度会显著提升，尤其是在处理大规模数据时。

@jit 与 @njit

• @jit: @jit 是 Numba 的通用装饰器，它允许函数包含一些 Python 特性，如错误处理等。它可以选择使用 nopython=True 参数来要求 Numba 在不使用 Python 解释器的情况下运行代码。
• @njit: @njit 是 @jit(nopython=True) 的简写，强制要求 Numba 使用 nopython 模式。nopython 模式下，所有代码都被编译为机器代码，不依赖 Python 解释器，从而获得更高的性能。

并行化

Numba 支持并行计算，通过 @njit(parallel=True) 可以将循环自动并行化，充分利用多核 CPU。需要使用 prange 来替代 Python 的 range 函数以支持并行。

from numba import njit, prange

@njit(parallel=True)
def parallel_sum(x):
    total = 0
    for i in prange(x.shape[0]):
        total += x[i]
    return total

arr = np.random.rand(1000000)
print(parallel_sum(arr))

在这个示例中，parallel_sum 函数中的循环被自动并行化，从而在多核 CPU 上运行更快。

GPU 加速

Numba 提供了对 CUDA 编程的支持，可以将计算卸载到 GPU 上。通过 @cuda.jit 装饰器，可以在 GPU 上运行 Python 函数。

from numba import cuda
import numpy as np

@cuda.jit
def gpu_addition(a, b, c):
    i = cuda.grid(1)
    if i < a.size:
        c[i] = a[i] + b[i]

n = 100000
a = np.random.rand(n)
b = np.random.rand(n)
c = np.zeros(n)

gpu_addition[n//256 + 1, 256](a, b, c)
print(c[:10])

在这个示例中，gpu_addition 函数在 GPU 上并行执行，使得对大规模数据的操作更为高效。

案例分析

接下来以矩阵乘法为例，以     矩阵为例，探究 CPython & Numba & 纯 Python & 内置函数，它们的运行速度。

import numpy as np
import timeit
from numba import jit, njit, prange
from matrix_multiply import matrix_multiply

# 使用纯 Python 实现的矩阵乘法
def matrix_multiply_python(A, B):
    n, m = A.shape
    m, p = B.shape
    result = np.zeros((n, p))
    for i in range(n):
        for j in range(p):
            for k in range(m):
                result[i, j] += A[i, k] * B[k, j]
    return result

# 使用 Numba 并行加速的矩阵乘法
@njit(parallel=True)
def matrix_multiply_numba_parallel(A, B):
    n = A.shape[0]
    m = A.shape[1]
    p = B.shape[1]
    result = np.zeros((n, p))
    
    for i in prange(n):
        for j in prange(p):
            for k in prange(m):
                result[i, j] += A[i, k] * B[k, j]
    
    return result

# 使用 Numba 加速的矩阵乘法
@jit(nopython=True)
def matrix_multiply_numba(A, B):
    n = A.shape[0]
    m = A.shape[1]
    p = B.shape[1]
    result = np.zeros((n, p))
    
    for i in range(n):
        for j in range(p):
            for k in range(m):
                result[i, j] += A[i, k] * B[k, j]
    
    return result

# 生成随机矩阵作为测试数据
A = np.random.rand(200, 200)
B = np.random.rand(200, 200)

print(f"矩阵 A 的尺寸: {A.shape}")
print(f"矩阵 B 的尺寸: {B.shape}")

# 测试纯 Python 的矩阵乘法运行时间
python_time = timeit.timeit('matrix_multiply_python(A, B)', globals=globals(), number=10)
print(f"纯 Python 矩阵乘法运行时间: {python_time:.4f} 秒")

# 测试 Numba 加速的矩阵乘法运行时间
numba_time = timeit.timeit('matrix_multiply_numba(A, B)', globals=globals(), number=10)
print(f"Numba 加速矩阵乘法运行时间: {numba_time:.4f} 秒")

# 测试 Numba 并行加速的矩阵乘法运行时间
numba_time = timeit.timeit('matrix_multiply_numba_parallel(A, B)', globals=globals(), number=10)
print(f"Numba 并行加速矩阵乘法运行时间: {numba_time:.4f} 秒")

# 测试 CPython 加速的矩阵乘法运行时间
cpython_time = timeit.timeit('matrix_multiply(A, B)', globals=globals(), number=10)
print(f"CPython 加速矩阵乘法运行时间: {cpython_time:.4f} 秒")

# 测试 Python 内置矩阵乘法 (numpy.dot) 的运行时间
numpy_time = timeit.timeit('np.dot(A, B)', globals=globals(), number=10)
print(f"Python 内置 (numpy.dot) 矩阵乘法运行时间: {numpy_time:.4f} 秒")

运行结果：

矩阵 A 的尺寸: (200, 200)
矩阵 B 的尺寸: (200, 200)
纯 Python 矩阵乘法运行时间: 45.7477 秒
Numba 加速矩阵乘法运行时间: 0.7108 秒
Numba 并行加速矩阵乘法运行时间: 0.7129 秒
CPython 加速矩阵乘法运行时间: 0.4730 秒
Python 内置 (numpy.dot) 矩阵乘法运行时间: 0.0051 秒

• Python 内置的 numpy.dot: 在绝大多数情况下，使用 numpy.dot 是执行矩阵乘法的最佳选择，其性能在所有方法中最为出色。这是因为 numpy.dot 内部使用了高度优化的线性代数库（如 BLAS 或 LAPACK），这些库在底层进行了大量的优化，如向量化、并行化、缓存利用等。
• Numba: 使用 Numba 加速的矩阵乘法比纯 Python 实现快了约 64 倍。这是因为 Numba 使用 Just-In-Time (JIT) 编译，将 Python 代码即时编译为高效的机器代码，从而减少了解释器的开销。尽管单线程执行，但 JIT 编译显著提高了代码的执行效率。
• CPython：Cython 将 Python 代码编译为 C 扩展模块后，结合静态类型声明和 C 语言的执行效率，能够在矩阵乘法等数值密集型计算中提供非常高效的解决方案。
• 并行化: 对于小规模的计算，并行化的优势可能不明显。只有在处理更大规模的数据集时，或者在多核、多线程环境下，Numba 的并行化才可能显现出更显著的性能提升。
• 纯 Python 实现: 纯 Python 实现尽管可以用于教学或小规模计算，但在性能要求较高的实际应用中，使用纯 Python 进行矩阵乘法是不可取的。

接下来以     的矩阵乘法为例：

矩阵 A 的尺寸: (1000, 1000)
矩阵 B 的尺寸: (1000, 1000)
Numba 加速矩阵乘法运行时间: 17.7945 秒
Numba 并行加速矩阵乘法运行时间: 12.3544 秒
CPython 加速矩阵乘法运行时间: 68.4504 秒
Python 内置 (numpy.dot) 矩阵乘法运行时间: 0.2398 秒

纯 Python 的时间就不计算在内了，很慢！

• 从上面运行结果来看，当计算量越来越大时，CPython 的时间并没有 Numba 快，可能是我没有设置好，只是简易化的使用
• 并行计算的优势就体现的比较明显，如果矩阵尺寸更大，并行的优势将更明显
• 当然，最快的还要看内置函数

结论

在计算量较大的程序中，应尽量避免自编的功能函数，如果有内置函数的，优先使用内置函数，倘若没有的话，可以借助 CPython 或者 Numba 等加速小技巧。