3d numpy 中的矩陣乘法

Question

我正在使用 numpy 執行矩陣乘法，但我無法弄清楚如何利用 numpy 進行 3d 矩陣乘法。

假設我有一個 3x3 矩陣 a，然后將它乘以一個 3x1 向量 b。 這將給出一個 3x1 向量 c。

這是在 numpy 中完成的：

# (3, 3) * (3, 1) -> (3, 1)
c = np.matmul(a, b)

好的，所以現在我想對本質上是 2500 個 3x3 矩陣的 3d 矩陣執行類似的操作。 現在我正在做一些事情：

# (2500, 3, 3) * (2500, 3, 1) -> list of (3, 1) vectors with length 2500
C = [np.matmul(a, b) for a, b in zip(A, B)]

它返回 (3, 1) 個向量的列表。

我寧願不循環，而是充分利用 numpy 的矢量化和矩陣/張量產品。 有沒有什么手術可以做...

# (2500, 3, 3) * (2500, 3, 1) -> (2500, 3, 1)
np.<function>(A, B, <args>)

我見過有關使用 np.tensordot 的東西，但我不知道如何設置軸。

np.tensordot(A, B, axes=???)

Answer 1

對於您擁有的 3 維數組（或 3 階張量），您可以使用np.einsum doc進行更復雜的矩陣乘法。 在您的特定情況下，您可以使用以下內容

>>> import numpy as np
>>> x = np.random.randint(0, 3, (3, 3, 3))  # shape (3, 3, 3)
>>> y = np.random.randint(0, 3, (3, 3, 3))  # shape (3, 3, 3)
>>> np.einsum('ijk,ikl->ijl', x, y)  # still shape (3, 3, 3)

特別是， einsum表達式'ijk,ikl->ijl'意味着對於每個第i個矩陣，執行正則矩陣乘法jk,kl->jl並將結果放入結果張量（ndarray）的第i個條目中。 此過程的更一般形式可能是

np.einsum('...jk,...kl->...jl', x, y)

你可以在你擁有的每個張量（ndarray）前面有任意數量的維度。

有關完整示例，請參見以下內容：

>>> import numpy as np
>>> x = np.random.randint(0, 3, (3, 3, 3))  # shape (3, 3, 3)
>>> x
array([[[0, 0, 1],
        [2, 2, 1],
        [2, 1, 1]],

       [[2, 0, 2],
        [2, 2, 1],
        [2, 2, 2]],

       [[2, 2, 2],
        [1, 1, 2],
        [0, 2, 2]]])
>>> y = np.random.randint(0, 3, (3, 3, 3))  # shape (3, 3, 3)
>>> y
array([[[0, 0, 1],
        [2, 1, 0],
        [0, 0, 2]],

       [[1, 2, 0],
        [2, 0, 1],
        [2, 2, 1]],

       [[0, 2, 1],
        [0, 1, 0],
        [0, 2, 1]]])
>>> np.einsum('ijk,ikl->ijl', x, y)
array([[[ 0,  0,  2],
        [ 4,  2,  4],
        [ 2,  1,  4]],

       [[ 6,  8,  2],
        [ 8,  6,  3],
        [10,  8,  4]],

       [[ 0, 10,  4],
        [ 0,  7,  3],
        [ 0,  6,  2]]])
>>> np.einsum('...ij,...jk->...ik', x, y)
array([[[ 0,  0,  2],
        [ 4,  2,  4],
        [ 2,  1,  4]],

       [[ 6,  8,  2],
        [ 8,  6,  3],
        [10,  8,  4]],

       [[ 0, 10,  4],
        [ 0,  7,  3],
        [ 0,  6,  2]]])

Answer 2

np.matmul(A,B)工作得很好。 你得到了什么錯誤？

In [263]: A,B = np.arange(24).reshape(2,3,4), np.arange(8).reshape(2,4,1)

einsum解決方案：

In [264]: np.einsum('ijk,ikl->ijl',A,B)
Out[264]: 
array([[[ 14],
        [ 38],
        [ 62]],

       [[302],
        [390],
        [478]]])
In [265]: _.shape
Out[265]: (2, 3, 1)

matmul解決方案：

In [266]: A@B
Out[266]: 
array([[[ 14],
        [ 38],
        [ 62]],

       [[302],
        [390],
        [478]]])

你的循環：

In [267]: [np.matmul(a, b) for a, b in zip(A, B)]
Out[267]: 
[array([[14],
        [38],
        [62]]),
 array([[302],
        [390],
        [478]])]

matmul文檔：

If either argument is N-D, N > 2, it is treated as a stack of
matrices residing in the last two indexes and broadcast accordingly.