回答此问题可获得 20 贡献值,回答如果被采纳可获得 50 分。
<p>我正在尝试优化以下循环:</p>
<pre><code>def numpy(nx, nz, c, rho):
for ix in range(2, nx-3):
for iz in range(2, nz-3):
a[ix, iz] = sum(c*rho[ix-1:ix+3, iz])
b[ix, iz] = sum(c*rho[ix-2:ix+2, iz])
return a, b
</code></pre>
<p>我尝试了不同的解决方案,发现使用numba计算产品总和可以获得更好的性能:</p>
^{pr2}$
<p>这导致</p>
<blockquote>
<p>Time numpy : 4.1595</p>
<p>Time numba1 : 0.6993</p>
<p>Time numba2 : 1.0135</p>
</blockquote>
<p>使用numba版本的sum函数(sum\u opt)性能非常好。但是我想知道为什么numba版本的双循环函数(numba2)会导致执行速度变慢。我试图使用jit而不是autojit,指定参数类型,但情况更糟。在</p>
<p>我还注意到,先在最小的循环上循环比先在最大的循环上循环慢。有什么解释吗?在</p>
<p>不管怎样,我确信这个双循环函数可以改进很多向量化问题(比如<a href="https://stackoverflow.com/questions/8299891/vectorization-of-this-numpy-double-loop">this</a>)或使用其他方法(map?)但我对这些方法有点困惑。在</p>
<p>在我代码的其他部分,我使用numba和numpy切片方法来替换所有显式循环,但是在这个特殊的例子中,我不知道如何设置它。在</p>
<p>有什么想法吗?在</p>
<p><strong>编辑</strong></p>
<p>谢谢你的评论。我在这个问题上做了一些工作:</p>
<pre><code>import numba as nb
import numpy as np
from scipy import signal
import time
@nb.jit(['float64(float64[:], float64[:])'], nopython=True)
def sum_opt(arr1, arr2):
s = arr1[0]*arr2[0]
for i in xrange(1, len(arr1)):
s+=arr1[i]*arr2[i]
return s
@nb.autojit
def numba1(nx, nz, c, rho, a, b):
for ix in range(2, nx-3):
for iz in range(2, nz-3):
a[ix, iz] = sum_opt(c, rho[ix-1:ix+3, iz])
b[ix, iz] = sum_opt(c, rho[ix-2:ix+2, iz])
return a, b
@nb.jit(nopython=True)
def numba2(nx, nz, c, rho, a, b):
for ix in range(2, nx-3):
for iz in range(2, nz-3):
a[ix, iz] = sum_opt(c, rho[ix-1:ix+3, iz])
b[ix, iz] = sum_opt(c, rho[ix-2:ix+2, iz])
return a, b
@nb.jit(['float64[:,:](int16, int16, float64[:], float64[:,:], float64[:,:])'], nopython=True)
def numba3a(nx, nz, c, rho, a):
for ix in range(2, nx-3):
for iz in range(2, nz-3):
a[ix, iz] = sum_opt(c, rho[ix-1:ix+3, iz])
return a
@nb.jit(['float64[:,:](int16, int16, float64[:], float64[:,:], float64[:,:])'], nopython=True)
def numba3b(nx, nz, c, rho, b):
for ix in range(2, nx-3):
for iz in range(2, nz-3):
b[ix, iz] = sum_opt(c, rho[ix-2:ix+2, iz])
return b
def convol(nx, nz, c, aa, bb):
s1 = rho[1:nx-1,2:nz-3]
s2 = rho[0:nx-2,2:nz-3]
kernel = c[:,None][::-1]
aa[2:nx-3,2:nz-3] = signal.convolve2d(s1, kernel, boundary='symm', mode='valid')
bb[2:nx-3,2:nz-3] = signal.convolve2d(s2, kernel, boundary='symm', mode='valid')
return aa, bb
nx = 1024
nz = 256
rho = np.random.rand(nx, nz)
c = np.random.rand(4)
a = np.zeros((nx, nz))
b = np.zeros((nx, nz))
ti = time.clock()
for i in range(1000):
a, b = numba1(nx, nz, c, rho, a, b)
print 'Time numba1 : ' + `round(time.clock() - ti, 4)`
ti = time.clock()
for i in range(1000):
a, b = numba2(nx, nz, c, rho, a, b)
print 'Time numba2 : ' + `round(time.clock() - ti, 4)`
ti = time.clock()
for i in range(1000):
a = numba3a(nx, nz, c, rho, a)
b = numba3b(nx, nz, c, rho, b)
print 'Time numba3 : ' + `round(time.clock() - ti, 4)`
ti = time.clock()
for i in range(1000):
a, b = convol(nx, nz, c, a, b)
print 'Time convol : ' + `round(time.clock() - ti, 4)`
</code></pre>
<p>您的解决方案非常优雅,但我必须在代码中大量使用此函数。因此,对于1000次迭代,这将导致</p>
<blockquote>
<p>Time numba1 : 3.2487</p>
<p>Time numba2 : 3.7012</p>
<p>Time numba3 : 3.2088</p>
<p>Time convol : 22.7696</p>
</blockquote>
<p><code>autojit</code>和{<cd2>}非常接近。
但是,在使用<code>jit</code>时,指定所有参数类型似乎很重要。在</p>
<p>当函数有多个输出时,我不知道是否有方法在<code>jit</code>修饰符中指定参数类型。有人吗?在</p>
<p>现在我没有找到其他的解决办法,除了使用numba。欢迎有新想法!在</p>