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

torch.mm(input, mat2) → Tensor

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)

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

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)

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

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

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])

# 慢速手动方式
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

RuntimeError: batch1 must be a 3D tensor

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)

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

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)