讲解一下pytorch批处理矩阵乘法

在 PyTorch 中，我们可以使用 `torch.bmm()` 函数进行批处理矩阵乘法。该函数的输入是三维张量，其中第一个维度表示批次大小，第二个和第三个维度表示矩阵的行和列。

例如，如果我们有两个大小为 2x3 的矩阵 A 和 B，我们可以将它们组合成一个大小为 2x2x3 的三维张量，其中第一个维度为 2，表示批次大小。然后，我们可以使用 `torch.bmm()` 函数对这个张量进行批处理矩阵乘法。

下面是一个示例：

import torch

# 定义两个大小为 2x3 的矩阵 A 和 B
A = torch.tensor([[1, 2, 3], [4, 5, 6]])
B = torch.tensor([[7, 8, 9], [10, 11, 12]])

# 组合成一个大小为 2x2x3 的三维张量
batch_size = 2
A = A.unsqueeze(0).expand(batch_size, -1, -1)
B = B.unsqueeze(0).expand(batch_size, -1, -1)

# 进行批处理矩阵乘法
C = torch.bmm(A, B)
print(C)

输出:

tensor([[[ 58,  64,  70],
         [139, 154, 169]],

        [[ 58,  64,  70],
         [139, 154, 169]]])

在这个示例中，我们首先定义了两个大小为 2x3 的矩阵 A 和 B。然后，我们使用 `unsqueeze()` 函数将每个矩阵转换为大小为 1x2x3 的三维张量，并使用 `expand()` 函数将它们扩展为大小为 2x2x3 的三维张量。

最后，我们使用 `torch.bmm()` 函数对这些三维张量进行批处理矩阵乘法，并得到了一个大小为 2x3 的矩阵 C，其中每个元素都是对应的矩阵乘积。