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

时间:2013-07-08 14:33:54

标签: python optimization numba

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

Here is my code,包含运行位于herehere的代码所需的图片。

1 个答案:

答案 0 :(得分:8)

注意到你的代码有一套四重嵌套的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倍,你将获得巨大的加速。
  6. 这就是我现在的意思。这会让你离开?假设计算机稍微好一些并且没有进一步的优化工作,即使优化的代码也需要大约一个月的时间来处理您的测试图像。如果你只需要做一次,也许那很好。如果您需要更频繁地执行此操作,或者需要在尝试不同的事情时迭代代码,您可能需要继续优化 - 从那个消耗超过一半时间的RBS函数开始。

    额外提示:如果代码没有几乎相同的变量名称(例如kK),并且使用j作为变量,那么您的代码将更容易处理名称以及复数后缀(0j)。