使用 Numba 更快地制作四个嵌套 for 循环

Question

我对 Numba 的工作有点陌生，但我已经了解了它的要点。我想知道是否有更高级的技巧可以使四个嵌套for循环比我现在的速度更快。特别是，我需要计算以下积分：

其中B是二维数组，S0和E是某些参数。我的代码如下：

import numpy as np
from numba import njit, double

def calc_gb_gauss_2d(b,s0,e,dx):
    n,m=b.shape
    norm = 1.0/(2*np.pi*s0**2)
    gb = np.zeros((n,m))
    for i in range(n):
        for j in range(m):
            for ii in range(n):
                for jj in range(m):
                    gb[i,j]+=np.exp(-(((i-ii)*dx)**2+((j-jj)*dx)**2)/(2.0*(s0*(1.0+e*b[i,j]))**2))
            gb[i,j]*=norm
    return gb

calc_gb_gauss_2d_nb = njit(double[:, :](double[:, :],double,double,double))(calc_gb_gauss_2d)

对于大小为输入的数组，256x256计算速度为：

In [4]: a=random.random((256,256))

In [5]: %timeit calc_gb_gauss_2d_nb(a,0.1,1.0,0.5)
The slowest run took 8.46 times longer than the fastest. This could mean that an intermediate result is being cached.
1 loop, best of 3: 1min 1s per loop

纯Python和Numba计算速度的比较给我这张图：

有什么方法可以优化我的代码以获得更好的性能吗？

Answer 1

通过使用 numpy 和一些数学，可以加速你的代码，因此它比当前的 numba 版本快一个数量级。我们还将看到，在改进的功能上使用 numba 使其速度更快。

\n\n

numba 经常被过度使用——通常可以编写非常高效的纯 numpy 代码——这里也是这种情况。

\n\n

手头的 numpy 代码存在一个问题：不应访问单个元素，而应利用 numpy 的内置函数 - 它们在大多数情况下都尽可能快。仅当无法使用这些 numpy 函数时，才会使用 numba 或 cython。

\n\n

然而，这里最大的问题是问题的表述。对于固定的i，j我们有以下公式来计算（我简化了一点）：

\n\n

 g[i,j]=sum_ii sum_jj exp(value_ii+value_jj)\n       =sum_ii sum_jj exp(value_ii)*exp(value_jj)\n       =sum_ii exp(value_ii) * sum_jj exp(value_jj)\n

\n\n

为了评估最后一个公式，我们需要O(n+m)进行运算，但对于第一个简单的公式O(n*m)- 差别很大！

\n\n

利用 numpy 功能的第一个版本可能类似于：

\n\n

def calc_ead(b,s0,e,dx):\n    n,m=b.shape\n    norm = 1.0/(2*np.pi*s0**2)\n    gb = np.zeros((n,m))\n    vI=np.arange(n)\n    vJ=np.arange(m)\n    for i in range(n):\n        for j in range(m):\n            II=(i-vI)*dx\n            JJ=(j-vJ)*dx\n            denom=2.0*(s0*(1.0+e*b[i,j]))**2\n            expII=np.exp(-II*II/denom)\n            expJJ=np.exp(-JJ*JJ/denom)\n            gb[i,j]=norm*(expII.sum()*expJJ.sum())\n    return gb\n

\n\n

现在，与原始的 numba 实现相比：

\n\n

>>> a=np.random.random((256,256))\n\n>>> print(calc_gb_gauss_2d_nb(a,0.1,1.0,0.5)[1,1])\n15.9160709993\n>>> %timeit -n1 -r1 calc_gb_gauss_2d_nb(a,0.1,1.0,0.5)\n1min 6s \xc2\xb1 0 ns per loop (mean \xc2\xb1 std. dev. of 1 run, 1 loop each)\n

\n\n

现在 numpy 函数：

\n\n

>>> print(calc_ead(a,0.1,1.0,0.5)[1,1])\n15.9160709993\n>>> %timeit -n1 -r1 calc_ead(a,0.1,1.0,0.5)\n1.8 s \xc2\xb1 0 ns per loop (mean \xc2\xb1 std. dev. of 1 run, 1 loop each)\n

\n\n

有两个观察结果：

\n\n

\n
1. 结果是一样的。
\n
2. numpy 版本快了 37 倍，对于更大的问题，这种差异会变得更大。
\n

\n\n

显然，您可以利用 numba 来实现更大的加速。然而，在可能的情况下使用 numpy 功能仍然是一个好主意 - 令人惊讶的是，最简单的事情可以如此微妙 - 例如甚至计算总和：

\n\n

>>> nb_calc_ead = njit(double[:, :](double[:, :],double,double,double))(calc_ead)\n>>>print(nb_calc_ead(a,0.1,1.0,0.5)[1,1])\n15.9160709993\n>>>%timeit -n1 -r1 nb_calc_ead(a,0.1,1.0,0.5)\n587 ms \xc2\xb1 0 ns per loop (mean \xc2\xb1 std. dev. of 1 run, 1 loop each)\n

\n\n

还有一个因素3！

\n\n

这个问题可以并行化，但是要正确地解决这个问题并不容易。我的廉价尝试使用显式循环并行化：

\n\n

from numba import njit, prange\nimport math\n\n@njit(parallel=True)                 #needed, so it is parallelized\ndef parallel_nb_calc_ead(b,s0,e,dx):\n    n,m=b.shape\n    norm = 1.0/(2*np.pi*s0**2)\n    gb = np.zeros((n,m))\n    vI=np.arange(n)\n    vJ=np.arange(m)\n    for i in prange(n):             #outer loop = explicit prange-loop\n        for j in range(m):\n            denom=2.0*(s0*(1.0+e*b[i,j]))**2\n            expII=np.zeros((n,))\n            expJJ=np.zeros((m,))\n            for k in range(n):\n                II=(i-vI[k])*dx\n                expII[k]=math.exp(-II*II/denom)\n\n            for k in range(m):\n                JJ=(j-vJ[k])*dx\n                expJJ[k]=math.exp(-JJ*JJ/denom)\n            gb[i,j]=norm*(expII.sum()*expJJ.sum())\n    return gb\n

\n\n

现在：

\n\n

>>> print(parallel_nb_calc_ead(a,0.1,1.0,0.5)[1,1])\n15.9160709993\n>>> %timeit -n1 -r1 parallel_nb_calc_ead(a,0.1,1.0,0.5)\n349 ms \xc2\xb1 0 ns per loop (mean \xc2\xb1 std. dev. of 1 run, 1 loop each)\n

\n\n

意味着几乎另一个因素 2（我的机器只有两个 CPU，根据硬件，加速可能会更大）。顺便说一句，我们的速度比原始版本快了近 200 倍。

\n\n

我打赌可以改进上面的代码，但我不会去那里。

\n\n\n\n

列出与之比较的当前版本calc_ead：

\n\n

import numpy as np\nfrom numba import njit, double\n\ndef calc_gb_gauss_2d(b,s0,e,dx):\n    n,m=b.shape\n    norm = 1.0/(2*np.pi*s0**2)\n    gb = np.zeros((n,m))\n    for i in range(n):\n        for j in range(m):\n            for ii in range(n):\n                for jj in range(m):\n                    gb[i,j]+=np.exp(-(((i-ii)*dx)**2+((j-jj)*dx)**2)/(2.0*(s0*(1.0+e*b[i,j]))**2))\n            gb[i,j]*=norm\n    return gb\n\ncalc_gb_gauss_2d_nb = njit(double[:, :](double[:, :],double,double,double))(calc_gb_gauss_2d)\n

\n