torch.bmm函数

torch.bmm函数是PyTorch中的一个函数，用于执行批量矩阵乘法（Batch Matrix Multiply）。它接受三个输入张量：input1，input2和out。其中，input1的形状是(batch_size, n, m)，input2的形状是(batch_size, m, p)，out的形状是(batch_size, n, p)。函数会计算input1和input2的批量矩阵乘法，并将结果写入out张量。

具体来说，对于每个batch，torch.bmm函数会将input1中的每个(n, m)矩阵与input2中的对应(m, p)矩阵相乘，得到一个(n, p)的结果矩阵，并将其存储在out中。因此，输出张量out的形状是(batch_size, n, p)。

以下是使用torch.bmm函数执行批量矩阵乘法的示例代码：

import torch

# 创建输入张量
batch_size = 2
n = 3
m = 4
p = 5
input1 = torch.randn(batch_size, n, m)
input2 = torch.randn(batch_size, m, p)

# 执行批量矩阵乘法
out = torch.bmm(input1, input2)

print(out.shape)  # 输出: torch.Size([2, 3, 5])

在上述示例中，我们创建了两个batch的输入张量input1和input2，它们的形状分别是(2, 3, 4)和(2, 4, 5)。然后，我们使用torch.bmm函数执行批量矩阵乘法，并将结果存储在out张量中。最后，我们打印出out张量的形状，得到(2, 3, 5)。