Numpy Einsum 路径差异和优化参数

Question

我有以下张量执行，

np.einsum('k,pjqk,yzjqk,yzk,ipqt->it', A, B, C, D, E)

而且我注意到，当 'z' 或 'q' 在维度上扩展时，执行时间确实受到影响，尽管我的直觉是它可能不应该那么糟糕 - 也许这是我的输入形式我可以通过手动张量收缩来优化。

经过一点挖掘，我发现优化有两种模式：'optimal' 和 'greedy'。如果我分别针对两种模式评估我的路径：

(['einsum_path', (0, 3), (1, 3), (0, 2), (0, 1)],
'  Complete contraction:  k,pjqk,yzjqk,yzk,ipqt->it\n'
'         Naive scaling:  8\n'
'     Optimized scaling:  5\n'
'      Naive FLOP count:  5.530e+04\n'
'  Optimized FLOP count:  7.930e+02\n'
'   Theoretical speedup:  69.730\n'
'  Largest intermediate:  2.400e+01 elements\n'
'--------------------------------------------------------------------------\n'
'scaling                  current                                remaining\n'
'--------------------------------------------------------------------------\n'
'   3                  yzk,k->yzk                  pjqk,yzjqk,ipqt,yzk->it\n'
'   5              yzk,yzjqk->jqk                        pjqk,ipqt,jqk->it\n'
'   4                jqk,pjqk->qp                              ipqt,qp->it\n'
'   4                 qp,ipqt->it                                   it->it')

和

(['einsum_path', (2, 3), (1, 3), (1, 2), (0, 1)],
'  Complete contraction:  k,pjqk,yzjqk,yzk,ipqt->it\n'
'         Naive scaling:  8\n'
'     Optimized scaling:  5\n'
'      Naive FLOP count:  5.530e+04\n'
'  Optimized FLOP count:  1.729e+03\n'
'   Theoretical speedup:  31.981\n'
'  Largest intermediate:  4.800e+01 elements\n'
'--------------------------------------------------------------------------\n'
'scaling                  current                                remaining\n'
'--------------------------------------------------------------------------\n'
'   5              yzk,yzjqk->jqk                      k,pjqk,ipqt,jqk->it\n'
'   4               jqk,pjqk->qkp                           k,ipqt,qkp->it\n'
'   5               qkp,ipqt->tik                                k,tik->it\n'
'   3                   tik,k->it                                   it->it')

测试结果是 'optimal' 对我来说要快得多，如图所示。

问题

谁能简单地解释一下区别是什么以及为什么 'greedy' 设置为默认值？

总是使用 'optimal' 的缺点是什么？

如果我的 einsum 计算将是运行 1000 次（这是优化迭代的一部分），我是否应该重组执行以自动受益于 'optimal' 路径而无需每次都重新计算它（或 'greedy' 路径）？

Answer 1

阅读更多内容后，任何发现此内容的人都会发现以下内容：

'greedy' 通常在大多数使用情况下产生 'optimal' 解决方案并且执行速度更快。对于可能无意中在迭代循环中使用 einsum 的一般用户，将 'greedy' 保留为默认值可能就足够了。

否则对于 one-off 计算，似乎 'optimal' 的最小额外开销意味着它可以有效地使用，除了，也许，对于大量索引，它可能会提供很大的提升（就我而言）。

在循环中最好的做法是 pre-calculate 它（或在第一次迭代中计算并更新非局部变量）并将其作为参数提供：

path, display = np.einsum_path('k,pjqk,yzjqk,yzk,ipqt->it', A, B, C, D, E, optimize='optimal')
for i in range(BIG_INT):
    # other things
    calculation = np.einsum_path('k,pjqk,yzjqk,yzk,ipqt->it', A, B, C, D, E, optimize=path)
    # more things

Numpy Einsum 路径差异和优化参数

Numpy Einsum Path Differences and the Optimize Argument

numpy

numpy-einsum

问题