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

Question

我试图通过首先插入数据来制作表面来进行双积分.我正在使用numba试图加快这个过程,但这只是花了太长时间.

这是我的代码,其中包含运行位于此处和此处的代码所需的图像.

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,因此最内层语句需要执行十六万亿次.这在使用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秒中,在rbs函数中花费0.68秒,即scipy.interpolate.RectBivariateSpline.可能还有进一步优化的空间 - 这里有一些想法:

1. 重新启用Numba.我的系统上没有它.在这一点上它可能没什么区别,但你很容易测试.
2. 进行更多特定于域的优化,例如尝试简化正在进行的基本计算.我的优化是无损的,我不知道你的问题域,所以我无法尽可能深入地进行优化.
3. 尝试向量化剩余的循环.除非您愿意用支持每次调用多次计算的东西替换scipy RBS函数,否则这可能很难.
4. 获得更快的CPU.我很慢; 只需使用比我的小笔记本电脑更好的电脑,你就可以获得至少2倍的加速.
5. 下采样数据.您的测试图像是2000x2000像素,但包含的细节相当少.如果你将线性尺寸减少2-10倍,你将获得巨大的加速.

所以这就是我现在的情况.这会让你离开？假设计算机稍好一些并且没有进一步的优化工作,即使优化的代码也需要大约一个月的时间来处理您的测试图像.如果你只需要做一次,也许那很好.如果你需要更频繁地完成它,或者需要在尝试不同的事情时迭代代码,你可能需要继续优化 - 从那个消耗超过一半时间的RBS函数开始.

特别提示:如果它不具有几乎相同的变量名称,如您的代码会更容易处理大量k和K,也没有使用j作为变量名,也可以作为一个复杂的数字后缀(0j).