np.einsum 4 个矩阵乘法的性能答案

【问题标题】：np.einsum performance of 4 matrix multiplicationsnp.einsum 4 个矩阵乘法的性能
【发布时间】：2018-07-23 09:03:45
【问题描述】：

给定以下 3 个矩阵：

M = np.arange(35 * 37 * 59).reshape([35, 37, 59])
A = np.arange(35 * 51 * 59).reshape([35, 51, 59])
B = np.arange(37 * 51 * 51 * 59).reshape([37, 51, 51, 59])
C = np.arange(59 * 27).reshape([59, 27])

我正在使用einsum 进行计算：

D1 = np.einsum('xyf,xtf,ytpf,fr->tpr', M, A, B, C, optimize=True);

但我发现它的性能要差得多：

tmp = np.einsum('xyf,xtf->tfy', A, M, optimize=True)
tmp = np.einsum('ytpf,yft->ftp', B, tmp, optimize=True)
D2 = np.einsum('fr,ftp->tpr', C, tmp, optimize=True)

我不明白为什么。
总的来说，我正在尽可能地优化这段代码。我已经阅读了有关 np.tensordot 函数的信息，但我似乎无法弄清楚如何将它用于给定的计算。

【问题讨论】：

看起来已经不错了。您可以在最后一步使用 tensordot 进行一些改进，但这不是瓶颈，因此不会对时间进行太多更改。
但是如何将einsum 拆分为 3 会产生更好的结果？那么瓶颈是什么？
嗯，瓶颈似乎是第二种方法（D2）的前两个步骤。 "But how come splitting einsum to 3 yields better results" - 不知道原因，但我在使用 einsum 时观察到，当使用更多输入时，einsum 会变慢，尤其是使用张量。如果我不得不猜测，我会认为它的内存拥塞，就像这里我们在 D1 方法中有 6 个循环。
@Divakar optimize=True 会将收缩拆分为多个较小的收缩以降低整体缩放比例。

标签： python numpy numpy-einsum

【解决方案1】：

您似乎偶然发现了greedy 路径提供非最佳缩放的情况。

>>> path, desc = np.einsum_path('xyf,xtf,ytpf,fr->tpr', M, A, B, C, optimize="greedy");
>>> print(desc)
  Complete contraction:  xyf,xtf,ytpf,fr->tpr
         Naive scaling:  6
     Optimized scaling:  5
      Naive FLOP count:  3.219e+10
  Optimized FLOP count:  4.165e+08
   Theoretical speedup:  77.299
  Largest intermediate:  5.371e+06 elements
--------------------------------------------------------------------------
scaling                  current                                remaining
--------------------------------------------------------------------------
   5              ytpf,xyf->xptf                         xtf,fr,xptf->tpr
   4               xptf,xtf->ptf                              fr,ptf->tpr
   4                 ptf,fr->tpr                                 tpr->tpr

>>> path, desc = np.einsum_path('xyf,xtf,ytpf,fr->tpr', M, A, B, C, optimize="optimal");
>>> print(desc)
  Complete contraction:  xyf,xtf,ytpf,fr->tpr
         Naive scaling:  6
     Optimized scaling:  4
      Naive FLOP count:  3.219e+10
  Optimized FLOP count:  2.744e+07
   Theoretical speedup:  1173.425
  Largest intermediate:  1.535e+05 elements
--------------------------------------------------------------------------
scaling                  current                                remaining
--------------------------------------------------------------------------
   4                xtf,xyf->ytf                         ytpf,fr,ytf->tpr
   4               ytf,ytpf->ptf                              fr,ptf->tpr
   4                 ptf,fr->tpr                                 tpr->tpr

使用np.einsum('xyf,xtf,ytpf,fr->tpr', M, A, B, C, optimize="optimal") 应该可以让您以最佳性能运行。我可以看看这个边缘，看看贪婪能不能抓住它。

【讨论】：

【解决方案2】：

虽然在这种情况下，一个贪心算法（有几个）确实可能找不到最优排序，但这与这里的谜题没有任何关系。当您执行 D2 方法时，您已经确定了操作的顺序，在这种情况下是 (((A,M),B),C) 或等效的 (((M,A),B),C)。这恰好是最佳路径。这 3 个 optimize=True 语句不需要并且被忽略，因为当有 2 个因素时没有使用优化。 D1 方法的减速是由于需要找到 4 数组操作的最优排序。如果您首先找到路径，然后使用 4 个数组将其传递给 einsum，使用 Optimize=path 我的猜测是这两种方法本质上是等效的。因此，减速是由于 D1 的优化步骤造成的。虽然我不确定如何找到最佳排序，但根据我所做的未发表的工作，这个任务通常会有 O(3^n) 最坏情况的行为，其中 n 是数组的数量。

【讨论】：