python - 将向量 (1 x N) 乘以张量 (N x M x M)

Question

我正在 numpy 中或最好在 pytorch 中寻找一种矩阵运算，它允许将向量 (1 x N) 乘以张量 (N x M x M) 并得到 (1 x M x M)。使用 for 循环可以轻松完成此操作，但 for 循环不允许在训练期间进行反向传播。我尝试在 numpy 和 pytorch(以及其他几个，例如 dot 和 bmm)中使用 matmul，但无法正常工作。这是我想要做的一个示例(其中 M=2，但在我的用例中是 256):

a = np.array([1,2,3])
b = np.array([[[1,2],[3,4]],[[5,6],[7,8]],[[9,10],[11,12]]])

我想执行以下操作:1*[[1,2],[3,4]] + 2*[[5,6],[7,8]] + 3*[[ 9,10],[11,12]]，这可以通过这样的 for 循环来实现:

for i in range(3):
  matrix_sum += a[i]*b[i]

任何建议或解决方案将不胜感激。

Answer 1

您可以使用简单的 einsum:

#this gives you 2-D array (M,M)
np.einsum('i,ijk->jk',a,b)

输出:

[[38 44]
 [50 56]]

或其他解决方案:

#this gives you 3-D array (1,M,M)
a[None,:]@b.swapaxes(0,1)

输出:

[[[38 44]
  [50 56]]]