PyTorch 中mm和bmm函数的使用详解

torch.mm 是 PyTorch 中用于 二维矩阵乘法（matrix-matrix multiplication） 的函数，等价于数学中的 A × B 矩阵乘积。

一、函数定义

torch.mm(input, mat2) → Tensor

执行的是两个 2D Tensor（矩阵）的标准矩阵乘法。

• input: 第一个二维张量，形状为 (n × m)
• mat2: 第二个二维张量，形状为 (m × p)
• 返回：形状为 (n × p) 的张量

  ---

  二、使用条件和注意事项

  条件	说明
  仅支持 2D 张量	一维或三维以上使用 torch.matmul 或 @ 操作符
  维度要匹配	即 input.shape[1] == mat2.shape[0]
  不支持广播	两个矩阵维度不匹配会直接报错
  结果是普通矩阵乘积	不是逐元素乘法（Hadamard），即不是 * 或 torch.mul()

  ---

  三、示例代码

  示例 1：基本矩阵乘法

  import torch
  A = torch.tensor([[1., 2.], [3., 4.]])   # 2x2
  B = torch.tensor([[5., 6.], [7., 8.]])   # 2x2
  C = torch.mm(A, B)
  print(C)
  

  输出：

  tensor([[19., 22.],
          [43., 50.]])
  

  计算步骤：

  C[0][0] = 1*5 + 2*7 = 19
  C[0][1] = 1*6 + 2*8 = 22
  ...
  

  ---

  示例 2：不匹配维度导致报错

  A = torch.rand(2, 3)
  B = torch.rand(4, 2)
  C = torch.mm(A, B)  # ❌ 会报错
  

  报错：

  RuntimeError: mat1 and mat2 shapes cannot be multiplied (2x3 and 4x2)
  

  ---

  示例 3：推荐写法（推荐使用 @ 或 matmul）

  A = torch.rand(3, 4)
  B = torch.rand(4, 5)
  C1 = torch.mm(A, B)
  C2 = A @ B                # 推荐用法
  C3 = torch.matmul(A, B)   # 推荐用法
  

  ---

  四、与其他乘法函数的比较

  函数名	支持维度	运算类型	支持广播
  torch.mm	仅限二维	矩阵乘法	❌ 不支持
  torch.matmul	1D, 2D, ND	自动判断点乘 / 矩阵乘	✅ 支持
  torch.bmm	批量二维乘法	3D Tensor batch × batch	❌ 不支持
  torch.mul	任意维度	元素乘（Hadamard）	✅ 支持
  * 运算符	任意维度	元素乘	✅ 支持
  @ 运算符	ND（推荐用）	矩阵乘法（和 matmul 一样）	✅

  ---

  五、典型应用场景

  • 神经网络权重乘法：output = torch.mm(W, x)
  • 点云 / 图像变换：x' = torch.mm(R, x) + t
  • 多层感知机中的矩阵计算
  • 注意力机制中 QK^T 乘积

    ---

    六、总结：什么时候用 mm？

    使用场景	用什么
    仅二维矩阵乘法	torch.mm
    高维或支持广播乘法	torch.matmul / @
    批量矩阵乘法 (如 batch_size×3×3)	torch.bmm
    元素乘	torch.mul or *

    ---

    在 PyTorch 中，torch.bmm 是 批量矩阵乘法（batch matrix multiplication） 的操作，专用于处理三维张量（batch of matrices）。它的主要作用是对一组矩阵成对进行乘法，效率远高于手动循环计算。

    ---

    一、torch.bmm 语法

    torch.bmm(input, mat2, *, out=None) → Tensor
    

    • input: Tensor，形状为 (B, N, M)
    • mat2: Tensor，形状为 (B, M, P)
    • 返回结果形状为 (B, N, P)

      这表示对 B 对 N×M 和 M×P 的矩阵进行成对相乘。

      ---

      二、示例演示

      示例 1：基础用法

      import torch
      # 定义两个 batch 矩阵
      A = torch.randn(4, 2, 3)  # shape: (B=4, N=2, M=3)
      B = torch.randn(4, 3, 5)  # shape: (B=4, M=3, P=5)
      # 批量矩阵乘法
      C = torch.bmm(A, B)       # shape: (4, 2, 5)
      print(C.shape)  # 输出: torch.Size([4, 2, 5])
      

      ---

      示例 2：手动循环 vs bmm 效率对比

      # 慢速手动方式
      C_manual = torch.stack([A[i] @ B[i] for i in range(A.size(0))])
      # 等效于 bmm
      C_bmm = torch.bmm(A, B)
      print(torch.allclose(C_manual, C_bmm))  # True
      

      ---

      三、注意事项

      1. 维度必须是三维张量

      • 否则会报错：

        RuntimeError: batch1 must be a 3D tensor
        

        你可以通过 .unsqueeze() 手动调整维度：

        a = torch.randn(2, 3)
        b = torch.randn(3, 4)
        # 升维
        a_batch = a.unsqueeze(0)  # (1, 2, 3)
        b_batch = b.unsqueeze(0)  # (1, 3, 4)
        c = torch.bmm(a_batch, b_batch)  # (1, 2, 4)
        

        ---

        2. 维度必须满足矩阵乘法规则

        • (B, N, M) × (B, M, P) → (B, N, P)
        • 若 M 不一致会报错：

          RuntimeError: Expected size for the second dimension of batch2 tensor to match the first dimension of batch1 tensor
          

          ---

          3. bmm 不支持广播（broadcasting）

          • 必须显式提供相同的 batch size。
          • 如果只有一个矩阵固定，可以使用 .expand()：

            A = torch.randn(1, 2, 3)  # 单个矩阵
            B = torch.randn(4, 3, 5)  # 4 个矩阵
            # 扩展 A 以进行 batch 乘法
            A_expand = A.expand(4, -1, -1)
            C = torch.bmm(A_expand, B)  # (4, 2, 5)
            

            ---

            四、在实际应用中的例子

            在点云变换中：批量乘旋转矩阵

            # 假设有 B 个旋转矩阵和点坐标
            R = torch.randn(B, 3, 3)       # 旋转矩阵
            points = torch.randn(B, 3, N)  # 点云
            # 先转置点坐标为 (B, N, 3)
            points_T = points.transpose(1, 2)  # (B, N, 3)
            # 用 bmm 做点变换：每组点乘旋转
            transformed = torch.bmm(points_T, R.transpose(1, 2))  # (B, N, 3)
            

            ---

            五、总结

            特性	torch.bmm
            操作对象	三维张量（batch of matrices）
            核心规则	(B, N, M) x (B, M, P) = (B, N, P)
            是否支持广播	❌ 不支持，需要手动 .expand()
            与 matmul 区别	matmul 支持更多广播，bmm 更高效用于纯批量矩阵乘法
            应用场景	批量线性变换、点云配准、神经网络前向传播等

            ---