本文针对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, prangeimport numpy as np # 引入 numpy
使用 @njit 装饰器:
度加剪辑
度加剪辑（原度咔剪辑），百度旗下AI创作工具
63 查看详情 
在要加速的函数上添加 @njit 装饰器。这将指示 Numba 编译该函数。
@njitdef 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 timeitfrom numba import njit, prangeimport numpy as npP_mean = 1500P_std = 100Q_mean = 1500Q_std = 100W = 1 # Number of matches won by PL = 0 # Number of matches lost by PL_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@njitdef probability_of_loss_numba(x): return 1 / (1 + np.exp(x / 67))@njitdef 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_pomega_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.427635274827480310 calls using Numba : 0.01395714003592729610 calls using Numba (+ parallel) : 0.003793451003730297
正如输出所示，使用 Numba 可以显著提高代码的执行速度。
注意事项
数据类型: Numba 在处理 NumPy 数组时效果最佳。确保你的数据存储在 NumPy 数组中。首次运行时间: Numba 需要一些时间来编译函数。因此，首次运行使用 @njit 装饰的函数可能会比未装饰的函数慢。但是，后续运行将会非常快。支持的 Python 功能: Numba 并非支持所有的 Python 功能。在使用 Numba 之前，请查阅 Numba 的官方文档，了解其支持的功能。错误处理: Numba 在编译时可能会报错。仔细阅读错误信息，并根据提示修改代码。并行化: 并非所有循环都适合并行化。确保循环的迭代之间没有依赖关系。fastmath 参数： 对于一些数学运算，可以尝试使用 @njit(fastmath=True)。fastmath 允许编译器进行更激进的优化，但这可能会导致一些精度损失。请根据你的应用场景权衡精度和性能。
总结
Numba 是一个强大的工具，可以显著提高 Python 中数值计算密集型代码的性能。通过使用 @njit 装饰器和并行化，可以轻松加速包含嵌套循环的代码。希望本教程能够帮助你优化 Python 代码，提高计算效率。
以上就是使用 Numba 加速 Python 嵌套循环：性能优化教程的详细内容，更多请关注创想鸟其它相关文章！
版权声明：本文内容由互联网用户自发贡献，该文观点仅代表作者本人。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。
如发现本站有涉嫌抄袭侵权/违法违规的内容， 请发送邮件至 chuangxiangniao@163.com 举报，一经查实，本站将立刻删除。
发布者：程序猿，转转请注明出处：https://www.chuangxiangniao.com/p/589676.html
微信扫一扫 
支付宝扫一扫 
