与 Fortran 相比，numpy 计算太慢了答案

【问题标题】：Numpy calculation too slow compared to Fortran与 Fortran 相比，numpy 计算太慢了
【发布时间】：2020-11-19 21:28:09
【问题描述】：

我以前主要使用 Fortran，但最近我开始使用 python 和 Numpy 进行深度学习应用。

但是，在 python 中计算双 for 循环比 Fortran 慢得多。我已经知道 Fortran 本来计算速度很快，但我想知道我的 python 代码是否有问题。这是我使用的代码如下：

for it in range(nt):
    if it%20 == 1:
        print(it,'//',nt)

    itimenum4 = "%.4i" %(it)
    
    ppsix2[:,:]=0.; ppsiz2[:,:]=0.
    apsix2[:,:]=0.; apsiz2[:,:]=0.
    ax[:,:]=0.; az[:,:]=0.
    p3[:,:]=0.
    
    for iz in range(2,nnz+1):
        for ix in range(2,nnx+1):

            pdx2=(p2[ix+1,iz]-p2[ix,iz])*a1+(p2[ix+2,iz]-p2[ix-1,iz])*a2
            pdz2=(p2[ix,iz+1]-p2[ix,iz])*a1+(p2[ix,iz+2]-p2[ix,iz-1])*a2

            dpml0=math.log(1./R)*3.*vp[ix,iz]/(2.*dx*pml)

            if ix <= pml:
                dpml=dpml0*(float(pml-ix+1)/float(pml))**2
                damp=math.exp(-dpml*dt)
                ppsix2[ix,iz]=damp*ppsix1[ix,iz]+(damp-1.)*pdx2

            if ix >nx+pml:
                dpml=dpml0*(float(ix-(nx+pml))/float(pml))**2
                damp=math.exp(-dpml*dt)
                ppsix2[ix,iz]=damp*ppsix1[ix,iz]+(damp-1.)*pdx2

            if iz > nz:
                dpml=dpml0*(float(iz-(nz))/float(pml))**2
                damp=math.exp(-dpml*dt)
                ppsiz2[ix,iz]=damp*ppsiz1[ix,iz]+(damp-1.)*pdz2

            ax[ix,iz]=pdx2+ppsix2[ix,iz]
            az[ix,iz]=pdz2+ppsiz2[ix,iz]
        
    az[:,1]=az[:,2]
    az[:,0]=az[:,3]

    for iz in range(2,nnz+1):
        for ix in range(2,nnx+1):
            adx=a1*(ax[ix,iz]-ax[ix-1,iz])+a2*(ax[ix+1,iz]-ax[ix-2,iz])
            adz=a1*(az[ix,iz]-az[ix,iz-1])+a2*(az[ix,iz+1]-az[ix,iz-2])

            dpml0=math.log(1./R)*3.*vp[ix,iz]/(2.*dx*pml)

            if ix <= pml:
                dpml=dpml0*(float(pml-ix+1)/float(pml))**2
                damp=math.exp(-dpml*dt)
                apsix2[ix,iz]=damp*apsix1[ix,iz]+(damp-1.)*adx

            if ix > nx+pml:
                dpml=dpml0*(float(ix-(nx+pml))/float(pml))**2
                damp=math.exp(-dpml*dt)
                apsix2[ix,iz]=damp*apsix1[ix,iz]+(damp-1.)*adx

            if iz > nz:
                dpml=dpml0*(float(iz-(nz))/float(pml))**2
                damp=math.exp(-dpml*dt)
                apsiz2[ix,iz]=damp*apsiz1[ix,iz]+(damp-1.)*adz

            px2=adx+apsix2[ix,iz]
            pz2=adz+apsiz2[ix,iz]

            p3[ix,iz]=2.*p2[ix,iz]-p1[ix,iz]+(dt*vp[ix,iz])**2*((px2+pz2)+srcf[ix,iz]*source[it])

            if iz == recz:
                mod_p[ix,it]=p3[ix,iz]

    p3[:,0:2]=0.

    p1[:,:]=p2[:,:]
    p2[:,:]=p3[:,:]
    ppsix1[:,:]=ppsix2[:,:]
    ppsiz1[:,:]=ppsiz2[:,:]
    apsix1[:,:]=apsix2[:,:]
    apsiz1[:,:]=apsiz2[:,:]

这是地球物理学中传统的一阶声波方程。无论如何，我使用了 Numpy 数组。是否有任何因素导致计算缓慢？

谢谢。

【问题讨论】：

在我对代码应用矢量化后，计算时间显着减少。感谢 cmets 的细心。

标签： python numpy

【解决方案1】：

几乎没有一项工作看起来像是矢量化的——例如，它正在做一些事情，比如提取单个元素 ppsize1[ix,iz] 以用于乘法。

Numpy 不会像 Fortran 那样快速编译成某种东西——它通过快速的例程来一次操作大块数据而变得很快。如果您没有明确编码算法的可向量化部分，那么 numpy 甚至会比 vanilla python 慢。

【讨论】：

【解决方案2】：

不幸的是，由于 python 是一种解释性语言（与已编译的 Fortran 相比），您在使用 python for 循环时将无法获得类似的速度。

有几种方法可以解决这个问题。一种方法是尝试使用 numpy 数组和 numpy 数学函数对你的 for 循环进行矢量化。通过这样做，您会将代码中计算量大的部分传递给预编译的 numpy 库。

另一种加快 Python 代码速度的方法是使用库来编译代码。 numba 就是一个这样的库，它提供了即时编译选项。

【讨论】：