Python中嵌套循环的替代方案：使用Numba加速计算-Python教程-PHP中文网

Python中嵌套循环的替代方案：使用Numba加速计算

聖光之護

发布： 2025-10-17 16:36:02

原创

1015人浏览过

python中嵌套循环的替代方案：使用numba加速计算

本文旨在提供一种优化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 支持，并根据实际情况选择合适的优化策略。

以上就是Python中嵌套循环的替代方案：使用Numba加速计算的详细内容，更多请关注php中文网其它相关文章！