1. 先把“矩阵运算”这件事的边界划清楚我第一次被人问“PyTorch张量怎么做矩阵运算”对方的真实困惑其实不是不会写而是不知道自己写的那个到底调到了哪一层为什么有时候形状对得上、有时候对不上为什么同样的数学式子换成另一种写法速度快了一倍。这个问题在项目里被放大得很厉害小规模 demo 里for循环累加完全跑得动一旦数据量上来同样的逻辑换成批量矩阵运算耗时能从几秒掉到几十毫秒差的不是机器性能是你有没有用对 API。所以这篇内容不讲“张量是什么”这类从零起步的概念而是聚焦在PyTorch 张量上的矩阵运算这一个点上把它拆成三块乘法族的 API 到底各管什么、非乘法的线性代数操作求逆、解方程、分解、范数怎么用、以及形状/精度/性能这些真正会在项目里咬人的细节。适合已经能写出简单模型、但在写自定义损失、实现注意力、做批量相似度计算时被形状报错和性能问题卡住的人也适合从 NumPy 转过来、发现两边行为不完全一致的同学。以下所有结论都是我在实际项目和反复 benchmark 中得到的凡是涉及参数选择的我都会把推演过程写出来方便你自己判断。1.1 标量、向量、矩阵、张量在代码里到底差在哪数学上矩阵是二维的张量是高维的但在 PyTorch 里它们全都是torch.Tensor区别只体现在shape和dim()上。一个shape(3,4)的张量是矩阵shape(3,)的是向量shape(2,3,4)的是三维张量。这个“统一表示”带来一个副作用你没法靠类型系统区分“这是个矩阵”还是“这是个批量矩阵”只能靠 API 的约定去约束。关键就在于batch 维度的位置约定。PyTorch 里几乎所有批量矩阵 API 都默认批量维在前矩阵的两个维度在后。比如torch.bmm(A, B)要求A的形状是(b, n, m)、B是(b, m, p)输出(b, n, p)。这个约定不是随便定的它和内存布局有关连续张量的最后一个维度步长最小批量维在最前面意味着每个 batch 的矩阵是连续存放的一整块GPU 上做批量 GEMM 时可以一次性调度效率最高。很多人踩的坑是把 batch 维放到中间比如(n, b, m)然后硬套bmm结果报错或者拿到错误结果。这时候要么permute一下变成(b, n, m)要么换matmul让它走广播。permute回来的是非连续张量后续.view()会失败得用.reshape()或者先.contiguous()——这是我在处理变长序列时被教育过好几次的地方。还有一个容易被忽略的点PyTorch 里没有“行向量”和“列向量”的区分。shape(3,)的张量在参与matmul时会根据另一侧的形状自动决定是当行向量还是列向量用。这个规则在torch.matmul文档里写得很清楚但实际读代码时经常看不出来尤其是当两边都是一维的时候——torch.matmul(a, b)对两个一维张量做的是点积返回标量而torch.outer(a, b)返回的是外积矩阵。同样的两个输入函数不同结果维度完全不同这是新手最容易混淆的地方之一。1.2 为什么 PyTorch 要搞出那么多乘法函数torch.mm、torch.mv、torch.dot、torch.bmm、torch.matmul、torch.einsum、torch.outer光乘法就有七八个入口。乍看是历史包袱实际上每一个都对应一种明确的形状契约用对了能少写很多判断用错了报错信息也更直接。torch.mm只接受二维张量不接受广播不做任何形状兜底。这个“严格”恰恰是它的价值在代码 review 时看到mm你就能确定两个操作数都是矩阵不会出现“某个分支下传进来的是一维张量”的隐藏风险。torch.bmm同理严格三维不广播。把这种严格 API 用在核心计算路径上相当于给自己加了一道类型检查。torch.matmul是“万能版”它内部会做形状规约和广播能处理一维、二维、多维的任意组合。方便是方便代价是错误被推迟了形状不匹配的时候报错信息往往指向matmul内部提示“size mismatch”但真正的 bug 可能在十行之前就已经埋下。我自己的习惯是核心路径优先用严格 API只有在确实需要广播或者维度不固定的地方才用matmul这样出错时定位成本低得多。torch.einsum是另一类东西它不针对某一种形状而是用爱因斯坦求和约定描述“哪些维度参与缩并、哪些保留”。一个einsum表达式可以替代多重循环、permute bmm、diagonal、trace等一大堆操作。它的缺点是字母标记容易写错、初学者读起来费劲而且早期版本的einsum在某些形状下性能不如手写bmm。不过近几年优化已经好很多尤其是配合optimizeoptimal参数时它会自动搜索最优的缩并顺序。2. 乘法族 API 的能力边界与形状推演这一节把乘法相关的函数放在一起对照重点不是背函数签名而是建立一个判断流程看到两个张量要相乘先问自己三个问题——各是几维、需不需要广播、批量维在哪。这三个问题回答完用哪个函数基本就定了。2.1 mm、mv、dot、bmm、matmul 的适用条件对照把常用乘法 API 整理成一张表方便随时查。这里列的是“能不能用”不是“推荐用”实际选择还要结合可读性和性能。API输入维度要求是否广播输出维度典型用途torch.dot两个一维长度相同否标量向量内积、相似度torch.mv矩阵 一维向量否一维线性变换y Axtorch.outer两个一维否二维矩阵外积、秩一更新torch.mm两个二维否二维标准矩阵乘torch.bmm两个三维批量维相同否三维批量矩阵乘torch.matmul/任意遵循广播与一维规约是视输入而定通用乘法torch.einsum任意按下标约定否但可表达批量由表达式决定复杂缩并torch.tensordot任意指定缩并维否由dims决定高阶张量缩并torch.dot有一个限制经常被忽略它只接受一维张量且要求dtype一致。如果你拿一个shape(3,)和一个shape(3,1)去做dot会直接报错因为后者是二维。这类错误在 batch size 为 1 的时候特别隐蔽因为(3,1)在肉眼上和(3,)看起来差不多但shape完全不同。我在调试损失函数时遇到过一次原因就是某次unsqueeze多了一个维度导致内积变成了矩阵乘输出形状从标量变成(3,1)广播到后面全部算错但因为有广播程序不报错只是 loss 曲线莫名其妙地不对。torch.mv的第二个参数必须是一维。如果你手上是(n, 1)的列向量得先squeeze(-1)。反过来matmul就没这个限制(n, m) (m, 1)得到(n, 1)(n, m) (m,)得到(n,)两种都能跑但结果形状不同。这个差异在下游要做broadcast的时候会有实际影响比如后面接一个(n,)的 bias 相加前者(n,1)(n,)会广播成(n,n)后者(n,)(n,)是正确的。这类“静默广播错误”是矩阵运算里最难查的一类 bug因为程序不报错、结果也不崩就是数值不对。torch.bmm的批量维必须严格相等不做广播。如果两个张量的批量维一个是 1、一个是 32bmm会报错这时候要么expand成 32要么换matmul。expand出来的是共享内存的视图不占额外显存但它要求被扩展的维度大小是 1且后续如果是 in-place 操作会出问题。我一般直接用matmul让它自己处理广播省得纠结。torch.matmul的一维规约规则值得单独记一下如果第一个参数是一维会在它前面补一个 1 变成二维如果第二个参数是一维会在它后面补一个 1算完之后再把补的维度去掉。所以(3,) (3,)其实是(1,3) (3,1) (1,1)去掉补的维变成标量。(3,) (3,4)是(1,3) (3,4) (1,4)去掉前面的 1 得到(4,)。这条规则理解了很多“为什么结果维度少了”的困惑就解决了。2.2 广播机制matmul 真正强的地方广播这个词听起来抽象其实就是“维度大小对不上时如果有一边是 1就自动复制成另一边的大小”。规则是从最后一个维度开始往前对齐每一维要么相等、要么其中一个是 1、要么其中一方没有这一维否则报错。举个具体例子A形状(32, 1, 8, 16)B形状(16, 32)。先对齐B是二维A是四维B前面补两个 1 变成(1, 1, 16, 32)。逐维比较最后一维16vs32不相等且都不是 1报错。如果B是(16, 16)补完是(1, 1, 16, 16)最后一维16 vs 16通过倒数第二维8 vs 16又不通过。所以要让这个乘法成立B的后两个维度得是(8, x)或者(1, x)配合A的8。实际编码中广播最常见的用法是“单头 vs 多头”。比如你有一个(heads, head_dim, head_dim)的旋转矩阵要作用到(batch, heads, seq, head_dim)的张量上。直接matmulPyTorch 会把旋转矩阵前面补一个 1 变成(1, heads, head_dim, head_dim)然后和(batch, heads, seq, head_dim)做批量乘。这里1和batch广播、heads对齐、seq和head_dim对齐最后一个维度做矩阵乘输出(batch, heads, seq, head_dim)。整个式子写出来就是一行x R.transpose(-1, -2)比手写循环清爽太多。需要注意的是广播出来的张量是虚拟扩展不真的分配内存但一旦参与计算硬件层面可能还是会按扩展后的大小去读。expand之后连续做多次运算性能不一定比repeat好因为访存模式是重复读取同一块。我实测过(1, 64, 128, 128)扩展成(32, 64, 128, 128)再matmul的情况和直接repeat再算在小尺寸下差别不明显在seq_len上千的时候expand版本反而略慢原因是缓存局部性差。所以“expand 省显存”是对的但别默认它一定更快。还有一类坑是“意外的标量广播”。(n, 1) (n,)会得到(n, n)这个前面提过。更隐蔽的是(n, 1) * (1, m)得到(n, m)如果本意是逐元素对齐就会算出一个完全错误的矩阵。防御办法很简单在关键位置assert一下形状尤其是损失函数、mask 处理这类地方。一行assert x.shape y.shape能省掉半夜调参的时间。2.3 用 einsum 统一表达所有乘法形态einsum的核心是下标字符串输入张量各用一个字母序列标注维度箭头右边写输出保留的维度没出现在右边的字母就是被缩并掉的。ij,jk-ik是矩阵乘i,i-是内积i,j-ij是外积ii-是迹ij-ji是转置。它最大的好处是把“形状变换”和“计算”合成一句代码的数学含义直接可读。比如实现多头注意力的打分Q形状(b, h, s, d)K形状(b, h, t, d)要算(b, h, s, t)的分数矩阵。用matmul的写法是Q K.transpose(-1, -2)用einsum是torch.einsum(bhsd,bhtd-bhst, Q, K)。后者一眼就能看出d被缩并掉了输出的四个维度分别来自哪里不需要在脑子里做转置推演。再看一个matmul不太好表达的批量加权求和。feat形状(b, n, d)weight形状(b, n)要对每个样本按权重把n个向量加权平均成(b, d)。matmul写法是(weight.unsqueeze(1) feat).squeeze(1)中间要插维度、再去掉。einsum写法是torch.einsum(bn,bnd-bd, weight, feat)直接表达“沿 n 缩并”。可读性和出错概率都低不少。性能上einsum内部会把它翻译成bmm、permute、sum的组合。默认的缩并顺序是按表达式从左到右不一定是全局最优。加optimizeoptimal会让它枚举所有缩并顺序选计算量最小的那个——代价是第一次调用有额外的规划开销适合形状固定的场景比如固定的模型结构。如果是动态形状每次调用形状都变规划开销会抵消收益这时候用默认或者optimizegreedy更稳妥。有一个einsum的坑要说清楚它不做广播。所有出现在表达式里的维度输入张量必须有大小必须严格匹配。bhsd,bhtd-bhst里如果某个输入的h是 1、另一个是 12einsum会报错而matmul会自动广播。所以在写可变头数的代码时要自己处理单头扩展的情况不能指望einsum兜底。3. 非乘法类矩阵运算分解、求逆与解方程乘法只解决“乘”的问题实际项目里更多时候要解方程、求逆、做分解。PyTorch 把这部分统一放在torch.linalg命名空间下接口和 NumPy 的linalg基本对齐行为上有一些差异需要注意。3.1 求逆是最不该用的操作torch.linalg.inv能用但绝大多数情况下不该用。原因有三层一是数值稳定性显式求逆再相乘的误差通常比直接解方程大一个量级二是性能解方程是O(n^3)的分解加O(n^2)的回代求逆是O(n^3)加一次完整的矩阵乘常数更大三是显存n×n的逆矩阵要额外占一块空间而解方程只需要一个n×n的分解缓存。正确的写法是torch.linalg.solve(A, B)直接解AX B。它在内部做 LU 分解对称正定时走 Cholesky速度更快然后回代不显式构造逆矩阵。如果确实需要逆矩阵本身比如要拿它做后续的多次乘法可以用torch.linalg.inv_ex或者先lu_factor再lu_solve把分解结果复用起来。我在一个做高斯过程回归的模块里踩过这个坑。原代码是K_inv torch.linalg.inv(K)然后alpha K_inv y在小数据上没问题n到 2000 的时候开始出现loss nan。排查下来是K的条件数接近 1e12求逆过程中小奇异值被放大数值彻底崩掉。改成torch.linalg.solve(K, y)并加了一点点jitter对角线加1e-6问题消失。这里的关键不是solve有多神而是它避免了显式构造一个数值上很脆弱的中间量。torch.linalg.pinv用在非方阵或者秩亏的情况下内部走 SVD速度比solve慢不少但鲁棒。它的rcond参数控制小奇异值的截断阈值默认值是max(m, n) * eps实际用的时候通常要手动调大比如1e-5否则噪声方向会被当成有效信息保留下来反而有害。3.2 解线性方程组、最小二乘与矩阵分解除了solve还有几个必须知道的入口。torch.linalg.lstsq(A, B)直接解最小二乘min ||AX - B||不用你手动写正规方程。这一点很重要因为正规方程(A^T A)^{-1} A^T B会把条件数平方精度损失严重。lstsq内部走 QR 或者 SVD稳定性好很多。它的返回值和 PyTorch 版本有关新版返回一个 namedtuple包含solution、residuals、rank、singular_values取结果时注意用属性名而不是位置索引避免版本升级后出错。torch.linalg.qr做 QR 分解torch.linalg.svd做奇异值分解torch.linalg.eigh做对称矩阵的特征分解torch.linalg.cholesky做 Cholesky 分解。这几个是构建更复杂算法的基石比如 PCA 就是中心化之后取 SVD 的前 k 个奇异向量。svd有一个参数full_matrices默认True返回完整的U和Vh对于(m, n)且m n的矩阵U是(m, m)显存开销巨大。做降维的时候设成False只返回经济型分解U是(m, n)能省下大量显存。这个参数我在处理图像特征矩阵时吃过亏一个(50000, 128)的矩阵full_matricesTrue直接 OOM改成False之后轻松跑完。cholesky只对对称正定矩阵有效但它比solve快将近一倍而且能顺便判断正定性——如果矩阵不正定分解会直接报错。做协方差矩阵相关的运算时先试cholesky失败再退回solve是一个很实用的策略。3.3 范数、条件数与秩判断矩阵“好不好用”在动手算之前先花一行代码看看矩阵的“体质”能省掉大量调试时间。torch.linalg.norm(A, ord...)支持各种范数矩阵场景下常用ord2谱范数即最大奇异值、ordfroF 范数所有元素平方和开根、ord1列和最大值、ordtorch.inf行和最大值。做梯度裁剪的时候用的是向量范数ord2就是把所有元素平方求和开根这个和矩阵的谱范数不是一回事别混。torch.linalg.cond(A)算条件数等于最大奇异值除以最小奇异值。经验值小于 1e3 很安全1e3 到 1e6 要留意超过 1e8 基本可以认为在 float32 下不可逆超过 1e12 在 float64 下也悬。我一般会在涉及求逆或解方程的模块里加一个可选的 debug 分支训练前跑一次cond数值大了就报警比等到 loss 变 nan 再回头查高效得多。torch.linalg.matrix_rank(A)算数值秩内部也是走 SVD阈值同样由rcond控制。判断一个矩阵是不是满秩、要不要加正则项用它比看det靠谱——行列式接近 0 不代表一定奇异行列式很大也不代表条件数好这两个指标经常误导人。4. 从零跑一遍线性回归的正规方程与向量化改造理论讲完跑一段能直接复现的代码。这一段的目标不是训一个多好的模型而是把“同一个数学式子有四种写法性能和稳定性各不相同”这件事演示清楚。4.1 环境与数据构造不需要 GPU 也能跑CPU 上足够看出差异。import torch import time torch.manual_seed(42) torch.set_default_dtype(torch.float32) n_samples, n_features 5000, 64 X torch.randn(n_samples, n_features) true_w torch.randn(n_features, 1) y X true_w 0.01 * torch.randn(n_samples, 1) print(X.shape, y.shape) # torch.Size([5000, 64]) torch.Size([5000, 1])数据构造里有一个细节值得说y的形状是(n_samples, 1)而不是(n_samples,)。这个选择会影响后面所有形状推演。如果y是一维的X.t() y得到(64,)而X.t() X是(64, 64)解出来的w是(64,)如果y是二维的w是(64, 1)。两种都能用但混着用就会出问题比如预测值X w的形状在一维w下是(n,)二维w下是(n, 1)和y相减时如果y是(n, 1)、预测是(n,)会广播成(n, n)loss 直接爆炸。我的习惯是回归任务统一用二维(n, 1)全程形状一致少一类 bug。4.2 用 inv 硬算与用 solve 稳算的对比正规方程的数学形式是w (X^T X)^{-1} X^T y。三种实现从“最直译”到“最推荐”。XTX X.t() X # (64, 64) XTy X.t() y # (64, 1) # 写法一显式求逆数学上最直译实际最不该用 w_inv torch.linalg.inv(XTX) XTy # 写法二解方程稳定性好速度快 w_solve torch.linalg.solve(XTX, XTy) # 写法三直接最小二乘不走正规方程条件数不被平方 w_lstsq torch.linalg.lstsq(X, y).solution三者结果应该非常接近因为数据是良态的。为了看出差异给X加一点共线性把第 2 列设成第 1 列加一点噪声人为把条件数拉高。X_bad X.clone() X_bad[:, 1] X_bad[:, 0] 1e-4 * torch.randn(n_samples) print(cond , torch.linalg.cond(X_bad.t() X_bad).item())条件数到 1e8 以上时inv版本的结果会明显偏离solve稍微好一点lstsq最稳。这个实验我建议你自己跑一遍把绝对误差打出来看比看任何文字描述都有说服力。原因在于X^T X的条件数是X的条件数平方本来 1e4 的条件数经过平方变成 1e8float32 只有大约 7 位有效十进制数字误差直接吃掉了大部分精度。这也是为什么在特征维度高、特征之间相关性强的场景下宁可直接用lstsq或者加 L2 正则把XTX换成XTX λI也不要硬解正规方程。加正则之后的矩阵条件数会被λ拉低λ取1e-4量级就能把 1e8 降到 1e4 左右效果立竿见影。4.3 把三重循环改写成 einsum前一个例子是一步矩阵乘体现不出向量化的价值。换一个更真实的场景计算一批样本两两之间的加权相似度。假设有b个批次每批n个向量每个向量d维另有一个d×d的度量矩阵M要算S[i,j] x_i^T M x_j。朴素写法是三重循环def pairwise_naive(Xb, M): b, n, d Xb.shape out torch.zeros(b, n, n) for bi in range(b): for i in range(n): for j in range(n): out[bi, i, j] Xb[bi, i] M Xb[bi, j] return outn200、b8的时候这个循环要跑 32 万次内层迭代Python 层面开销巨大。改写分两步先把M吸收进X即Z X M形状(b, n, d)然后S Z X.transpose(-1, -2)形状(b, n, n)。用einsum表达就是def pairwise_einsum(Xb, M): Z torch.einsum(bnd,de-bne, Xb, M) return torch.einsum(bnd,bmd-bnm, Z, Xb)注意einsum的第二个表达式里两个输入的d是同一个字母表示缩并输出的n和m分别来自两个输入所以得到(b, n, n)。这里用不同字母n和m而不是都用n是为了避免einsum误解成“对角”或者“逐元素”——同一个字母出现在两个输入的非缩并位置语义是“这两个维度必须相等且逐元素对应”不会产生两两组合。这个细节很多人第一次写会搞错写成bnd,bnd-bnn得到的结果完全不同。实测下来n200, b8, d64朴素循环大约 1.8 秒einsum版本 3 毫秒左右差了两个数量级。而且循环版本还有一个隐患out[bi, i, j] ...是原地赋值如果out参与自动求导这种逐元素写入的方式在某些情况下会导致梯度计算变慢或者报错向量化版本完全没有这个问题。5. 形状、dtype、device 与数值精度的高频塌方点前面讲的是“怎么写对”这一节讲“为什么明明写对了还是出错”。矩阵运算的报错信息通常很简短size mismatch、expected scalar type、device mismatch但根因经常在很远的地方。5.1 形状报错的定位套路遇到形状报错别急着改代码先做三件事。第一把报错行涉及的所有张量形状打出来包括中间变量不要只打最终的两个操作数。很多时候根因是某个unsqueeze加错了位置或者某个sum少了dim参数导致维度凭空多出来一个。第二用matmul的规则手动推一遍广播。从最后一维往前对齐逐维检查把每一维的来源写下来。这个过程慢但能定位到具体哪一维不匹配。比如报错说Expected size for argument #2 mat2 (64) to be 32那就是第二个操作数的倒数第二维是 64第一个操作数的最后一维是 32它们要做矩阵乘必须相等。第三检查有没有“隐式 squeeze”。torch.sum(x, dim0)会真正去掉那一维而torch.sum(x, dim0, keepdimTrue)保留成 1。代码里混用这两种写法形状就会飘。我现在的习惯是凡是在流程中间做 reduction一律keepdimTrue到最后确定要输出的时候再显式squeeze这样形状变化都有迹可循。还有一种情况是 reshape 顺序问题。(b, n, d)直接reshape(b*n, d)是对的但如果中间做过transpose变成(b, d, n)再reshape(b*n, d)就会把内存顺序搞乱得到错误的数据不报错但结果全错。判断方法是用.is_contiguous()检查或者干脆用permute之后再.reshape()reshape会在必要时自动拷贝一份保证语义正确代价是多一次内存分配。这个坑在实现 Transformer 的多头拆分时特别常见(b, s, h*d)拆成(b, s, h, d)再转成(b, h, s, d)如果不用reshape而用view或者转置顺序写错后面所有注意力计算全是垃圾数据。5.2 dtype 与 device 不一致的连锁反应float32和float64混用矩阵乘会直接报错PyTorch 不做隐式升级。这一点和 NumPy 不同NumPy 会自动提升到更高精度。PyTorch 的选择是明确的混精度必须显式转换避免不知不觉中损失性能。排查套路是在报错处打印x.dtype和y.dtype。常见来源有三个一是数据加载时numpy默认float64torch.from_numpy之后没转二是某个常数写成了 Pythonfloat参与运算时被当成float32但如果另一边是float64就报错三是torch.zeros没指定dtype用了全局默认值而全局默认被某处set_default_dtype改过。device 不一致的报错更直接Expected all tensors to be on the same device。麻烦的是多卡场景模型在某张卡上数据在另一张卡上或者某个 buffer 没有跟着.to(device)走。我的做法是写一个工具函数在训练循环的第一步检查所有关键张量的 device 和 dtype一次性打印出来而不是等报错出现再逐个查。混合精度训练时还要注意autocast上下文里的矩阵乘会被自动转成半精度但如果手动调用了torch.linalg下的函数有些操作在半精度下不支持会报错或者悄悄降级。linalg.inv、linalg.svd这类操作在半精度下精度损失严重建议显式float()之后再算算完再转回来。5.3 浮点精度、TF32 与混合精度float32有 24 位有效尾数大约 7 位十进制精度。矩阵乘是累加操作d维内积累加d次误差大致按sqrt(d)增长。d1024时相对误差可能在 1e-5 量级对大多数任务够用但对需要高精度对齐的场景比如和参考实现做数值比对就会出问题。这时候要么用float64要么把比对阈值放宽到 1e-4。GPU 上还有一个额外变量TF32。这是较新架构的 GPU 在矩阵乘上默认开启的一种折中精度格式尾数只有 10 位比float32少很多但比float16多。它能显著提升矩阵乘吞吐代价是精度下降。PyTorch 里由torch.backends.cuda.matmul.allow_tf32控制不同版本默认值变过做数值复现实验时要显式设成False否则你会发现同样的代码在两张不同卡上结果不一致。混合精度里float16的动态范围小容易出现下溢。做矩阵乘之前如果输入的量级差异大建议用bfloat16它的指数位和float32一样宽不容易溢出代价是尾数更短、精度更低。实际项目里通常是bfloat16更省心尤其是在做注意力里的 softmax 之前float16经常需要把分数先缩放一下防溢出bfloat16基本不用管。5.4 原地操作与自动求导的冲突矩阵运算里最容易触发自动求导报错的是 in-place 操作。比如x.t_()、x.add_()、x[:, 0] 0如果x是叶子张量或者被其他计算依赖autograd 会报a leaf Variable that requires grad is being used in an in-place operation或者更隐蔽的版本号不匹配错误。t_()和transpose_()在较新版本里已经不推荐用了因为它们改变的是 stride 而不是数据和后续的视图、求导混在一起容易出问题。统一用t()和transpose()返回新视图虽然多了一个 Python 对象但没有副作用。填充 mask 的时候x.masked_fill_(mask, 0)是原地操作如果x需要梯度会报错。改成x.masked_fill(mask, 0)返回新张量就行或者在torch.no_grad()上下文里做原地填充。后者用在推理阶段更省显存训练阶段老老实实用非原地版本。6. 性能实测不同 API 的速度差异从哪来讲性能不能空谈得有数据。下面这段 benchmark 脚本可以直接复制运行CPU 和 GPU 都适配。6.1 一个简单的 benchmark 脚本import torch import time def bench(fn, warmup10, rep50): for _ in range(warmup): fn() if torch.cuda.is_available(): torch.cuda.synchronize() t0 time.perf_counter() for _ in range(rep): fn() if torch.cuda.is_available(): torch.cuda.synchronize() return (time.perf_counter() - t0) / rep * 1000 # ms device cuda if torch.cuda.is_available() else cpu b, n, m, p 32, 256, 256, 256 A torch.randn(b, n, m, devicedevice) B torch.randn(b, m, p, devicedevice) print(bmm :, bench(lambda: torch.bmm(A, B))) print(matmul :, bench(lambda: torch.matmul(A, B))) print(einsum :, bench(lambda: torch.einsum(bij,bjk-bik, A, B))) print(loop :, bench(lambda: torch.stack([A[i] B[i] for i in range(b)])))warmup 不能省。第一次调用某个形状的矩阵乘时底层库要做算法选择、分配缓存耗时会明显偏高。不做 warmup 直接计时测出来的数字完全没有参考价值。GPU 上还要加torch.cuda.synchronize()因为核函数是异步提交的不同步的话计时只量到了“提交任务”的时间不是“算完”的时间。6.2 结果解读在 GPU 上bmm和matmul的耗时基本一致因为matmul在批量维严格相等时会直接转调bmm没有额外开销。einsum稍慢一点大概 5% 到 15%主要是表达式解析和中间调度的开销。loop版本慢一个数量级以上因为 32 次独立的核函数启动每次都有调度延迟。在 CPU 上结论不同小矩阵比如 32×32时loop版本可能反而更快因为每次调用的矩阵小多线程调度的固定开销占比高批量 API 的优势体现不出来。矩阵尺寸超过 128 之后批量 API 才开始明显领先。这个交叉点在不同 CPU 上不一样你的机器上具体是多少跑一下就知道别照搬别人的数字。还有一个容易被忽略的因素是非连续张量。如果A是permute出来的.is_contiguous()是Falsebmm在内部要先拷一份连续版本耗时可能翻倍。做 benchmark 时要区分连续和非连续两种情况否则会得出“某个 API 很慢”的错误结论实际上是输入布局的问题。修法是提前.contiguous()把拷贝成本单独承担后续多次复用就摊薄了。具体到工程选择我的排序是核心路径上的固定形状批量矩阵乘用bmm语义明确、性能稳定维度不固定或者需要广播的用matmul涉及三个以上张量缩并、或者表达式本身就复杂比如注意力里的QK^T之后再乘V用einsum换可读性必要时加optimize至于手写循环只在 batch 很小、调试阶段用上生产前替换掉。我个人的体会是矩阵运算的坑大多不在“不会写”而在“写错了但不报错”。广播把形状不匹配变成静默计算reshape把内存顺序问题藏起来float32把精度损失摊薄到你察觉不到直到某一天 loss 曲线诡异地上飘或者和参考实现比对时差了 1e-3才开始逐个环节回查。所以我现在的习惯是每个涉及矩阵运算的模块写完先构造一个手算得出来的小例子3×3 的矩阵足够把期望结果和实际结果打出来对一遍再放大到真实尺寸。这一步多花十分钟能省掉后面几天的排查。另外凡是形状变化的地方加一行assert或者注释写清楚变换前后的形状也是我自己强制执行的一条规矩——过两周再回来看代码那行注释比任何记忆都可靠。 
长河君 · 2026/10/26 3:09:46 · 阅读 1.2k PyTorch张量矩阵运算：乘法API、广播机制与数值稳定性
1. 先把“矩阵运算”这件事的边界划清楚我第一次被人问“PyTorch张量怎么做矩阵运算”对方的真实困惑其实不是不会写而是不知道自己写的那个到底调到了哪一层为什么有时候形状对得上、有时候对不上为什么同样的数学式子换成另一种写法速度快了一倍。这个问题在项目里被放大得很厉害小规模 demo 里for循环累加完全跑得动一旦数据量上来同样的逻辑换成批量矩阵运算耗时能从几秒掉到几十毫秒差的不是机器性能是你有没有用对 API。所以这篇内容不讲“张量是什么”这类从零起步的概念而是聚焦在PyTorch 张量上的矩阵运算这一个点上把它拆成三块乘法族的 API 到底各管什么、非乘法的线性代数操作求逆、解方程、分解、范数怎么用、以及形状/精度/性能这些真正会在项目里咬人的细节。适合已经能写出简单模型、但在写自定义损失、实现注意力、做批量相似度计算时被形状报错和性能问题卡住的人也适合从 NumPy 转过来、发现两边行为不完全一致的同学。以下所有结论都是我在实际项目和反复 benchmark 中得到的凡是涉及参数选择的我都会把推演过程写出来方便你自己判断。1.1 标量、向量、矩阵、张量在代码里到底差在哪数学上矩阵是二维的张量是高维的但在 PyTorch 里它们全都是torch.Tensor区别只体现在shape和dim()上。一个shape(3,4)的张量是矩阵shape(3,)的是向量shape(2,3,4)的是三维张量。这个“统一表示”带来一个副作用你没法靠类型系统区分“这是个矩阵”还是“这是个批量矩阵”只能靠 API 的约定去约束。关键就在于batch 维度的位置约定。PyTorch 里几乎所有批量矩阵 API 都默认批量维在前矩阵的两个维度在后。比如torch.bmm(A, B)要求A的形状是(b, n, m)、B是(b, m, p)输出(b, n, p)。这个约定不是随便定的它和内存布局有关连续张量的最后一个维度步长最小批量维在最前面意味着每个 batch 的矩阵是连续存放的一整块GPU 上做批量 GEMM 时可以一次性调度效率最高。很多人踩的坑是把 batch 维放到中间比如(n, b, m)然后硬套bmm结果报错或者拿到错误结果。这时候要么permute一下变成(b, n, m)要么换matmul让它走广播。permute回来的是非连续张量后续.view()会失败得用.reshape()或者先.contiguous()——这是我在处理变长序列时被教育过好几次的地方。还有一个容易被忽略的点PyTorch 里没有“行向量”和“列向量”的区分。shape(3,)的张量在参与matmul时会根据另一侧的形状自动决定是当行向量还是列向量用。这个规则在torch.matmul文档里写得很清楚但实际读代码时经常看不出来尤其是当两边都是一维的时候——torch.matmul(a, b)对两个一维张量做的是点积返回标量而torch.outer(a, b)返回的是外积矩阵。同样的两个输入函数不同结果维度完全不同这是新手最容易混淆的地方之一。1.2 为什么 PyTorch 要搞出那么多乘法函数torch.mm、torch.mv、torch.dot、torch.bmm、torch.matmul、torch.einsum、torch.outer光乘法就有七八个入口。乍看是历史包袱实际上每一个都对应一种明确的形状契约用对了能少写很多判断用错了报错信息也更直接。torch.mm只接受二维张量不接受广播不做任何形状兜底。这个“严格”恰恰是它的价值在代码 review 时看到mm你就能确定两个操作数都是矩阵不会出现“某个分支下传进来的是一维张量”的隐藏风险。torch.bmm同理严格三维不广播。把这种严格 API 用在核心计算路径上相当于给自己加了一道类型检查。torch.matmul是“万能版”它内部会做形状规约和广播能处理一维、二维、多维的任意组合。方便是方便代价是错误被推迟了形状不匹配的时候报错信息往往指向matmul内部提示“size mismatch”但真正的 bug 可能在十行之前就已经埋下。我自己的习惯是核心路径优先用严格 API只有在确实需要广播或者维度不固定的地方才用matmul这样出错时定位成本低得多。torch.einsum是另一类东西它不针对某一种形状而是用爱因斯坦求和约定描述“哪些维度参与缩并、哪些保留”。一个einsum表达式可以替代多重循环、permute bmm、diagonal、trace等一大堆操作。它的缺点是字母标记容易写错、初学者读起来费劲而且早期版本的einsum在某些形状下性能不如手写bmm。不过近几年优化已经好很多尤其是配合optimizeoptimal参数时它会自动搜索最优的缩并顺序。2. 乘法族 API 的能力边界与形状推演这一节把乘法相关的函数放在一起对照重点不是背函数签名而是建立一个判断流程看到两个张量要相乘先问自己三个问题——各是几维、需不需要广播、批量维在哪。这三个问题回答完用哪个函数基本就定了。2.1 mm、mv、dot、bmm、matmul 的适用条件对照把常用乘法 API 整理成一张表方便随时查。这里列的是“能不能用”不是“推荐用”实际选择还要结合可读性和性能。API输入维度要求是否广播输出维度典型用途torch.dot两个一维长度相同否标量向量内积、相似度torch.mv矩阵 一维向量否一维线性变换y Axtorch.outer两个一维否二维矩阵外积、秩一更新torch.mm两个二维否二维标准矩阵乘torch.bmm两个三维批量维相同否三维批量矩阵乘torch.matmul/任意遵循广播与一维规约是视输入而定通用乘法torch.einsum任意按下标约定否但可表达批量由表达式决定复杂缩并torch.tensordot任意指定缩并维否由dims决定高阶张量缩并torch.dot有一个限制经常被忽略它只接受一维张量且要求dtype一致。如果你拿一个shape(3,)和一个shape(3,1)去做dot会直接报错因为后者是二维。这类错误在 batch size 为 1 的时候特别隐蔽因为(3,1)在肉眼上和(3,)看起来差不多但shape完全不同。我在调试损失函数时遇到过一次原因就是某次unsqueeze多了一个维度导致内积变成了矩阵乘输出形状从标量变成(3,1)广播到后面全部算错但因为有广播程序不报错只是 loss 曲线莫名其妙地不对。torch.mv的第二个参数必须是一维。如果你手上是(n, 1)的列向量得先squeeze(-1)。反过来matmul就没这个限制(n, m) (m, 1)得到(n, 1)(n, m) (m,)得到(n,)两种都能跑但结果形状不同。这个差异在下游要做broadcast的时候会有实际影响比如后面接一个(n,)的 bias 相加前者(n,1)(n,)会广播成(n,n)后者(n,)(n,)是正确的。这类“静默广播错误”是矩阵运算里最难查的一类 bug因为程序不报错、结果也不崩就是数值不对。torch.bmm的批量维必须严格相等不做广播。如果两个张量的批量维一个是 1、一个是 32bmm会报错这时候要么expand成 32要么换matmul。expand出来的是共享内存的视图不占额外显存但它要求被扩展的维度大小是 1且后续如果是 in-place 操作会出问题。我一般直接用matmul让它自己处理广播省得纠结。torch.matmul的一维规约规则值得单独记一下如果第一个参数是一维会在它前面补一个 1 变成二维如果第二个参数是一维会在它后面补一个 1算完之后再把补的维度去掉。所以(3,) (3,)其实是(1,3) (3,1) (1,1)去掉补的维变成标量。(3,) (3,4)是(1,3) (3,4) (1,4)去掉前面的 1 得到(4,)。这条规则理解了很多“为什么结果维度少了”的困惑就解决了。2.2 广播机制matmul 真正强的地方广播这个词听起来抽象其实就是“维度大小对不上时如果有一边是 1就自动复制成另一边的大小”。规则是从最后一个维度开始往前对齐每一维要么相等、要么其中一个是 1、要么其中一方没有这一维否则报错。举个具体例子A形状(32, 1, 8, 16)B形状(16, 32)。先对齐B是二维A是四维B前面补两个 1 变成(1, 1, 16, 32)。逐维比较最后一维16vs32不相等且都不是 1报错。如果B是(16, 16)补完是(1, 1, 16, 16)最后一维16 vs 16通过倒数第二维8 vs 16又不通过。所以要让这个乘法成立B的后两个维度得是(8, x)或者(1, x)配合A的8。实际编码中广播最常见的用法是“单头 vs 多头”。比如你有一个(heads, head_dim, head_dim)的旋转矩阵要作用到(batch, heads, seq, head_dim)的张量上。直接matmulPyTorch 会把旋转矩阵前面补一个 1 变成(1, heads, head_dim, head_dim)然后和(batch, heads, seq, head_dim)做批量乘。这里1和batch广播、heads对齐、seq和head_dim对齐最后一个维度做矩阵乘输出(batch, heads, seq, head_dim)。整个式子写出来就是一行x R.transpose(-1, -2)比手写循环清爽太多。需要注意的是广播出来的张量是虚拟扩展不真的分配内存但一旦参与计算硬件层面可能还是会按扩展后的大小去读。expand之后连续做多次运算性能不一定比repeat好因为访存模式是重复读取同一块。我实测过(1, 64, 128, 128)扩展成(32, 64, 128, 128)再matmul的情况和直接repeat再算在小尺寸下差别不明显在seq_len上千的时候expand版本反而略慢原因是缓存局部性差。所以“expand 省显存”是对的但别默认它一定更快。还有一类坑是“意外的标量广播”。(n, 1) (n,)会得到(n, n)这个前面提过。更隐蔽的是(n, 1) * (1, m)得到(n, m)如果本意是逐元素对齐就会算出一个完全错误的矩阵。防御办法很简单在关键位置assert一下形状尤其是损失函数、mask 处理这类地方。一行assert x.shape y.shape能省掉半夜调参的时间。2.3 用 einsum 统一表达所有乘法形态einsum的核心是下标字符串输入张量各用一个字母序列标注维度箭头右边写输出保留的维度没出现在右边的字母就是被缩并掉的。ij,jk-ik是矩阵乘i,i-是内积i,j-ij是外积ii-是迹ij-ji是转置。它最大的好处是把“形状变换”和“计算”合成一句代码的数学含义直接可读。比如实现多头注意力的打分Q形状(b, h, s, d)K形状(b, h, t, d)要算(b, h, s, t)的分数矩阵。用matmul的写法是Q K.transpose(-1, -2)用einsum是torch.einsum(bhsd,bhtd-bhst, Q, K)。后者一眼就能看出d被缩并掉了输出的四个维度分别来自哪里不需要在脑子里做转置推演。再看一个matmul不太好表达的批量加权求和。feat形状(b, n, d)weight形状(b, n)要对每个样本按权重把n个向量加权平均成(b, d)。matmul写法是(weight.unsqueeze(1) feat).squeeze(1)中间要插维度、再去掉。einsum写法是torch.einsum(bn,bnd-bd, weight, feat)直接表达“沿 n 缩并”。可读性和出错概率都低不少。性能上einsum内部会把它翻译成bmm、permute、sum的组合。默认的缩并顺序是按表达式从左到右不一定是全局最优。加optimizeoptimal会让它枚举所有缩并顺序选计算量最小的那个——代价是第一次调用有额外的规划开销适合形状固定的场景比如固定的模型结构。如果是动态形状每次调用形状都变规划开销会抵消收益这时候用默认或者optimizegreedy更稳妥。有一个einsum的坑要说清楚它不做广播。所有出现在表达式里的维度输入张量必须有大小必须严格匹配。bhsd,bhtd-bhst里如果某个输入的h是 1、另一个是 12einsum会报错而matmul会自动广播。所以在写可变头数的代码时要自己处理单头扩展的情况不能指望einsum兜底。3. 非乘法类矩阵运算分解、求逆与解方程乘法只解决“乘”的问题实际项目里更多时候要解方程、求逆、做分解。PyTorch 把这部分统一放在torch.linalg命名空间下接口和 NumPy 的linalg基本对齐行为上有一些差异需要注意。3.1 求逆是最不该用的操作torch.linalg.inv能用但绝大多数情况下不该用。原因有三层一是数值稳定性显式求逆再相乘的误差通常比直接解方程大一个量级二是性能解方程是O(n^3)的分解加O(n^2)的回代求逆是O(n^3)加一次完整的矩阵乘常数更大三是显存n×n的逆矩阵要额外占一块空间而解方程只需要一个n×n的分解缓存。正确的写法是torch.linalg.solve(A, B)直接解AX B。它在内部做 LU 分解对称正定时走 Cholesky速度更快然后回代不显式构造逆矩阵。如果确实需要逆矩阵本身比如要拿它做后续的多次乘法可以用torch.linalg.inv_ex或者先lu_factor再lu_solve把分解结果复用起来。我在一个做高斯过程回归的模块里踩过这个坑。原代码是K_inv torch.linalg.inv(K)然后alpha K_inv y在小数据上没问题n到 2000 的时候开始出现loss nan。排查下来是K的条件数接近 1e12求逆过程中小奇异值被放大数值彻底崩掉。改成torch.linalg.solve(K, y)并加了一点点jitter对角线加1e-6问题消失。这里的关键不是solve有多神而是它避免了显式构造一个数值上很脆弱的中间量。torch.linalg.pinv用在非方阵或者秩亏的情况下内部走 SVD速度比solve慢不少但鲁棒。它的rcond参数控制小奇异值的截断阈值默认值是max(m, n) * eps实际用的时候通常要手动调大比如1e-5否则噪声方向会被当成有效信息保留下来反而有害。3.2 解线性方程组、最小二乘与矩阵分解除了solve还有几个必须知道的入口。torch.linalg.lstsq(A, B)直接解最小二乘min ||AX - B||不用你手动写正规方程。这一点很重要因为正规方程(A^T A)^{-1} A^T B会把条件数平方精度损失严重。lstsq内部走 QR 或者 SVD稳定性好很多。它的返回值和 PyTorch 版本有关新版返回一个 namedtuple包含solution、residuals、rank、singular_values取结果时注意用属性名而不是位置索引避免版本升级后出错。torch.linalg.qr做 QR 分解torch.linalg.svd做奇异值分解torch.linalg.eigh做对称矩阵的特征分解torch.linalg.cholesky做 Cholesky 分解。这几个是构建更复杂算法的基石比如 PCA 就是中心化之后取 SVD 的前 k 个奇异向量。svd有一个参数full_matrices默认True返回完整的U和Vh对于(m, n)且m n的矩阵U是(m, m)显存开销巨大。做降维的时候设成False只返回经济型分解U是(m, n)能省下大量显存。这个参数我在处理图像特征矩阵时吃过亏一个(50000, 128)的矩阵full_matricesTrue直接 OOM改成False之后轻松跑完。cholesky只对对称正定矩阵有效但它比solve快将近一倍而且能顺便判断正定性——如果矩阵不正定分解会直接报错。做协方差矩阵相关的运算时先试cholesky失败再退回solve是一个很实用的策略。3.3 范数、条件数与秩判断矩阵“好不好用”在动手算之前先花一行代码看看矩阵的“体质”能省掉大量调试时间。torch.linalg.norm(A, ord...)支持各种范数矩阵场景下常用ord2谱范数即最大奇异值、ordfroF 范数所有元素平方和开根、ord1列和最大值、ordtorch.inf行和最大值。做梯度裁剪的时候用的是向量范数ord2就是把所有元素平方求和开根这个和矩阵的谱范数不是一回事别混。torch.linalg.cond(A)算条件数等于最大奇异值除以最小奇异值。经验值小于 1e3 很安全1e3 到 1e6 要留意超过 1e8 基本可以认为在 float32 下不可逆超过 1e12 在 float64 下也悬。我一般会在涉及求逆或解方程的模块里加一个可选的 debug 分支训练前跑一次cond数值大了就报警比等到 loss 变 nan 再回头查高效得多。torch.linalg.matrix_rank(A)算数值秩内部也是走 SVD阈值同样由rcond控制。判断一个矩阵是不是满秩、要不要加正则项用它比看det靠谱——行列式接近 0 不代表一定奇异行列式很大也不代表条件数好这两个指标经常误导人。4. 从零跑一遍线性回归的正规方程与向量化改造理论讲完跑一段能直接复现的代码。这一段的目标不是训一个多好的模型而是把“同一个数学式子有四种写法性能和稳定性各不相同”这件事演示清楚。4.1 环境与数据构造不需要 GPU 也能跑CPU 上足够看出差异。import torch import time torch.manual_seed(42) torch.set_default_dtype(torch.float32) n_samples, n_features 5000, 64 X torch.randn(n_samples, n_features) true_w torch.randn(n_features, 1) y X true_w 0.01 * torch.randn(n_samples, 1) print(X.shape, y.shape) # torch.Size([5000, 64]) torch.Size([5000, 1])数据构造里有一个细节值得说y的形状是(n_samples, 1)而不是(n_samples,)。这个选择会影响后面所有形状推演。如果y是一维的X.t() y得到(64,)而X.t() X是(64, 64)解出来的w是(64,)如果y是二维的w是(64, 1)。两种都能用但混着用就会出问题比如预测值X w的形状在一维w下是(n,)二维w下是(n, 1)和y相减时如果y是(n, 1)、预测是(n,)会广播成(n, n)loss 直接爆炸。我的习惯是回归任务统一用二维(n, 1)全程形状一致少一类 bug。4.2 用 inv 硬算与用 solve 稳算的对比正规方程的数学形式是w (X^T X)^{-1} X^T y。三种实现从“最直译”到“最推荐”。XTX X.t() X # (64, 64) XTy X.t() y # (64, 1) # 写法一显式求逆数学上最直译实际最不该用 w_inv torch.linalg.inv(XTX) XTy # 写法二解方程稳定性好速度快 w_solve torch.linalg.solve(XTX, XTy) # 写法三直接最小二乘不走正规方程条件数不被平方 w_lstsq torch.linalg.lstsq(X, y).solution三者结果应该非常接近因为数据是良态的。为了看出差异给X加一点共线性把第 2 列设成第 1 列加一点噪声人为把条件数拉高。X_bad X.clone() X_bad[:, 1] X_bad[:, 0] 1e-4 * torch.randn(n_samples) print(cond , torch.linalg.cond(X_bad.t() X_bad).item())条件数到 1e8 以上时inv版本的结果会明显偏离solve稍微好一点lstsq最稳。这个实验我建议你自己跑一遍把绝对误差打出来看比看任何文字描述都有说服力。原因在于X^T X的条件数是X的条件数平方本来 1e4 的条件数经过平方变成 1e8float32 只有大约 7 位有效十进制数字误差直接吃掉了大部分精度。这也是为什么在特征维度高、特征之间相关性强的场景下宁可直接用lstsq或者加 L2 正则把XTX换成XTX λI也不要硬解正规方程。加正则之后的矩阵条件数会被λ拉低λ取1e-4量级就能把 1e8 降到 1e4 左右效果立竿见影。4.3 把三重循环改写成 einsum前一个例子是一步矩阵乘体现不出向量化的价值。换一个更真实的场景计算一批样本两两之间的加权相似度。假设有b个批次每批n个向量每个向量d维另有一个d×d的度量矩阵M要算S[i,j] x_i^T M x_j。朴素写法是三重循环def pairwise_naive(Xb, M): b, n, d Xb.shape out torch.zeros(b, n, n) for bi in range(b): for i in range(n): for j in range(n): out[bi, i, j] Xb[bi, i] M Xb[bi, j] return outn200、b8的时候这个循环要跑 32 万次内层迭代Python 层面开销巨大。改写分两步先把M吸收进X即Z X M形状(b, n, d)然后S Z X.transpose(-1, -2)形状(b, n, n)。用einsum表达就是def pairwise_einsum(Xb, M): Z torch.einsum(bnd,de-bne, Xb, M) return torch.einsum(bnd,bmd-bnm, Z, Xb)注意einsum的第二个表达式里两个输入的d是同一个字母表示缩并输出的n和m分别来自两个输入所以得到(b, n, n)。这里用不同字母n和m而不是都用n是为了避免einsum误解成“对角”或者“逐元素”——同一个字母出现在两个输入的非缩并位置语义是“这两个维度必须相等且逐元素对应”不会产生两两组合。这个细节很多人第一次写会搞错写成bnd,bnd-bnn得到的结果完全不同。实测下来n200, b8, d64朴素循环大约 1.8 秒einsum版本 3 毫秒左右差了两个数量级。而且循环版本还有一个隐患out[bi, i, j] ...是原地赋值如果out参与自动求导这种逐元素写入的方式在某些情况下会导致梯度计算变慢或者报错向量化版本完全没有这个问题。5. 形状、dtype、device 与数值精度的高频塌方点前面讲的是“怎么写对”这一节讲“为什么明明写对了还是出错”。矩阵运算的报错信息通常很简短size mismatch、expected scalar type、device mismatch但根因经常在很远的地方。5.1 形状报错的定位套路遇到形状报错别急着改代码先做三件事。第一把报错行涉及的所有张量形状打出来包括中间变量不要只打最终的两个操作数。很多时候根因是某个unsqueeze加错了位置或者某个sum少了dim参数导致维度凭空多出来一个。第二用matmul的规则手动推一遍广播。从最后一维往前对齐逐维检查把每一维的来源写下来。这个过程慢但能定位到具体哪一维不匹配。比如报错说Expected size for argument #2 mat2 (64) to be 32那就是第二个操作数的倒数第二维是 64第一个操作数的最后一维是 32它们要做矩阵乘必须相等。第三检查有没有“隐式 squeeze”。torch.sum(x, dim0)会真正去掉那一维而torch.sum(x, dim0, keepdimTrue)保留成 1。代码里混用这两种写法形状就会飘。我现在的习惯是凡是在流程中间做 reduction一律keepdimTrue到最后确定要输出的时候再显式squeeze这样形状变化都有迹可循。还有一种情况是 reshape 顺序问题。(b, n, d)直接reshape(b*n, d)是对的但如果中间做过transpose变成(b, d, n)再reshape(b*n, d)就会把内存顺序搞乱得到错误的数据不报错但结果全错。判断方法是用.is_contiguous()检查或者干脆用permute之后再.reshape()reshape会在必要时自动拷贝一份保证语义正确代价是多一次内存分配。这个坑在实现 Transformer 的多头拆分时特别常见(b, s, h*d)拆成(b, s, h, d)再转成(b, h, s, d)如果不用reshape而用view或者转置顺序写错后面所有注意力计算全是垃圾数据。5.2 dtype 与 device 不一致的连锁反应float32和float64混用矩阵乘会直接报错PyTorch 不做隐式升级。这一点和 NumPy 不同NumPy 会自动提升到更高精度。PyTorch 的选择是明确的混精度必须显式转换避免不知不觉中损失性能。排查套路是在报错处打印x.dtype和y.dtype。常见来源有三个一是数据加载时numpy默认float64torch.from_numpy之后没转二是某个常数写成了 Pythonfloat参与运算时被当成float32但如果另一边是float64就报错三是torch.zeros没指定dtype用了全局默认值而全局默认被某处set_default_dtype改过。device 不一致的报错更直接Expected all tensors to be on the same device。麻烦的是多卡场景模型在某张卡上数据在另一张卡上或者某个 buffer 没有跟着.to(device)走。我的做法是写一个工具函数在训练循环的第一步检查所有关键张量的 device 和 dtype一次性打印出来而不是等报错出现再逐个查。混合精度训练时还要注意autocast上下文里的矩阵乘会被自动转成半精度但如果手动调用了torch.linalg下的函数有些操作在半精度下不支持会报错或者悄悄降级。linalg.inv、linalg.svd这类操作在半精度下精度损失严重建议显式float()之后再算算完再转回来。5.3 浮点精度、TF32 与混合精度float32有 24 位有效尾数大约 7 位十进制精度。矩阵乘是累加操作d维内积累加d次误差大致按sqrt(d)增长。d1024时相对误差可能在 1e-5 量级对大多数任务够用但对需要高精度对齐的场景比如和参考实现做数值比对就会出问题。这时候要么用float64要么把比对阈值放宽到 1e-4。GPU 上还有一个额外变量TF32。这是较新架构的 GPU 在矩阵乘上默认开启的一种折中精度格式尾数只有 10 位比float32少很多但比float16多。它能显著提升矩阵乘吞吐代价是精度下降。PyTorch 里由torch.backends.cuda.matmul.allow_tf32控制不同版本默认值变过做数值复现实验时要显式设成False否则你会发现同样的代码在两张不同卡上结果不一致。混合精度里float16的动态范围小容易出现下溢。做矩阵乘之前如果输入的量级差异大建议用bfloat16它的指数位和float32一样宽不容易溢出代价是尾数更短、精度更低。实际项目里通常是bfloat16更省心尤其是在做注意力里的 softmax 之前float16经常需要把分数先缩放一下防溢出bfloat16基本不用管。5.4 原地操作与自动求导的冲突矩阵运算里最容易触发自动求导报错的是 in-place 操作。比如x.t_()、x.add_()、x[:, 0] 0如果x是叶子张量或者被其他计算依赖autograd 会报a leaf Variable that requires grad is being used in an in-place operation或者更隐蔽的版本号不匹配错误。t_()和transpose_()在较新版本里已经不推荐用了因为它们改变的是 stride 而不是数据和后续的视图、求导混在一起容易出问题。统一用t()和transpose()返回新视图虽然多了一个 Python 对象但没有副作用。填充 mask 的时候x.masked_fill_(mask, 0)是原地操作如果x需要梯度会报错。改成x.masked_fill(mask, 0)返回新张量就行或者在torch.no_grad()上下文里做原地填充。后者用在推理阶段更省显存训练阶段老老实实用非原地版本。6. 性能实测不同 API 的速度差异从哪来讲性能不能空谈得有数据。下面这段 benchmark 脚本可以直接复制运行CPU 和 GPU 都适配。6.1 一个简单的 benchmark 脚本import torch import time def bench(fn, warmup10, rep50): for _ in range(warmup): fn() if torch.cuda.is_available(): torch.cuda.synchronize() t0 time.perf_counter() for _ in range(rep): fn() if torch.cuda.is_available(): torch.cuda.synchronize() return (time.perf_counter() - t0) / rep * 1000 # ms device cuda if torch.cuda.is_available() else cpu b, n, m, p 32, 256, 256, 256 A torch.randn(b, n, m, devicedevice) B torch.randn(b, m, p, devicedevice) print(bmm :, bench(lambda: torch.bmm(A, B))) print(matmul :, bench(lambda: torch.matmul(A, B))) print(einsum :, bench(lambda: torch.einsum(bij,bjk-bik, A, B))) print(loop :, bench(lambda: torch.stack([A[i] B[i] for i in range(b)])))warmup 不能省。第一次调用某个形状的矩阵乘时底层库要做算法选择、分配缓存耗时会明显偏高。不做 warmup 直接计时测出来的数字完全没有参考价值。GPU 上还要加torch.cuda.synchronize()因为核函数是异步提交的不同步的话计时只量到了“提交任务”的时间不是“算完”的时间。6.2 结果解读在 GPU 上bmm和matmul的耗时基本一致因为matmul在批量维严格相等时会直接转调bmm没有额外开销。einsum稍慢一点大概 5% 到 15%主要是表达式解析和中间调度的开销。loop版本慢一个数量级以上因为 32 次独立的核函数启动每次都有调度延迟。在 CPU 上结论不同小矩阵比如 32×32时loop版本可能反而更快因为每次调用的矩阵小多线程调度的固定开销占比高批量 API 的优势体现不出来。矩阵尺寸超过 128 之后批量 API 才开始明显领先。这个交叉点在不同 CPU 上不一样你的机器上具体是多少跑一下就知道别照搬别人的数字。还有一个容易被忽略的因素是非连续张量。如果A是permute出来的.is_contiguous()是Falsebmm在内部要先拷一份连续版本耗时可能翻倍。做 benchmark 时要区分连续和非连续两种情况否则会得出“某个 API 很慢”的错误结论实际上是输入布局的问题。修法是提前.contiguous()把拷贝成本单独承担后续多次复用就摊薄了。具体到工程选择我的排序是核心路径上的固定形状批量矩阵乘用bmm语义明确、性能稳定维度不固定或者需要广播的用matmul涉及三个以上张量缩并、或者表达式本身就复杂比如注意力里的QK^T之后再乘V用einsum换可读性必要时加optimize至于手写循环只在 batch 很小、调试阶段用上生产前替换掉。我个人的体会是矩阵运算的坑大多不在“不会写”而在“写错了但不报错”。广播把形状不匹配变成静默计算reshape把内存顺序问题藏起来float32把精度损失摊薄到你察觉不到直到某一天 loss 曲线诡异地上飘或者和参考实现比对时差了 1e-3才开始逐个环节回查。所以我现在的习惯是每个涉及矩阵运算的模块写完先构造一个手算得出来的小例子3×3 的矩阵足够把期望结果和实际结果打出来对一遍再放大到真实尺寸。这一步多花十分钟能省掉后面几天的排查。另外凡是形状变化的地方加一行assert或者注释写清楚变换前后的形状也是我自己强制执行的一条规矩——过两周再回来看代码那行注释比任何记忆都可靠。