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

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)。