<p>我不确定您是否希望它都是-<a href="/questions/tagged/numpy" class="post-tag" title="show questions tagged 'numpy'" rel="tag">numpy</a>,但我一直使用<a href="/questions/tagged/numba" class="post-tag" title="show questions tagged 'numba'" rel="tag">numba</a>来表示缓慢但易于实现基于循环的函数。循环密集型任务的加速速度惊人。首先,我只是把你的<code>all_loopy</code>变体,它已经给了我有竞争力的结果:</p>
<pre><code>m,n,N = 20,20,20
a = np.random.rand(m,n,N)
b = np.random.rand(n,m,N)
%timeit numba_all_loopy(a,b)
1000 loops, best of 3: 476 µs per loop # 3 times faster than everything else
%timeit tensordot_twoloop(a,b)
100 loops, best of 3: 16.1 ms per loop
%timeit einsum_twoloop(a,b)
100 loops, best of 3: 4.02 ms per loop
%timeit einsum_oneloop(a,b)
1000 loops, best of 3: 1.52 ms per loop
%timeit fully_vectorized(a,b)
1000 loops, best of 3: 1.67 ms per loop
</code></pre>
<p>然后我用你的100,100,100个案子来测试它:</p>
^{pr2}$
<p>除了注意到我的电脑比你的电脑慢得多之外,numba版本的电脑也越来越慢。怎么了?在</p>
<p>Numpy使用编译版本,根据编译器选项,Numpy可能会优化循环,而numba则不会。假设C-连续数组,最里面的循环应该是数组的最后一个轴。它是变化最快的轴,所以缓存位置会更好。在</p>
<pre><code>@nb.njit
def numba_all_loopy2(a,b):
P,Q,N = a.shape
d = np.zeros(N)
# First axis a, second axis b
for k in range(P):
# first axis b, second axis a
for n in range(Q):
# third axis a
for i in range(N):
# third axis b
A = a[k,n,i] # so we have less lookups of the same variable
for j in range(i):
d[i] += A * b[n,k,j]
return d
</code></pre>
<p>那么这个“优化”的numba函数的时间安排是什么呢?它能和其他人相比甚至打败他们吗?在</p>
<pre><code>m = n = N = 20
%timeit numba_all_loopy(a,b)
1000 loops, best of 3: 476 µs per loop
%timeit numba_all_loopy2(a,b)
1000 loops, best of 3: 379 µs per loop # New one is a bit faster
</code></pre>
<p>所以对于小矩阵来说,速度稍微快一点,大矩阵呢:</p>
<pre><code>m = n = N = 100
%timeit numba_all_loopy(a,b)
1 loop, best of 3: 2.34 s per loop
%timeit numba_all_loopy2(a,b)
1 loop, best of 3: 213 ms per loop # More then ten times faster now!
</code></pre>
<p>所以我们对大型阵列有一个惊人的加速。这个函数现在比矢量化方法快4-5倍,结果是一样的。在</p>
<p>但令人惊讶的是,排序似乎似乎以某种方式依赖于计算机,因为<code>fully_vectorized</code>最快,而{<cd4>}-选项在@Divakar的机器上更快。如果这些结果真的很快就公开了。在</p>
<p>为了好玩,我尝试了<code>n=m=N=200</code>:</p>
<pre><code>%timeit numba_all_loopy2(a,b)
1 loop, best of 3: 3.38 s per loop # still 5 times faster
%timeit einsum_oneloop(a,b)
1 loop, best of 3: 51.4 s per loop
%timeit fully_vectorized(a,b)
1 loop, best of 3: 16.7 s per loop
</code></pre>