3D張量的內部尺寸的矩陣乘法？

Question

我有兩個分別為(386, 3, 4)和(386, 4, 3) (386, 3, 4) numpy矩陣。 我想產生一個輸出尺寸為(386, 3, 3) 。 換句話說，我希望以向量化方式執行以下循環-

for i in range(len(input1)):
    output[i] = np.matmul(input1[i], input2[i])

最好的方法是什么？

Answer 1

matmul還可以工作：

a = np.random.random((243,3,4))
b = np.random.random((243,4,3))
np.matmul(a,b).shape
# (243, 3, 3)

Answer 2

我們需要使第一個軸保持對齊，因此我建議對np.einsum使用一種方法-

np.einsum('ijk,ikl->ijl',input1,input2)

樣品運行以驗證形狀-

In [106]: a = np.random.rand(386, 3, 4)

In [107]: b = np.random.rand(386, 4, 3)

In [108]: np.einsum('ijk,ikl->ijl',a,b).shape
Out[108]: (386, 3, 3)

樣本運行以驗證較小輸入上的值-

In [174]: a = np.random.rand(2, 3, 4)

In [175]: b = np.random.rand(2, 4, 3)

In [176]: output = np.zeros((2,3,3))

In [177]: for i in range(len(a)):
     ...:     output[i] = np.matmul(a[i], b[i])
     ...:     

In [178]: output
Out[178]: 
array([[[ 1.43473795,  0.860279  ,  1.17855877],
        [ 1.91036828,  1.23063125,  1.5319063 ],
        [ 1.06489098,  0.86868941,  0.84986621]],

       [[ 1.07178572,  1.020091  ,  0.63070531],
        [ 1.34033495,  1.26641131,  0.79911685],
        [ 1.68916831,  1.63009854,  1.14612462]]])

In [179]: np.einsum('ijk,ikl->ijl',a,b)
Out[179]: 
array([[[ 1.43473795,  0.860279  ,  1.17855877],
        [ 1.91036828,  1.23063125,  1.5319063 ],
        [ 1.06489098,  0.86868941,  0.84986621]],

       [[ 1.07178572,  1.020091  ,  0.63070531],
        [ 1.34033495,  1.26641131,  0.79911685],
        [ 1.68916831,  1.63009854,  1.14612462]]])

樣本運行以驗證較大輸入上的值-

In [180]: a = np.random.rand(386, 3, 4)

In [181]: b = np.random.rand(386, 4, 3)

In [182]: output = np.zeros((386,3,3))

In [183]: for i in range(len(a)):
     ...:     output[i] = np.matmul(a[i], b[i])
     ...:     

In [184]: np.allclose(np.einsum('ijk,ikl->ijl',a,b), output)
Out[184]: True