使用 Numba 加速 Python 嵌套循环-Python教程-PHP中文网

使用 Numba 加速 Python 嵌套循环

花韻仙語

发布： 2025-08-27 16:17:01

原创

206人浏览过

使用 numba 加速 python 嵌套循环

本文将探讨如何使用 Numba 库中的 Just-In-Time (JIT) 编译器来显著提升 Python 中嵌套循环的执行速度。通过简单的装饰器 @njit 和 prange，可以将耗时的循环计算加速数十倍，尤其是在涉及大量数值计算的场景中。此外，文章还展示了如何通过存储中间结果来进一步优化代码，充分利用并行计算的优势。

利用 Numba 加速嵌套循环

Python 是一种解释型语言，在执行循环时效率相对较低，尤其是在多层嵌套循环中。对于需要大量数值计算的任务，这种低效率会变得非常明显。Numba 是一个开源的 JIT 编译器，可以将 Python 代码转换为优化的机器码，从而显著提高执行速度。

使用 @njit 装饰器

Numba 的核心功能是通过 @njit 装饰器实现的。这个装饰器告诉 Numba 编译被装饰的函数，使其能够以接近 C 或 Fortran 的速度运行。

以下是一个示例，展示了如何使用 @njit 来加速一个包含四个嵌套循环的函数：

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

from numba import njit

@njit
def fn():
    for a in range(-100, 101):
        for b in range(-100, 101):
            for c in range(-100, 101):
                for d in range(-100, 101):
                    n = (2.0**a) * (3.0**b) * (5.0**c) * (7.0**d)
                    v = n - 0.3048
                    if abs(v) <= 1e-06:
                        print(
                            "a=",
                            a,
                            ", b=",
                            b,
                            ", c=",
                            c,
                            ", d=",
                            d,
                            ", the number=",
                            n,
                            ", error=",
                            abs(n - 3.048),
                        )

fn()

登录后复制

在这个例子中，@njit 装饰器指示 Numba 将 fn 函数编译为机器码。这将显著提高循环的执行速度。

度加剪辑

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

查看详情

使用 prange 实现并行化

对于计算密集型任务，还可以利用多核 CPU 的优势，通过并行化来进一步提高性能。Numba 提供了 prange 函数，它是 range 函数的并行版本。要使用 prange，需要在 @njit 装饰器中启用 parallel=True 选项。

from numba import njit, prange

@njit(parallel=True)
def fn():
    for a in prange(-100, 101):
        i_a = 2.0**a
        for b in prange(-100, 101):
            i_b = i_a * 3.0**b
            for c in prange(-100, 101):
                i_c = i_b * 5.0**c
                for d in prange(-100, 101):
                    n = i_c * (7.0**d)
                    v = n - 0.3048
                    if abs(v) <= 1e-06:
                        print(
                            "a=",
                            a,
                            ", b=",
                            b,
                            ", c=",
                            c,
                            ", d=",
                            d,
                            ", the number=",
                            n,
                            ", error=",
                            abs(n - 3.048),
                        )

fn()

登录后复制

在这个例子中，prange 替换了 range，并且 @njit 装饰器中设置了 parallel=True。这将使 Numba 将循环并行化，从而利用多核 CPU 的优势。

注意： 为了获得最佳性能，建议将计算量最大的循环放在最外层，并使用 prange 进行并行化。