如何沿第一维在两个 3D 张量之间执行矩阵乘法？答案

【问题标题】：How to perform matrix multiplication between two 3D tensors along the first dimension?如何沿第一维在两个 3D 张量之间执行矩阵乘法？
【发布时间】：2021-09-22 18:24:58
【问题描述】：

我希望沿第一维计算两个 3D 张量之间的点积。我尝试了以下 einsum 表示法：

import numpy as np

a = np.random.randn(30).reshape(3, 5, 2)
b = np.random.randn(30).reshape(3, 2, 5)

# Expecting shape: (3, 5, 5)
np.einsum("ijk,ikj->ijj", a, b)

不幸的是它返回了这个错误：

ValueError: einstein sum subscripts string includes output subscript 'j' multiple times

在np.tensordot 失败后，我选择了 Einstein sum。非常欢迎提出想法和后续问题！

【问题讨论】：

沿第一个轴/维度的矩阵乘法听起来就像你只想要a @ b，它将为a 和b 的三对中的每一对执行矩阵乘法(5, 2) x (2, 5)。跨度>

标签： python numpy numpy-einsum

【解决方案1】：

您的两个尺寸 5 和 5 不对应于相同的轴。因此，您需要使用两个不同的下标来指定它们。例如，您可以这样做：

>>> res = np.einsum('ijk,ilm->ijm', a, b)

>>> res.shape
(3, 5, 5)

请注意，您还需要更改大小为 2 和 2 的轴的下标。这是因为您正在计算批量外积（即我们同时在两个轴上迭代），而不是点积（即我们同时在两个轴上迭代)。

外品：
```
>>> np.einsum('ijk,ilm->ijm', a, b)
```
下标k上的点积，即a的axis=2和b的axis=1：
```
>>> np.einsum('ijk,ikm->ijm', a, b)
```
相当于a@b。

【讨论】：

OP 需要什么？我以为他想要a@b。您的第一个示例得到了正确的形状，但首先将 a 和 b 减少到 (3,5,1) 和 (3,1,5)，然后进行广播外部。

【解决方案2】：

dot product ... along the first dimension 有点不清楚。第一个维度是“批次”维度，其余维度是 3 个点吗？还是别的什么？

In [103]: a = np.random.randn(30).reshape(3, 5, 2)
     ...: b = np.random.randn(30).reshape(3, 2, 5)
In [104]: (a@b).shape
Out[104]: (3, 5, 5)
In [105]: np.einsum('ijk,ikl->ijl',a,b).shape
Out[105]: (3, 5, 5)

@Ivan's 答案不同：

In [106]: np.einsum('ijk,ilm->ijm', a, b).shape
Out[106]: (3, 5, 5)
In [107]: np.allclose(np.einsum('ijk,ilm->ijm', a, b), a@b)
Out[107]: False

In [108]: np.allclose(np.einsum('ijk,ikl->ijl', a, b), a@b)
Out[108]: True

Ivan 将一个维度的k 和另一个维度的l 相加，然后按元素进行广播。那不是矩阵乘法：

In [109]: (a.sum(axis=-1,keepdims=True)* b.sum(axis=1,keepdims=True)).shape
Out[109]: (3, 5, 5)
In [110]: np.allclose((a.sum(axis=-1,keepdims=True)* b.sum(axis=1,keepdims=True)),np.einsum('ijk,ilm->ijm', a,
     ...:  b))
Out[110]: True

批处理的另一个测试：

In [112]: res=np.zeros((3,5,5))
     ...: for i in range(3):
     ...:    res[i] = a[i]@b[i]
     ...: np.allclose(res, a@b)
Out[112]: True

【讨论】：