我正在使用 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=???)

最佳答案

对于您拥有的 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]]])

https://stackoverflow.com/questions/65583672/