长期的难题,如何在python中优化多级循环?

发布于 2021-01-29 15:00:53

我已经在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

我删除了一些常量以使其看起来更简单。

有谁能帮助我优化此脚本以提高效率?

关注者
0
被浏览
74
1 个回答
  • 面试哥
    面试哥 2021-01-29
    为面试而生,有面试问题,就找面试哥。

    只需编译

    为了获得最佳性能,我建议使用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
    

    由于“如果”,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倍。



知识点
面圈网VIP题库

面圈网VIP题库全新上线,海量真题题库资源。 90大类考试,超10万份考试真题开放下载啦

去下载看看