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

How to perform matrix multiplication between two 3D tensors along the first dimension?

我想计算两个 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)

遗憾的是 returns 这个错误:

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

np.tensordot 失败后,我选择了爱因斯坦求和。非常欢迎提出想法和跟进问题!

您的尺寸 55 的两个维度不对应于相同的轴。因此,您需要使用两个不同的下标来指定它们。例如,您可以这样做:

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

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

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

  • 外积:

    >>> np.einsum('ijk,ilm->ijm', a, b)
    
  • 下标 k 上的点积,即 aaxis=2baxis=1:

    >>> np.einsum('ijk,ikm->ijm', a, b)
    

    相当于a@b.

dot product ... along the first dimension有点不清楚。第一个维度是 'batch' 维度,其余维度是 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