深入理解PyTorch中的torch.mm, torch.bmm和torch.matmul函数

在PyTorch中，torch.mm、torch.bmm和torch.matmul都是用于矩阵乘法的函数，但它们在使用和功能上有一些重要的区别。

1. torch.mm (Matrix Multiply)

torch.mm函数用于执行两个矩阵之间的乘法操作。它要求输入的两个矩阵是兼容的，即它们的维度必须满足矩阵乘法的规则。具体来说，第一个矩阵的列数必须等于第二个矩阵的行数。

import torch
# 创建两个矩阵
A = torch.tensor([[1, 2], [3, 4]])
B = torch.tensor([[5, 6], [7, 8]])
# 使用torch.mm进行矩阵乘法
C = torch.mm(A, B)

1. torch.bmm (Batch Matrix Multiply)

torch.bmm函数用于执行批量矩阵之间的乘法操作。它接受三个维度大于或等于2的张量，并将前两个维度视为批量维度，对每个批量维度执行矩阵乘法操作。因此，输入张量的第0维和第1维的尺寸必须满足矩阵乘法的规则。

# 创建三个矩阵组成的批量矩阵
A = torch.tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
B = torch.tensor([[[9, 10], [11, 12]], [[13, 14], [15, 16]]])
# 使用torch.bmm进行批量矩阵乘法
C = torch.bmm(A, B)

1. torch.matmul (Matrix Multiplication)

torch.matmul函数是PyTorch中默认的矩阵乘法运算符，它既可以用于执行两个矩阵之间的乘法，也可以用于执行批量矩阵之间的乘法。它根据输入张量的维度自动选择使用torch.bmm还是torch.mm。如果输入张量的第0维和第1维的尺寸满足矩阵乘法的规则，则使用torch.mm；否则，使用torch.bmm。

# 创建两个矩阵
D = torch.tensor([[1, 2], [3, 4]])
E = torch.tensor([[5, 6], [7, 8]])
# 使用torch.matmul进行矩阵乘法（相当于torch.mm）
F = torch.matmul(D, E)

总结：在选择使用torch.mm、torch.bmm还是torch.matmul时，需要考虑你的具体需求。如果你需要执行两个矩阵之间的简单乘法操作，可以使用torch.mm。如果你需要执行批量矩阵之间的乘法操作，可以使用torch.bmm。如果你不确定输入张量的维度应该如何处理，可以使用torch.matmul，它会自动根据输入的维度选择合适的函数。