numpy 中的批量矩阵乘法

Question

我有两个 numpy 数组a和 ，b形状分别为[5, 5, 5]和[5, 5]。对于两者来说a，b形状中的第一个条目是批量大小。当我执行矩阵乘法选项时，我得到一个 shape 数组[5, 5, 5]。MWE如下。

import numpy as np

a = np.ones((5, 5, 5))
b = np.random.randint(0, 10, (5, 5))
c = a @ b
# c.shape is (5, 5, 5)

假设我要对批量大小运行一个循环，即a[0] @ b[0].T，它将产生一个 shape 的数组[5, 1]。最后，如果我沿着轴 1 连接所有结果，我将得到一个 shape 的结果数组[5, 5]。下面的代码更好地描述了这些行。

a = np.ones((5, 5, 5))
b = np.random.randint(0, 10, (5, 5))
c = []
for i in range(5):
    c.append(a[i] @ b[i].T)
c = np.concatenate([d[:, None] for d in c], axis=1).T
# c.shape evaluates to be (5, 5)

我可以在不使用循环的情况下获得上述功能吗？例如，PyTorch 提供了一个名为 的函数torch.bmm来计算此值。谢谢。

Answer 1

添加额外的维度以b使矩阵乘法批量兼容，并通过挤压删除最后多余的维度：

c = np.matmul(a, b[:, :, None]).squeeze(-1)

或者等价地：

c = (a @ b[:, :, None]).squeeze(-1)

a在您的示例中，两者都通过将和b重塑b为5 x 5 x 1来进行矩阵乘法。

Answer 2

您可以使用 numpy einsum 来解决这个问题。

c = np.einsum('BNi,Bi ->BN', a, b)

Pytorch 还提供了这个 einsum 函数，但语法略有变化。所以你可以很容易地解决它。它也可以轻松处理其他形状。

那么你就不必担心转置或挤压操作。它还可以节省内存，因为内部不会创建现有矩阵的副本。