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

时间:2020-07-29 06:23:16

标签: python arrays numpy pytorch matrix-multiplication

我正在寻找一种以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]

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

2 个答案:

答案 0 :(得分:3)

您可以使用简单的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]]]

答案 1 :(得分:1)

Numpy和pytorch是uppon矩阵乘法!

火炬示例:

A = torch.rand(1, N)
B = torch.rand(N, M, M)

C = A @ B.transpose(0, 1) 
C.transpose_(0, 1)
C.shape
torch.size(1, M, M)

对于numpy同样如此:

A = np.random.randn(1, N)
B = np.random.randn(N, M, M)

C = A @ B.transpose(1, 0, 2)
C = C.transpose(1, 0, 2)
C.shape
(1, M, M)

编辑:针对爱因索姆恋人:

Pytorch和numpy几乎以相同的方式处理einsum:

torch.einsum('i,ijk->jk', A, B)
np.einsum('i,ijk->jk', A, B)

Pytorch einsum文档:https://pytorch.org/docs/master/generated/torch.einsum.html 大量的Einsum文档:https://numpy.org/doc/stable/reference/generated/numpy.einsum.html