PyTorch中两个二维张量的批量矩阵乘法

在PyTorch中，批量矩阵乘法通常是指执行两个张量的矩阵乘法操作，其中这两个张量可以具有一个额外的批次维度。为了进行批量矩阵乘法，张量的尺寸需要满足特定的规则，其中一个张量的最后两个维度是矩阵的维度，而另一个张量的前两个维度是矩阵的维度。这样，每个矩阵对就可以独立地执行乘法操作。

PyTorch中的`torch.bmm`函数就是用来执行这种批量矩阵乘法的。具体来说，如果有两个张量`A`和`B`，它们的尺寸分别是`(batch_size, n, m)`和`(batch_size, m, p)`，那么使用`torch.bmm(A, B)`将会返回一个尺寸为`(batch_size, n, p)`的张量，其中每个`(i, :, :)`切片是`A[i, :, :]`和`B[i, :, :]`的矩阵乘积。

批量矩阵乘法的一个简单例子如下：

import torch

# 创建两个大小为(batch_size, n, m)和(batch_size, m, p)的张量
A = torch.randn(3, 2, 3)  # 例如，一个批次大小为3，矩阵大小为2x3的张量
B = torch.randn(3, 3, 4)  # 例如，一个批次大小为3，矩阵大小为3x4的张量

# 执行批量矩阵乘法
result = torch.bmm(A, B)

# 打印结果张量的大小
print(result.size())  # 输出应为torch.Size([3, 2, 4])

在这个例子中，我们创建了两个随机张量`A`和`B`，它们各自有3个矩阵。然后我们使用`torch.bmm`函数进行批量矩阵乘法，最终得到的结果也是一个包含3个矩阵的张量，每个矩阵是`A`中的对应矩阵与`B`中的对应矩阵的乘积。