多维矩阵向量积

扎眼的阳光 2年前 pytorch 185

原文标题 ：Matrix Vector Product across Multiple Dimensions

我有两个数组：

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)

原文链接：https://stackoverflow.com//questions/71983170/matrix-vector-product-across-multiple-dimensions

我来回复

hpaulj 评论

该回答已被采纳！

这是带有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

2年前 0条评论

多维矩阵向量积

回复

相关问题