Question

有些时候回复this question（现已删除，但10K +代表用户仍然可以查看）已发布。它对我来说很有趣，我在试图解决它时学到了一些新东西，我认为值得分享。我想发布这些想法/解决方案，并希望看到人们发布其他可能的解决方法。我正在发布问题的要点。

因此，我们有两个NumPy ndarrays a和b个形状：

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

我们假设我们正在处理m，n＆amp; N具有可比性。

问题是解决以下乘法和求和，重点关注性能：

def all_loopy(a,b):
    P,Q,N = a.shape
    d = np.zeros(N)
    for i in range(N):
        for j in range(i):
            for k in range(P):
                for n in range(Q):
                    d[i] += a[k,n,i] * b[n,k,j]
    return d

Answer 1

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

1）首先，"for j in range(i)"处的迭代器存在依赖关系。根据我之前的经验，特别是尝试在MATLAB上解决此类问题时，似乎可以使用lower triangular matrix来处理此类依赖关系，因此np.tril应该在那里工作。因此，一个完全向量化的解决方案而不是内存有效的解决方案（因为它在最终减少到(N,N)形状的数组之前创建一个中间(N,)形状的数组）将是 -

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

2）下一个技巧/想法是为i中的迭代器for i in range(N)保留一个循环，但是使用索引插入该依赖项并使用np.einsum执行所有这些乘法和求和。优点是内存效率。实现看起来像这样 -

def einsum_oneloop(a,b):
    d = np.zeros(N)
    for i in range(N):
        d[i] = np.einsum('ij,jik->',a[:,:,i],b[:,:,np.arange(i)])
    return d

有两种更明显的方法可以解决它。因此，如果我们从原始的all_loopy解决方案开始工作，可以保留外部的两个循环并使用np.einsum或np.tensordot来执行这些操作，从而删除内部的两个循环，就像这样 -

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（因为非常缓慢而退出all_loopy）：

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看起来对我很好，而fully_vectorized可以在处理小到大小的数组时使用！

Answer 2

我不确定你是不是希望它全部 - numpy但我总是使用numba来缓慢但很容易实现基于循环的功能。循环密集型任务的加速是惊人的。首先，我numba.njit引导您的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案例进行了测试：

m,n,N = 100,100,100
a = np.random.rand(m,n,N)
b = np.random.rand(n,m,N)

%timeit numba_all_loopy(a,b)
1 loop, best of 3: 2.35 s per loop
%timeit tensordot_twoloop(a,b)
1 loop, best of 3: 3.54 s per loop
%timeit einsum_twoloop(a,b)
1 loop, best of 3: 2.58 s per loop
%timeit einsum_oneloop(a,b)
1 loop, best of 3: 2.71 s per loop
%timeit fully_vectorized(a,b)
1 loop, best of 3: 1.08 s per loop

除了注意到我的电脑比你的电脑慢得多 - 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是最快的，einsum - 选项在@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

矩阵乘法与迭代器依赖 - NumPy

2 个答案: