介绍

torch.matmul(input, other) 是一个通用的矩阵乘法函数，它会根据输入张量的维度（阶数）自动调整乘法规则，比 torch.mm()（仅支持二维矩阵）和 torch.mul()（逐元素乘法）更灵活。

简单来说：

• 标量 × 标量：返回标量乘积
• 一维向量 × 一维向量：返回点积（标量）
• 二维矩阵 × 二维矩阵：标准矩阵乘法（需满足列数=行数）
• 高维张量（批量矩阵）：对最后两个维度做矩阵乘法，前面的维度作为“批量维度”自动广播

矩阵乘法计算过程

(m,i)的矩阵 * (i,n)的矩阵 = (m,n)的矩阵

例子

先看基础用法，所有示例都基于以下前置代码：

import torch

# 初始化不同维度的张量
# 标量
a_scalar = torch.tensor(2.0)
b_scalar = torch.tensor(3.0)

# 一维向量 (shape: [n])
a_vec = torch.tensor([1, 2, 3])  # shape: [3]
b_vec = torch.tensor([4, 5, 6])  # shape: [3]

# 二维矩阵 (shape: [m, n])
a_mat = torch.tensor([[1, 2], [3, 4]])  # shape: [2, 2]
b_mat = torch.tensor([[5, 6], [7, 8]])  # shape: [2, 2]

# 高维批量矩阵 (shape: [batch, m, n])
a_batch = torch.randn(2, 3, 4)  # 2个 3×4 的矩阵
b_batch = torch.randn(2, 4, 5)  # 2个 4×5 的矩阵

标量 × 标量

result = torch.matmul(a_scalar, b_scalar)
print(result)  # tensor(6.)

一维向量 × 一维向量（点积）

result = torch.matmul(a_vec, b_vec)
print(result)  # tensor(1*4 + 2*5 + 3*6) = tensor(32)

提示：相当于 pytorch api文档：torch.dot()点积

二维矩阵 × 二维矩阵（标准矩阵乘法）

result = torch.matmul(a_mat, b_mat)
print(result)
# 输出：
# tensor([[19, 22],  # 1*5+2*7=19, 1*6+2*8=22
#         [43, 50]])  # 3*5+4*7=43, 3*6+4*8=50

高维批量矩阵乘法

result = torch.matmul(a_batch, b_batch)
print(result.shape)  # torch.Size([2, 3, 5])
# 解释：对每个batch（第0维），执行 3×4 × 4×5 = 3×5 的矩阵乘法

广播机制（维度不匹配时自动扩展）

# 场景：1个 3×4 矩阵 与 2个 4×5 矩阵相乘（批量广播）
a = torch.randn(3, 4)    # shape: [3,4]
b = torch.randn(2, 4, 5) # shape: [2,4,5]
result = torch.matmul(a, b)
print(result.shape)      # torch.Size([2, 3, 5])
# 解释：a 被自动广播为 [2,3,4]，再与 b [2,4,5] 做批量矩阵乘法

注意事项

• 维度兼容性：对于二维矩阵 A(m×n) 和 B(n×p)，A 的列数必须等于 B 的行数，否则会报错。
• 数据类型：输入张量需为浮点型（float32/float64）或复数型，整数张量乘法需用 torch.matmul 或 torch.int_matmul（PyTorch 2.0+）。

等价 @ 运算符

@ 和 torch.matmul() 完全等价，唯一的区别是写法和使用场景的便捷性。

例子

先通过代码直观验证，你可以直接运行：

import torch

# 示例1：一维向量点积
a_vec = torch.tensor([1, 2, 3])
b_vec = torch.tensor([4, 5, 6])
print(torch.matmul(a_vec, b_vec))  # tensor(32)
print(a_vec @ b_vec)               # tensor(32)

# 示例2：二维矩阵乘法
a_mat = torch.tensor([[1,2], [3,4]])
b_mat = torch.tensor([[5,6], [7,8]])
print(torch.matmul(a_mat, b_mat))
# tensor([[19, 22],
#         [43, 50]])
print(a_mat @ b_mat)
# 输出和上面完全一致

# 示例3：高维批量矩阵乘法
a_batch = torch.randn(2, 3, 4)
b_batch = torch.randn(2, 4, 5)
print((torch.matmul(a_batch, b_batch)).shape)  # torch.Size([2, 3, 5])
print((a_batch @ b_batch).shape)               # torch.Size([2, 3, 5])

无论一维、二维还是高维张量，@ 的计算结果和 torch.matmul() 完全相同。

细微区别：写法与使用场景

虽然功能等价，但二者在写法和适用场景上有小差异：

不同场景的写法对比

# 场景1：链式矩阵乘法
a = torch.randn(2,3)
b = torch.randn(3,4)
c = torch.randn(4,5)

# matmul 写法（嵌套，易乱）
res1 = torch.matmul(torch.matmul(a, b), c)
# @ 写法（链式，清晰）
res2 = a @ b @ c
print(torch.equal(res1, res2))  # True

# 场景2：作为参数传入函数
list_a = [torch.randn(2,2) for _ in range(3)]
list_b = [torch.randn(2,2) for _ in range(3)]

# matmul 可直接传参
res3 = list(map(torch.matmul, list_a, list_b))
# @ 需包装成lambda
res4 = list(map(lambda x,y: x@y, list_a, list_b))
print(torch.equal(res3[0], res4[0]))  # True

与其他乘法的区别

• torch.mul(a, b)：逐元素乘法（对应位置相乘），要求维度完全一致（或可广播）。
• torch.mm(a, b)：仅支持二维矩阵乘法，无广播，功能是 torch.matmul 的子集。

总结

1. torch.matmul() 是 PyTorch 中通用的张量乘法函数，支持标量、向量、矩阵、批量矩阵的乘法，自动适配维度。
2. 核心规则：对最后两个维度执行矩阵乘法，前面的维度作为批量维度，支持广播。
3. 区分 matmul（矩阵乘法）、mul（逐元素乘法）、mm（仅二维矩阵乘法）的使用场景，避免混用。

原文出处：/show_1GW2b5HTn2z0.html