如何避免使用张量或 einsum 的 for 循环？答案

【问题标题】：How to avoid using a for loop using either tensors or einsum?如何避免使用张量或 einsum 的 for 循环？
【发布时间】：2020-12-27 21:01:19
【问题描述】：

我手头有以下问题。 F 是一个维度为 2 X 100 X 65 的 NumPy 数组。我想生成另一个数组V，其维度为2 X 2 X 65。此数组V 必须按以下方式计算：

对于每个t、V[:, :, t] = F[:, :, t] @ F[:, :, t].T，其中.T 表示矩阵转置，@ 表示通常的矩阵乘法。现在，我正在使用以下方法：

aux_matrix = np.matmul(np.transpose(F, (2, 0 ,1)), np.transpose(F, (2, 1, 0)))
V = np.transpose(aux_matrix, (1, 2, 0))

我知道np.tensordot 和np.einsum 可以帮助解决这种情况，让事情变得更快、更优雅。然而，我是张量的新手，我不习惯爱因斯坦符号。有人可以提供一些关于如何执行此计算的信息，并可能在此处链接指导初学者的参考吗？谢谢！

【问题讨论】：

numpy.org/doc/1.19/reference/generated/numpy.einsum.html
我读到了，但我仍然无法解决我的问题@wwii。我不知道 np.eisum 是否真的可以在这里使用，或者我是否需要张量积。
我以为问题是link a reference guiding a beginner here?
感谢您提供该链接。我试图将我的问题放在其中一个示例中，但我没有这样做。我认为这是“批量乘法”的情况，但我不确定。
您的matmul 使用没有问题。它可能会比einsum 更快。如果批量维度是第一个，它会更简单。在 numpy 中，前导维度是最外层的，最适合批量使用。

标签： python numpy transpose tensor numpy-einsum

【解决方案1】：

正如评论中的解决方案所说，einsum 相当于解决方案，

np.einsum("ijk,njk->ink", F, F)

按照 einsum 的规则，axis=1 和 axis=0（对应于标签 j 和 k 的轴）将按元素相乘。其中，最终输出中缺少j，而k 存在。这意味着与j 对应的轴将在最终解决方案中相加，k 将停止在逐元素乘法中。

一般来说，如果标签在符号中重复，它们将按元素相乘，如果最终输出中缺少标签，则在按元素相加的基础上进行相加。

这正是这里发生的事情。 F 的形状为 (2, 100, 65)。 "ijk,njk->ink" 在这种情况下会执行以下操作 -

第一个数组的axis=0 的元素将与第二个数组的axis=0 的每个元素一起操作。这就是字符串中的i,n->in 所代表的内容。这里，i=2 和 n=2 和最终矩阵的前两个维度由in 给出，因此最终数组的前两个维度将具有(2, 2) 的形状。
与j 对应的轴在两个数组中重复。因此，它们将按元素相乘。但是，最终输出中缺少这一点，因此该结果将沿该方向求和。也就是说，最终输出中将缺少对应于 100 的维度。
与k 对应的标签也会发生类似的事情，但最终输出中仍然存在，因此不会发生求和。因此，最终输出的轴对应于 65。

如果您检查最终输出的形状，如预期的那样，它是 (2, 2, 65)。

我希望这能消除 OP 在评论中表达的疑问。

但是，假设这在性能方面自动优于matmul 格式是不正确的。就可读性而言，也许。但实际性能可能取决于数组的大小和相对维度，以及许多其他因素。

如果您将optimize=True 键添加到einsum，则值得检查性能是否发生变化，因为我已经看到这在某些情况下会产生巨大的性能差异。然而，对于这个数组的大小，这似乎让事情变得更糟（这可能是因为 einsum 找出优化数组的好方法所花费的时间，考虑到相对较小的大小，这可能是不合理的数组）。

我的经验法则是，如果您能够在不使用任何额外的 for 循环的情况下使用 matmul 找到解决方案，那么如果您主要关心的是性能，请坚持使用它。另一方面，如果你的程序有一堆for 循环，试试einsum，有没有optimize=True。但是，即使在这种情况下，there are some instances 的解决方案也具有原生 for 循环的性能优于 einsum，具体取决于数组的相对尺寸。

【讨论】：