python - numpy.einsum 大大加快了计算速度 - 但 numpy.einsum_path 没有显示加速，我错过了什么？

Question

我有一个奇怪的情况，我可以看到 numpy.einsum 加速计算，但在 einsum_path 中看不到相同的结果。我想量化/解释这种可能的加速，但在某处遗漏了一些东西......

简而言之，我有一个矩阵乘法，其中只需要最终产品的对角线。

a = np.arange(9).reshape(3,3)
print('input array')
print(a)
print('normal method')
print(np.diag(a.dot(a)))
print('einsum method')
print(np.einsum('ij,ji->i', a, a))

产生输出:

input array
[[0 1 2]
 [3 4 5]
 [6 7 8]]
normal method
[ 15  54 111]
einsum method
[ 15  54 111]

在大型矩阵上运行时，numpy.einsum 速度要快得多。

A = np.random.randn(2000, 300)
B = np.random.randn(300, 2000)
print('normal method')
%timeit np.diag(A.dot(B))
print('einsum method')
%timeit np.einsum('ij,ji->i', A, B)

产生:

normal method
17.2 ms ± 131 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
einsum method
1.02 ms ± 7.82 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

我的直觉是，这种加速是可能的，因为 numpy.einsum 能够放弃最终通过采用对角线而放弃的计算 - 但是，如果我正确地阅读它，输出numpy.einsum_path 根本没有显示任何加速。

print(np.einsum_path('ij,ji->i',A,B,optimize=True)[1])
  Complete contraction:  ij,ji->i
         Naive scaling:  2
     Optimized scaling:  2
      Naive FLOP count:  1.200e+06
  Optimized FLOP count:  1.200e+06
   Theoretical speedup:  1.000
  Largest intermediate:  2.000e+03 elements
--------------------------------------------------------------------------
scaling                  current                                remaining
--------------------------------------------------------------------------
   2                    ji,ij->i                                     i->i

问题:

1. 为什么我可以看到实际的加速，但没有反射(reflect)在计算路径中？
2. 有没有办法量化 numpy.einsum 中的 ij,ji->i 路径的加速？

Answer 1

该路径仅在处理两个以上参数时查看替代顺序。只有两个参数，分析就没有任何作用。您的诊断(点)

In [113]: np.diag(a.dot(a))
Out[113]: array([ 15,  54, 111])

使用einsum的等效项是:

In [115]: np.einsum('ii->i',np.einsum('ij,jk->ik',a,a))
Out[115]: array([ 15,  54, 111])

但是我们可以跳过中间步骤:

In [116]: np.einsum('ij,ji->i',a,a)
Out[116]: array([ 15,  54, 111])

索引表示法足够灵活，不需要进行完整的点计算。

获得相同结果的另一种方法是:

In [117]: (a*a.T).sum(axis=1)
Out[117]: array([ 15,  54, 111])

使用matmul，我们可以在没有diag的情况下进行计算，将第一个维度视为“批处理”。但首先需要进行一些 reshape :

In [121]: a[:,None,:]@a.T[:,:,None]
Out[121]: 
array([[[ 15]],

       [[ 54]],

       [[111]]])

In [122]: np.squeeze(a[:,None,:]@a.T[:,:,None])
Out[122]: array([ 15,  54, 111])

我的时光

normal method
135 ms ± 3.87 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
einsum method
3.09 ms ± 46.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [130]: timeit (A*B.T).sum(axis=0)
8.06 ms ± 78.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [131]: timeit np.squeeze(A[:,None,:]@B.T[:,:,None])
3.52 ms ± 195 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

dot 创建一个 (2000,2000) 数组，然后提取 2000 个元素。

就元素乘法而言，点是:

In [136]: (a[:,:,None]*a[None,:,:]).sum(axis=1)
Out[136]: 
array([[ 15,  18,  21],
       [ 42,  54,  66],
       [ 69,  90, 111]])

对于 A 和 B，中间乘积将为 (2000,300,2000)，其总和为 (2000,2000)。 einsum(有效地)进行了 2000 次计算，大小 (300,300) 减少到 (1,)。

einsum 比 diag/dot 更接近此计算，将大小 2000 的维度视为 1d dot 的“批处理”计算:

In [140]: timeit np.array([A[i,:].dot(B[:,i]) for i in range(2000)])
9.46 ms ± 270 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)