跨多个维度的矩阵向量积

Question

我有两个数组：

A = torch.rand((64, 128, 10, 10))
B = torch.rand((64, 128, 10))

我想计算由 C 表示的乘积，我们在 A 和 B 的第一维和第二维上进行矩阵向量乘法，因此：

# C should have shape: (64, 128, 10)
for i in range(0, 64):
   for j in range(0, 128):
       C[i,j] = torch.matmul(A[i,j], B[i,j])

有谁知道如何使用 torch.einsum 来做到这一点？我尝试了以下方法，但得到的结果不正确。

C = torch.einsum('ijkl, ijk -> ijk', A, B)

Answer 1

这是 numpy 的选项。（我没有torch）

In [120]: A = np.random.random((64, 128, 10, 10))
     ...: B = np.random.random((64, 128, 10))

您的迭代参考案例：

In [122]: C = np.zeros((64,128,10))
     ...: # C should have shape: (64, 128, 10)
     ...: for i in range(0, 64):
     ...:    for j in range(0, 128):
     ...:        C[i,j] = np.matmul(A[i,j], B[i,j])
     ...:

matmul 完整广播:

In [123]: D  = np.matmul(A, B[:,:,:,None])
In [125]: C.shape
Out[125]: (64, 128, 10)
In [126]: D.shape            # D has an extra size 1 dimension
Out[126]: (64, 128, 10, 1)
In [127]: np.allclose(C,D[...,0])    # or use squeeze
Out[127]: True

相当于einsum：

In [128]: E = np.einsum('ijkl,ijl->ijk', A, B)
In [129]: np.allclose(C,E)
Out[129]: True

跨多个维度的矩阵向量积

Matrix Vector Product across Multiple Dimensions

numpy

pytorch

matrix-multiplication

numpy-einsum