迭代依赖数矩阵乘法

Question

有时有人发布了back^{}（现在已删除，但10K+用户仍可以查看）。我觉得这很有趣，我在那里学到了一些新的东西，同时试图解决它，我认为这值得分享。我想发表这些想法/解决方案，并希望看到人们张贴其他可能的方法来解决它。下一步我将把问题的要旨贴出来。在

所以，我们有两个数量庞大的a和{}的形状：

a : (m,n,N)
b : (n,m,N)

假设我们处理的是m，n&；N是可比的。在

问题是解决以下乘法和和运算，重点是性能：

^{pr2}$

Answer 1

在寻找矢量化和更快的解决方法的过程中，我学到了一些东西。在

1）首先，在"for j in range(i)"处存在迭代器的依赖关系。根据我以前的经验，尤其是在试图解决MATLAB上的问题时，似乎可以用^{}来处理这种依赖关系，因此{a2}应该可以在那里工作。因此，一个完全矢量化的解决方案和内存效率不高的解决方案（因为它在最终还原为(N,)形状的数组之前创建了一个中间的(N,N)形数组）将是-

def fully_vectorized(a,b):
    return np.tril(np.einsum('ijk,jil->kl',a,b),-1).sum(1)

2）下一个技巧/想法是在for i in range(N)中为迭代器i保留一个循环，但使用索引插入该依赖关系，并使用^{}执行所有这些乘法和求和。这样做的好处是内存效率。执行应该是这样的-

^{pr2}$

有两种更明显的解决方法。如果我们可以用两个外循环来执行两个循环

def tensordot_twoloop(a,b):
    d = np.zeros(N)
    for i in range(N):
        for j in range(i):
            d[i] += np.tensordot(a[:,:,i],b[:,:,j], axes=([1,0],[0,1]))        
    return d

def einsum_twoloop(a,b):
    d = np.zeros(N)
    for i in range(N):
        for j in range(i):
            d[i] += np.einsum('ij,ji->',a[:,:,i],b[:,:,j])
    return d

运行时测试

让我们比较一下迄今为止发布的五种解决问题的方法，包括在问题中发布的方法。在

案例1：

In [26]: # Input arrays with random elements
    ...: m,n,N = 20,20,20
    ...: a = np.random.rand(m,n,N)
    ...: b = np.random.rand(n,m,N)
    ...: 

In [27]: %timeit all_loopy(a,b)
    ...: %timeit tensordot_twoloop(a,b)
    ...: %timeit einsum_twoloop(a,b)
    ...: %timeit einsum_oneloop(a,b)
    ...: %timeit fully_vectorized(a,b)
    ...: 
10 loops, best of 3: 79.6 ms per loop
100 loops, best of 3: 4.97 ms per loop
1000 loops, best of 3: 1.66 ms per loop
1000 loops, best of 3: 585 µs per loop
1000 loops, best of 3: 684 µs per loop

案例2：

In [28]: # Input arrays with random elements
    ...: m,n,N = 50,50,50
    ...: a = np.random.rand(m,n,N)
    ...: b = np.random.rand(n,m,N)
    ...: 

In [29]: %timeit all_loopy(a,b)
    ...: %timeit tensordot_twoloop(a,b)
    ...: %timeit einsum_twoloop(a,b)
    ...: %timeit einsum_oneloop(a,b)
    ...: %timeit fully_vectorized(a,b)
    ...: 
1 loops, best of 3: 3.1 s per loop
10 loops, best of 3: 54.1 ms per loop
10 loops, best of 3: 26.2 ms per loop
10 loops, best of 3: 27 ms per loop
10 loops, best of 3: 23.3 ms per loop

案例3（由于速度太慢而忽略所有的漏洞）：

In [30]: # Input arrays with random elements
    ...: m,n,N = 100,100,100
    ...: a = np.random.rand(m,n,N)
    ...: b = np.random.rand(n,m,N)
    ...: 

In [31]: %timeit tensordot_twoloop(a,b)
    ...: %timeit einsum_twoloop(a,b)
    ...: %timeit einsum_oneloop(a,b)
    ...: %timeit fully_vectorized(a,b)
    ...: 
1 loops, best of 3: 1.08 s per loop
1 loops, best of 3: 744 ms per loop
1 loops, best of 3: 568 ms per loop
1 loops, best of 3: 866 ms per loop

从数字上看，einsum_oneloop在我看来相当不错，而{}可以在处理小到中等大小的数组时使用！在

Answer 2

我不确定您是否希望它都是-numpy，但我一直使用numba来表示缓慢但易于实现基于循环的函数。循环密集型任务的加速速度惊人。首先，我只是把你的all_loopy变体，它已经给了我有竞争力的结果：

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

然后我用你的100，100，100个案子来测试它：

^{pr2}$

除了注意到我的电脑比你的电脑慢得多之外，numba版本的电脑也越来越慢。怎么了？在

Numpy使用编译版本，根据编译器选项，Numpy可能会优化循环，而numba则不会。假设C-连续数组，最里面的循环应该是数组的最后一个轴。它是变化最快的轴，所以缓存位置会更好。在

@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

那么这个“优化”的numba函数的时间安排是什么呢？它能和其他人相比甚至打败他们吗？在

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

所以对于小矩阵来说，速度稍微快一点，大矩阵呢：

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!

所以我们对大型阵列有一个惊人的加速。这个函数现在比矢量化方法快4-5倍，结果是一样的。在

但令人惊讶的是，排序似乎似乎以某种方式依赖于计算机，因为fully_vectorized最快，而{}-选项在@Divakar的机器上更快。如果这些结果真的很快就公开了。在

为了好玩，我尝试了n=m=N=200：

%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