一个长期的难题,如何优化python中的多级循环？

Question

我在python中编写了一个函数来计算高斯展宽中的Delta函数,它涉及4级循环.但是,效率非常低,比以类似方式使用Fortran慢约10倍.

def Delta_Gaussf(Nw, N_bd, N_kp, hw, eigv):
    Delta_Gauss = np.zeros((Nw,N_kp,N_bd,N_bd),dtype=float)
    for w1 in range(Nw):
        for k1 in range(N_kp):
            for i1 in range(N_bd):
                for j1 in range(N_bd):
                    if ( j1 >= i1 ):
                        Delta_Gauss[w1][k1][i1][j1] = np.exp(pow((eigv[k1][j1]-eigv[k1][i1]-hw[w1])/width,2))
    return Delta_Gauss

我删除了一些常量,使它看起来更简单.

任何人都可以帮我优化这个脚本以提高效率吗？

Answer 1

只需编译它

为了获得最佳性能,我推荐使用Numba(使用方便,性能良好).或者Cython可能是一个好主意,但对代码进行了一些更改.

实际上你已经把一切都做对了,并且实现了一个易于理解(对于人类而言最重要的编译器)解决方案.

基本上有获得性能的方法

1. 像@scnerd所示,将代码矢量化.这通常比简单编译一个非常简单的代码更慢,更复杂,只使用一些for循环.不要向量化代码而不是使用编译器.从一个简单的循环方法来看,这通常是一些工作要做,并导致更慢和更复杂的结果.这个过程的优点是你只需要numpy,这是几乎每个处理一些数值计算的Python项目的标准依赖.

2. 编译代码.如果您已经有一个带有几个循环而没有其他循环的解决方案,或者只涉及一些非numpy函数,那么这通常是最简单和最快速的解决方案.

使用Numba的解决方案

你不必改变太多,我将pow函数更改为np.power并稍微更改了numpy中访问的数组的方式(这不是必需的).

import numba as nb
import numpy as np

#performance-debug info
import llvmlite.binding as llvm
llvm.set_option('', '--debug-only=loop-vectorize')

@nb.njit(fastmath=True)
def Delta_Gaussf_nb(Nw, N_bd, N_kp, hw, width,eigv):
    Delta_Gauss = np.zeros((Nw,N_kp,N_bd,N_bd),dtype=float)
    for w1 in range(Nw):
        for k1 in range(N_kp):
            for i1 in range(N_bd):
                for j1 in range(N_bd):
                    if ( j1 >= i1 ):
                        Delta_Gauss[w1,k1,i1,j1] = np.exp(np.power((eigv[k1,j1]-eigv[k1,i1]-hw[w1])/width,2))
    return Delta_Gauss

由于'if',SIMD矢量化失败.在下一步中我们可以删除它(可能需要在njited函数之外调用np.triu(Delta_Gauss)).我还并行化了这个功能.

@nb.njit(fastmath=True,parallel=True)
def Delta_Gaussf_1(Nw, N_bd, N_kp, hw, width,eigv):
    Delta_Gauss = np.zeros((Nw,N_kp,N_bd,N_bd),dtype=np.float64)
    for w1 in nb.prange(Nw):
        for k1 in range(N_kp):
            for i1 in range(N_bd):
                for j1 in range(N_bd):
                    Delta_Gauss[w1,k1,i1,j1] = np.exp(np.power((eigv[k1,j1]-eigv[k1,i1]-hw[w1])/width,2))
    return Delta_Gauss

性能

Nw = 20
N_bd = 20
N_kp = 20
width=20
hw = np.linspace(0., 1.0, Nw) 
eigv = np.zeros((N_kp, N_bd),dtype=np.float) 

Your version:           0.5s
first_compiled version: 1.37ms
parallel version:       0.55ms

这些简单的优化可以实现大约1000倍的加速.

Answer 2

BLUF：使用 Numpy 的全部功能，再加上另一个简洁的模块，您可以将 Python 代码比这个原始 for 循环代码快 100 倍以上。但是，使用@max9111 的答案，您可以通过更简洁的代码和更少的工作获得更快的速度。

生成的代码看起来与原始代码完全不同，因此我将一步一步地进行优化，以使过程和最终代码有意义。本质上，我们将使用大量广播来让 Numpy 在后台执行循环（这总是比在 Python 中循环快）。结果计算结果的完整平方，这意味着我们必须复制一些工作，因为结果是对称的，但以高性能方式完成这项工作比if在最深层次上更容易，老实说可能更快循环以避免计算。这在 Fortran 中可能可以避免，但在 Python 中可能无法避免。如果您希望结果与您提供的来源相同，我们需要取上三角形下面我的代码的结果（我在下面的示例代码中做的......triu在实际生产中随意删除调用，这是没有必要的）。

首先，我们会注意到一些事情。主方程的分母为np.sqrt，但该计算的内容在循环的任何迭代中都不会改变，因此我们将计算一次并重新使用结果。事实证明这是次要的，但无论如何我们都会这样做。接下来，内部两个循环的主要功能是执行eigv[k1][j1] - eigv[k1][i1]，这很容易矢量化。如果eigv是一个矩阵，则eigv[k1] - eigv[k1].T生成一个矩阵，其中result[i1, j1] = eigv[k1][j1] - eigv[k1][i1]。这允许我们完全删除最里面的两个循环：

