使用 Numba 加速 Python 嵌套循环：性能优化教程-Python教程-PHP中文网

使用 Numba 加速 Python 嵌套循环：性能优化教程

心靈之曲

发布： 2025-10-17 11:48:02

原创

452人浏览过

使用 numba 加速 python 嵌套循环：性能优化教程

本文针对Python中嵌套循环计算密集型任务的性能瓶颈，提供了一种有效的解决方案：使用Numba库进行即时编译(JIT)。通过Numba的`@njit`装饰器和并行计算特性，可以显著提升代码执行速度，尤其是在处理大型数据集时。本文将详细介绍如何使用Numba加速嵌套循环，并提供性能对比示例，帮助读者优化Python代码，提高计算效率。

Numba 简介

Numba 是一个开源的 Python 编译器，它使用 LLVM 将 Python 代码转换为优化的机器代码。Numba 的核心在于其即时编译 (JIT) 能力，这意味着它可以在运行时编译 Python 代码，从而显著提高性能。Numba 特别擅长加速数值计算密集型的代码，例如包含循环、数组操作和数学函数的代码。

优化嵌套循环的步骤

以下是如何使用 Numba 加速 Python 中嵌套循环的步骤：

安装 Numba:

立即学习“Python免费学习笔记（深入）”；

首先，确保你已经安装了 Numba。可以使用 pip 进行安装：
```
pip install numba
```
登录后复制
导入 Numba:

在你的 Python 脚本中导入 numba 库。
```
from numba import njit, prange
import numpy as np # 引入 numpy
```
登录后复制
使用 @njit 装饰器:

度加剪辑
度加剪辑（原度咔剪辑），百度旗下AI创作工具

63

查看详情

在要加速的函数上添加 @njit 装饰器。这将指示 Numba 编译该函数。
```
@njit
def your_function(args):
    # 包含嵌套循环的代码
    ...
    return result
```
登录后复制
考虑并行化 (可选):

对于可以并行执行的循环，可以使用 prange 替换 range，并使用 @njit(parallel=True) 装饰器。这将允许 Numba 在多个 CPU 核心上并行执行循环。
```
@njit(parallel=True)
def your_function(args):
    # 包含嵌套循环的代码
    for i in prange(len(data)):
        ...
    return result
```
登录后复制

示例代码

以下是一个使用 Numba 加速嵌套循环的示例。该示例基于问题中提供的代码，并展示了如何使用 @njit 和并行化来提高性能。

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

P_mean = 1500
P_std = 100
Q_mean = 1500
Q_std = 100
W = 1  # Number of matches won by P
L = 0  # Number of matches lost by P
L_P = np.exp(-0.5 * ((np.arange(0, 3501, 10) - P_mean) / P_std) ** 2) / (
    P_std * np.sqrt(2 * np.pi)
)
L_Q = np.exp(-0.5 * ((np.arange(0, 3501, 10) - Q_mean) / Q_std) ** 2) / (
    Q_std * np.sqrt(2 * np.pi)
)


def probability_of_loss(x):
    return 1 / (1 + np.exp(x / 67))


def U_p_law(W, L, L_P, L_Q):
    omega = np.arange(0, 3501, 10)

    U_p = np.zeros_like(omega, dtype=float)

    for p_idx, p in enumerate(omega):
        for q_idx, q in enumerate(omega):
            U_p[p_idx] += (
                probability_of_loss(q - p) ** W
                * probability_of_loss(p - q) ** L
                * L_Q[q_idx]
                * L_P[p_idx]
            )

    normalization_factor = np.sum(U_p)
    U_p /= normalization_factor

    return omega, U_p


@njit
def probability_of_loss_numba(x):
    return 1 / (1 + np.exp(x / 67))


@njit
def U_p_law_numba(W, L, L_P, L_Q):
    omega = np.arange(0, 3501, 10, dtype=np.float64)

    U_p = np.zeros_like(omega)

    for p_idx, p in enumerate(omega):
        for q_idx, q in enumerate(omega):
            U_p[p_idx] += (
                probability_of_loss_numba(q - p) ** W
                * probability_of_loss_numba(p - q) ** L
                * L_Q[q_idx]
                * L_P[p_idx]
            )

    normalization_factor = np.sum(U_p)
    U_p /= normalization_factor

    return omega, U_p


@njit(parallel=True)
def U_p_law_numba_parallel(W, L, L_P, L_Q):
    omega = np.arange(0, 3501, 10, dtype=np.float64)

    U_p = np.zeros_like(omega)

    for p_idx in prange(len(omega)):
        p = omega[p_idx]
        for q_idx in prange(len(omega)):
            q = omega[q_idx]
            U_p[p_idx] += (
                probability_of_loss_numba(q - p) ** W
                * probability_of_loss_numba(p - q) ** L
                * L_Q[q_idx]
                * L_P[p_idx]
            )

    normalization_factor = np.sum(U_p)
    U_p /= normalization_factor

    return omega, U_p


omega_1, U_p_1 = U_p_law(W, L, L_P, L_Q)
omega_2, U_p_2 = U_p_law_numba(W, L, L_P, L_Q)
omega_3, U_p_3 = U_p_law_numba_parallel(W, L, L_P, L_Q)

assert np.allclose(omega_1, omega_2)
assert np.allclose(omega_1, omega_3)
assert np.allclose(U_p_1, U_p_2)
assert np.allclose(U_p_1, U_p_3)

t1 = timeit("U_p_law(W, L, L_P, L_Q)", number=10, globals=globals())
t2 = timeit("U_p_law_numba(W, L, L_P, L_Q)", number=10, globals=globals())
t3 = timeit("U_p_law_numba_parallel(W, L, L_P, L_Q)", number=10, globals=globals())

print("10 calls using vanilla Python     :", t1)
print("10 calls using Numba              :", t2)
print("10 calls using Numba (+ parallel) :", t3)

登录后复制

代码解释:

probability_of_loss_numba: 使用 @njit 装饰器加速 probability_of_loss 函数。
U_p_law_numba: 使用 @njit 装饰器加速原始函数。
U_p_law_numba_parallel: 使用 @njit(parallel=True) 装饰器加速原始函数，并使用 prange 进行并行化。
assert np.allclose(...): 验证 Numba 加速后的函数结果与原始函数结果是否一致，确保正确性。
timeit: 使用 timeit 模块测量不同版本的函数执行时间，进行性能比较。

输出示例 (AMD 5700x):

10 calls using vanilla Python     : 2.4276352748274803
10 calls using Numba              : 0.013957140035927296
10 calls using Numba (+ parallel) : 0.003793451003730297

登录后复制

正如输出所示，使用 Numba 可以显著提高代码的执行速度。

注意事项

数据类型: Numba 在处理 NumPy 数组时效果最佳。确保你的数据存储在 NumPy 数组中。
首次运行时间: Numba 需要一些时间来编译函数。因此，首次运行使用 @njit 装饰的函数可能会比未装饰的函数慢。但是，后续运行将会非常快。
支持的 Python 功能: Numba 并非支持所有的 Python 功能。在使用 Numba 之前，请查阅 Numba 的官方文档，了解其支持的功能。
错误处理: Numba 在编译时可能会报错。仔细阅读错误信息，并根据提示修改代码。
并行化: 并非所有循环都适合并行化。确保循环的迭代之间没有依赖关系。
fastmath 参数： 对于一些数学运算，可以尝试使用 @njit(fastmath=True)。fastmath 允许编译器进行更激进的优化，但这可能会导致一些精度损失。请根据你的应用场景权衡精度和性能。