使用 numba 加速 for 循环

Posted 2023-02-16

【中文标题】使用 numba 加速 for 循环

【英文标题】：Use numba to speed up for loop

【发布时间】：2016-06-17 09:59:51

【问题描述】：

据我所知，numba 可以显着加快 Python 程序的速度。使用 numba 可以提高程序的时间效率吗？

import numpy as np

def f_big(A, k, std_A, std_k, mean_A=10, mean_k=0.2, hh=100):
    return ( 1 / (std_A * std_k * 2 * np.pi) ) * A * (hh/50) ** k * np.exp( -1*(k - mean_k)**2 / (2 * std_k **2 ) - (A - mean_A)**2 / (2 * std_A**2))

outer_sum = 0
dk = 0.000001
for k in np.arange(dk,0.4, dk):
    inner_sum = 0
    for A in np.arange(dk, 20, dk):
        inner_sum += dk * f_big(A, k, 1e-5, 1e-5)
    outer_sum += inner_sum * dk

print outer_sum

【参考方案1】：

是的，这是 Numba 真正解决的问题。我更改了您的 dk 值，因为它对于简单的演示来说是不明智的。代码如下：

import numpy as np
import numba as nb

def f_big(A, k, std_A, std_k, mean_A=10, mean_k=0.2, hh=100):
    return ( 1 / (std_A * std_k * 2 * np.pi) ) * A * (hh/50) ** k * np.exp( -1*(k - mean_k)**2 / (2 * std_k **2 ) - (A - mean_A)**2 / (2 * std_A**2))

def func():
    outer_sum = 0
    dk = 0.01 #0.000001
    for k in np.arange(dk, 0.4, dk):
        inner_sum = 0
        for A in np.arange(dk, 20, dk):
            inner_sum += dk * f_big(A, k, 1e-5, 1e-5)
        outer_sum += inner_sum * dk

    return outer_sum

@nb.jit(nopython=True)
def f_big_nb(A, k, std_A, std_k, mean_A=10, mean_k=0.2, hh=100):
    return ( 1 / (std_A * std_k * 2 * np.pi) ) * A * (hh/50) ** k * np.exp( -1*(k - mean_k)**2 / (2 * std_k **2 ) - (A - mean_A)**2 / (2 * std_A**2))

@nb.jit(nopython=True)
def func_nb():
    outer_sum = 0
    dk = 0.01 #0.000001
    X = np.arange(dk, 0.4, dk)
    Y = np.arange(dk, 20, dk)
    for i in xrange(X.shape[0]):
        k = X[i] # faster to do lookup than iterate over an array directly
        inner_sum = 0
        for j in xrange(Y.shape[0]):
            A = Y[j]
            inner_sum += dk * f_big_nb(A, k, 1e-5, 1e-5)
        outer_sum += inner_sum * dk

    return outer_sum

然后是时间安排：

In [7]: np.allclose(func(), func_nb())
Out[7]: True

In [8]: %timeit func()
1 loops, best of 3: 222 ms per loop

In [9]: %timeit func_nb()
The slowest run took 419.10 times longer than the fastest. This could mean that an intermediate result is being cached 
1000 loops, best of 3: 362 µs per loop

所以 numba 版本在我的笔记本电脑上快了大约 600 倍。

【讨论】：

可能很挑剔，但您可以使用 @nb.njit 而不使用 nopython=True，而不是使用 @nb.jit(nopython=True)。

@MSeifert 我倾向于习惯性地使用这种形式，因为我会经常对其进行参数化，以便在测试期间轻松地来回切换

xrange 应该是range，对吧？

使用python3时应该是range。早在 2016 年，当我最初做出回应时，我仍然主要使用 python2，其中xrange 是有效的。