def mine_Delta_Gaussf(Nw, N_bd, N_kp, hw, width, eigv):
    Delta_Gauss = np.zeros((Nw, N_kp, N_bd, N_bd), dtype=float)
    denom = np.sqrt(2.0 * np.pi) * width
    eigv = np.matrix(eigv)
    for w1 in range(Nw):
        for k1 in range(N_kp):
            this_eigv = (eigv[k1] - eigv[k1].T - hw[w1])
            v = np.power(this_eigv / width, 2)
            Delta_Gauss[w1, k1, :, :] = np.exp(-0.5 * v) / denom

    # Take the upper triangle to make the result exactly equal to the original code
    return np.triu(Delta_Gauss)

好吧，既然我们已经加入了广播的潮流，那么剩下的两个循环似乎真的可以用同样的方式删除了。碰巧，这很容易！我们唯一需要k1做的就是从eigv我们试图成对减去的行中取出行……那么为什么不同时对所有行执行此操作呢？我们目前基本上是为（其中is ）中的(1, B) - (B, 1)每一N行减去形状矩阵。我们可以通过减去形状矩阵（其中是）来滥用广播同时对所有行执行此操作：eigvBN_bdeigv(N, 1, B) - (N, B, 1)NN_kp

def mine_Delta_Gaussf(Nw, N_bd, N_kp, hw, width, eigv):
    Delta_Gauss = np.zeros((Nw, N_kp, N_bd, N_bd), dtype=float)
    denom = np.sqrt(2.0 * np.pi) * width
    for w1 in range(Nw):
        this_eigv = np.expand_dims(eigv, 1) - np.expand_dims(eigv, 2) - hw[w1]
        v = np.power(this_eigv / width, 2)
        Delta_Gauss[w1, :, :, :] = np.exp(-0.5 * v) / denom
    return np.triu(Delta_Gauss)

现在下一步应该很清楚了。我们只使用w1to index hw，所以让我们做更多的广播来numpy代替循环。我们目前正在从 shape 矩阵中减去一个标量值(N, B, B)，因此要获得 中每个W值的结果矩阵hw，我们需要对形状矩阵执行减法(1, N, B, B) - (W, 1, 1, 1)，numpy并将广播所有内容以生成形状矩阵(W, N, B, B)：

def Delta_Gaussf(hw, width, eigv):
    eigv_sub = np.expand_dims(eigv, 1) - np.expand_dims(eigv, 2)
    w_sub = np.expand_dims(eigv_sub, 0) - np.reshape(hw, (0, 1, 1, 1))
    v = np.power(w_sub / width, 2)
    denom = np.sqrt(2.0 * np.pi) * width
    Delta_Gauss = np.exp(-0.5 * v) / denom
    return np.triu(Delta_Gauss)

在我的示例数据中，此代码快了 ~100 倍（~900 毫秒到 ~10 毫秒）。您的里程可能会有所不同。

可是等等！还有更多！由于我们的代码都是数字/numpy/python，我们可以使用另一个方便的模块调用numba来将此函数编译为具有更高性能的等效函数。在幕后，它基本上是读取我们正在调用的函数并将函数转换为 C 类型和 C 调用以消除 Python 函数调用开销。它的作用不止于此，但这给出了我们将在何处受益的要点。在这种情况下，获得这种好处是微不足道的：

import numba

@numba.jit
def Delta_Gaussf(hw, width, eigv):
    eigv_sub = np.expand_dims(eigv, 1) - np.expand_dims(eigv, 2)
    w_sub = np.expand_dims(eigv_sub, 0) - np.reshape(hw, (0, 1, 1, 1))
    v = np.power(w_sub / width, 2)
    denom = np.sqrt(2.0 * np.pi) * width
    Delta_Gauss = np.exp(-0.5 * v) / denom
    return np.triu(Delta_Gauss)

仅通过添加该装饰器，结果函数在我的示例数据上从大约 10 毫秒下降到大约 7 毫秒。不费吹灰之力就不错了。

编辑：@max9111 给出了一个更好的答案，指出numba循环语法比numpy广播代码更有效。除了删除内部if语句之外几乎没有任何工作，他表明numba.jit可以更快地获得几乎原始的代码。结果更加清晰，因为您仍然只有一个最里面的等式来显示每个值是什么，并且您不必遵循上面使用的神奇广播。我强烈建议使用他的回答。

结论

对于我给定的示例数据（Nw = 20，N_bd = 20，N_kp = 20），我的最终运行时间如下（我在同一台计算机上为@max9111 的解决方案包含了计时，首先不使用并行执行，然后使用它我的 2 核 VM）：

Original code:               ~900 ms
Fortran estimate:            ~90 ms (based on OP saying it was ~10x faster)
Final numpy code:            ~10 ms
Final code with numba.jit:   ~7 ms
max9111's solution (serial): ~4ms
max9111 (parallel 2-core):   ~3ms

Overall vectorized speedup: ~130x
max9111's numba speedup: ~300x (potentially more with more cores)

我不知道您的 Fortran 代码到底有多快，但看起来正确使用 ofnumpy可以让您轻松地将其击败一个数量级，而@max9111 的numba解决方案可能会给您提供另一个数量级。