PyTorch数学函数避坑指南：torch.mm和torch.matmul到底有什么区别？

PyTorch矩阵乘法函数深度解析：torch.mm与torch.matmul的工程实践指南

1. 矩阵乘法基础与PyTorch实现概览

矩阵乘法是深度学习中最核心的数学运算之一，PyTorch提供了多种矩阵乘法实现函数，其中torch.mm和torch.matmul是最常用的两种。理解它们的差异对于编写高效、正确的深度学习代码至关重要。

在PyTorch中，矩阵乘法函数可以分为三个主要类别：

• 基础矩阵乘法：torch.mm
• 广播矩阵乘法：torch.matmul
• 批量矩阵乘法：torch.bmm

import torch

# 基础矩阵乘法示例
a = torch.randn(2, 3)
b = torch.randn(3, 4)
c = torch.mm(a, b)  # 输出形状：(2, 4)

# 广播矩阵乘法示例
a = torch.randn(5, 1, 2, 3)
b = torch.randn(3, 4)
c = torch.matmul(a, b)  # 输出形状：(5, 1, 2, 4)

2. torch.mm函数深度解析

torch.mm是PyTorch中最基础的矩阵乘法函数，专门用于执行严格的二维矩阵乘法运算。

2.1 函数定义与参数说明

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

参数说明：

• input (Tensor): 第一个相乘矩阵，形状为(n×m)
• mat2 (Tensor): 第二个相乘矩阵，形状为(m×p)
• out (Tensor, optional): 输出张量

2.2 严格维度要求与典型应用场景

torch.mm对输入张量有严格的二维要求：

• 输入必须是精确的2D张量
• 第一个矩阵的列数必须等于第二个矩阵的行数

典型应用场景：

• 全连接层的权重计算
• 传统线性代数运算
• 小规模矩阵乘法（<1000维）

2.3 性能特点与内存占用分析

torch.mm在实现上有以下特点：

1. 无广播开销：由于严格的维度限制，省去了广播检查的开销
2. 专用内核优化：PyTorch对其有专门的优化实现
3. 内存效率高：不会产生额外的内存占用

性能对比实验（在RTX 3090上）：