本文旨在提供一种优化Python中嵌套循环计算效率的方法，特别是针对计算密集型任务。通过使用Numba库的即时编译（JIT）技术，可以显著提升代码的执行速度，避免传统嵌套循环带来的性能瓶颈。文章将展示如何使用Numba加速原始代码，并提供并行化的优化方案，以及性能对比。
在Python中，嵌套循环是常见的编程结构，但当循环次数较多时，其执行效率会显著下降。对于需要处理大量数据的科学计算或数据分析任务，优化嵌套循环至关重要。Numba是一个开源的即时编译器，它可以将Python代码转换为优化的机器码，从而显著提高程序的运行速度。本文将介绍如何使用Numba来优化包含嵌套循环的Python函数。
使用Numba加速计算
首先，我们来看一个包含嵌套循环的示例函数 U_p_law，该函数计算两个概率密度函数之间的关系：
import numpy as np
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为了使用Numba加速这个函数，我们需要导入 numba 库，并使用 @njit 装饰器修饰函数。同时，为了获得更好的性能，我们也可以对 probability_of_loss 函数进行加速。
立即学习“Python免费学习笔记（深入）”；
from numba import njit
@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 装饰器告诉Numba将该函数编译为机器码。
- 为了保证Numba能够成功编译，我们需要确保函数中使用的所有操作和数据类型都受Numba支持。例如，如果使用NumPy数组，需要确保数组的数据类型是Numba支持的类型。
- Numba 推荐使用 np.float64 作为浮点数类型，以获得更好的性能。
使用并行化进一步加速
对于计算量更大的任务，我们可以利用多核CPU的优势，使用Numba的并行化功能。为了实现并行化，我们需要使用 parallel=True 参数修饰 @njit 装饰器，并将外层循环替换为 prange。
from numba import njit, prange
@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注意:
- prange 是 Numba 提供的并行循环，它会将循环迭代分配到多个线程上执行。
- 并行化可以显著提高程序的运行速度，但也会带来一些额外的开销，例如线程创建和同步。因此，只有当计算量足够大时，并行化才能带来明显的性能提升。
性能测试
为了验证Numba的加速效果，我们可以使用 timeit 模块来测试不同版本的函数的运行时间。
from timeit import timeit
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)
)
# 确保结果一致
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)在我的机器上（AMD 5700x），运行结果如下：
10 calls using vanilla Python : 2.4276352748274803 10 calls using Numba : 0.013957140035927296 10 calls using Numba (+ parallel) : 0.003793451003730297
从结果可以看出，使用 Numba 可以显著提高程序的运行速度。在这个例子中，使用 Numba JIT 可以提速约 170 倍，而使用多线程 Numba JIT 可以提速约 640 倍。
总结
Numba 是一个强大的工具，可以用来优化 Python 代码的性能，特别是对于包含嵌套循环的计算密集型任务。通过使用 Numba 的即时编译和并行化功能，我们可以显著提高程序的运行速度，从而更快地完成计算任务。在使用 Numba 时，需要注意确保函数中使用的所有操作和数据类型都受 Numba 支持，并根据实际情况选择合适的优化策略。
