本文将探讨如何使用 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 函数编译为机器码。这将显著提高循环的执行速度。
使用 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 进行并行化。
存储中间结果以优化性能
在某些情况下，可以在循环中存储中间结果，以避免重复计算，从而进一步提高性能。在上面的例子中，可以在每次迭代中存储 2.0**a、i_a * 3.0**b 和 i_b * 5.0**c 的值，并在后续迭代中重用这些值。
总结
通过使用 Numba 的 @njit 装饰器和 prange 函数，可以显著提高 Python 中嵌套循环的执行速度。这对于需要大量数值计算的任务尤其有用。此外，通过存储中间结果，可以进一步优化代码，从而获得更好的性能。在编写计算密集型 Python 代码时，Numba 是一个非常有用的工具。
