torch.bmm()

torch.bmm() 是 PyTorch 中用于批量矩阵乘法的函数。全称是 batch matrix-matrix product，用于对一批矩阵进行乘法运算。

功能

torch.bmm(input, mat2) 进行批量矩阵乘法运算，即对两个三维张量 input 和 mat2 中对应的二维矩阵进行乘法操作。要求 input 和 mat2 都是形状为 [b, n, m] 和 [b, m, p] 的三维张量，b 是批量的大小，n 和 p 分别是矩阵的行数和列数。

形状

• input：形状为 [b, n, m] 的三维张量，其中 b 是批量大小，n 是矩阵的行数，m 是矩阵的列数。
• mat2：形状为 [b, m, p] 的三维张量，其中 b 是批量大小，m 是矩阵的行数，p 是矩阵的列数。
• 输出：形状为 [b, n, p] 的三维张量。

解释

torch.bmm() 实现的是如下批量矩阵乘法：

其中 i 范围为 0 到 b-1，即对每一个批次的矩阵分别进行矩阵乘法运算。

注意

1. 维度匹配：确保 input 的第二个维度和 mat2 的第一个维度匹配（即 m 值相同）。
2. 三维张量：torch.bmm() 只能用于三维张量。如果需要对更高维度的张量进行操作，可以考虑使用 torch.matmul() 或 torch.einsum()。

应用

torch.bmm() 在深度学习中常用于需要同时对多个矩阵进行乘法操作的场景，如：

• 批量处理多个样本的线性变换。
• 在 RNN、LSTM 等循环神经网络中处理多个时间步的矩阵乘法。
• 需要对多个特征映射进行并行计算时。