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

如何使用 Numba 加速 Python 中的嵌套循环

DDD

发布： 2025-08-27 15:34:01

原创

305人浏览过

如何使用 numba 加速 python 中的嵌套循环

本文旨在提供一种使用 Numba 库加速 Python 中嵌套循环计算的方法。通过使用 Numba 的 JIT 编译和并行计算功能，可以将原本耗时较长的嵌套循环代码的执行速度显著提高，从而更高效地完成计算任务。文章将提供详细的代码示例和优化技巧，帮助读者理解和应用这些技术。

使用 Numba 加速嵌套循环

在 Python 中，嵌套循环可能会导致代码执行速度缓慢，尤其是在循环次数较多时。为了解决这个问题，可以使用 Numba 库，它是一个即时 (JIT) 编译器，可以将 Python 代码编译成机器码，从而提高执行速度。

Numba 的 njit 装饰器

Numba 提供了一个名为 njit 的装饰器，可以将其应用于函数，以指示 Numba 对该函数进行 JIT 编译。以下是一个简单的示例：

from numba import njit

@njit
def my_function():
    # 你的代码
    pass

登录后复制

将 @njit 装饰器添加到函数后，Numba 将在首次调用该函数时对其进行编译。后续调用将直接执行编译后的机器码，从而提高执行速度。

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

示例：加速四重嵌套循环

考虑以下四重嵌套循环的代码：

def original_function():
    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),
                        )

登录后复制

这段代码计算 2**a * 3**b * 5**c * 7**d 的值，并检查其是否接近 0.3048。如果没有 Numba，这段代码的执行时间会很长。

豆绘AI

豆绘AI是国内领先的AI绘图与设计平台，支持照片、设计、绘画的一键生成。

485

查看详情

现在，使用 Numba 的 njit 装饰器来加速这段代码：

from numba import njit

@njit
def optimized_function():
    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),
                        )

登录后复制

通过添加 @njit 装饰器，可以将代码的执行速度显著提高。

使用 prange 进行并行计算

Numba 还提供了一个名为 prange 的函数，可以用于并行执行循环。这对于具有大量迭代的循环非常有用。要使用 prange，需要将 @njit 装饰器的 parallel 参数设置为 True。

from numba import njit, prange

@njit(parallel=True)
def parallel_function():
    for a in prange(-100, 101):
        # 你的代码
        pass

登录后复制

以下是如何使用 prange 加速四重嵌套循环的代码：

from numba import njit, prange

@njit(parallel=True)
def optimized_parallel_function():
    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),
                        )

登录后复制

在这个例子中，我们使用 prange 替换了 range，并将 @njit 装饰器的 parallel 参数设置为 True。此外，为了减少重复计算，我们存储了中间结果 i_a, i_b, i_c。这使得 Numba 能够并行执行循环，从而进一步提高执行速度。