python - 需要帮助矢量化代码或优化

Question

我试图通过首先对数据进行插值来制作曲面来进行二重积分。我正在使用 numba 来尝试加快这个过程，但它花费的时间太长了。

Here is my code,运行代码所需的图像位于 here和 here .

Answer 1

注意到您的代码有一组四重嵌套的 for 循环，我专注于优化内部对。这是旧代码:

for i in xrange(K.shape[0]):
    for j in xrange(K.shape[1]):

        print(i,j)
        '''create an r vector '''
        r=(i*distX,j*distY,z)

        for x in xrange(img.shape[0]):
            for y in xrange(img.shape[1]):
                '''create an ksi vector, then calculate
                   it's norm, and the dot product of r and ksi'''
                ksi=(x*distX,y*distY,z)
                ksiNorm=np.linalg.norm(ksi)
                ksiDotR=float(np.dot(ksi,r))

                '''calculate the integrand'''
                temp[x,y]=img[x,y]*np.exp(1j*k*ksiDotR/ksiNorm)

        '''interpolate so that we can do the integral and take the integral'''
        temp2=rbs(a,b,temp.real)
        K[i,j]=temp2.integral(0,n,0,m)

由于 K 和 img 都是大约 2000x2000，所以最里面的语句需要执行 16 万亿次。这在使用 Python 时根本不切实际，但我们可以使用 NumPy 将工作转移到 C 和/或 Fortran 中进行矢量化。我一次小心地执行了这一步骤，以确保结果匹配；这是我最终得到的:

'''create all r vectors'''
R = np.empty((K.shape[0], K.shape[1], 3))
R[:,:,0] = np.repeat(np.arange(K.shape[0]), K.shape[1]).reshape(K.shape) * distX
R[:,:,1] = np.arange(K.shape[1]) * distY
R[:,:,2] = z

'''create all ksi vectors'''
KSI = np.empty((img.shape[0], img.shape[1], 3))
KSI[:,:,0] = np.repeat(np.arange(img.shape[0]), img.shape[1]).reshape(img.shape) * distX
KSI[:,:,1] = np.arange(img.shape[1]) * distY
KSI[:,:,2] = z

# vectorized 2-norm; see http://stackoverflow.com/a/7741976/4323                                                    
KSInorm = np.sum(np.abs(KSI)**2,axis=-1)**(1./2)

# loop over entire K, which is same shape as img, rows first                                                        
# this loop populates K, one pixel at a time (so can be parallelized)                                               
for i in xrange(K.shape[0]):                                                                                    
    for j in xrange(K.shape[1]):                                                                                

        print(i, j)

        KSIdotR = np.dot(KSI, R[i,j])
        temp = img * np.exp(1j * k * KSIdotR / KSInorm)

        '''interpolate so that we can do the integral and take the integral'''
        temp2 = rbs(a, b, temp.real)
        K[i,j] = temp2.integral(0, n, 0, m)

内部的一对循环现在完全消失了，取而代之的是预先完成的向量化操作(空间成本与输入大小成线性关系)。

在我的 Macbook Air 1.6 GHz i5 上，这将外部两个循环的每次迭代时间从 340 秒减少到 1.3 秒，而无需使用 Numba。在每次迭代的 1.3 秒中，0.68 秒用于 rbs 函数，即 scipy.interpolate.RectBivariateSpline。可能还有进一步优化的空间——这里有一些想法:

1. 重新启用 Numba。我的系统上没有它。在这一点上可能没有太大区别，但您可以轻松测试。
2. 进行更多针对特定领域的优化，例如尝试简化正在进行的基础计算。我的优化旨在实现无损，并且我不知道您的问题领域，因此我无法像您那样进行深度优化。
3. 尝试向量化剩余的循环。这可能很难，除非您愿意用支持每次调用多次计算的函数替换 scipy RBS 函数。
4. 获得更快的 CPU。我的速度很慢；只需使用比我的小型笔记本电脑更好的计算机，您就可能获得至少 2 倍的加速。
5. 对数据进行缩减采样。您的测试图像为 2000x2000 像素，但包含的细节相当少。如果将它们的线性尺寸减少 2-10 倍，您将获得巨大的加速。

我现在就到此为止。这会给你留下什么？假设一台稍微好一点的计算机并且没有进一步的优化工作，即使是优化的代码也需要大约一个月的时间来处理您的测试图像。如果您只需要执行一次，也许就可以了。如果您需要更频繁地执行此操作，或者需要在尝试不同的事情时迭代代码，您可能需要不断优化——从现在消耗一半以上时间的 RBS 函数开始。

额外提示:如果您的代码没有像 k 和 K 这样几乎相同的变量名，也不使用 j 作为变量名，也作为复数后缀 (0j)。