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

Question

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

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

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

假设我们正在处理m、n和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