这篇教程本来是 PyTorch 系列的第 2 课上一课讲完环境搭建和 Tensor 的创建、索引、变形、类型转换之后接下来自然要进入运算部分。张量运算直接决定了你后面能不能看懂模型源码里的attention、loss、embedding这些实现。很多人卡在张量形状不对、维度对不上、矩阵乘法报错、广播机制看不懂上本质原因就是没有把逐元素计算、矩阵乘法、广播机制这三块规则一次理清。这次我们把这三件事放在一起讲配可运行的代码你把代码复制到自己的环境里跑一遍应该比只看不练效果好得多。本课核心内容逐元素计算的算术运算符和函数torch.mm、torch.matmul、三种矩阵乘法的区别广播机制的规则和常见报错GPU 上的张量运算验证方式以及 PyTorch 2.6 之后加载模型时weights_only参数变化对张量操作的影响。老规矩文章里所有代码都默认你已经完成了 PyTorch 环境安装CPU 版本就能跑通不需要 GPU代码里涉及 GPU 的部分我会单独标注。1. 本课核心能力速览能力项说明涉及版本PyTorch 2.x 通用示例可在 CPU 环境下直接运行硬件要求CPU 即可运行GPU 可选核心内容逐元素计算、矩阵乘法、广播机制需要掌握的前置知识Tensor 创建、shape 查看、设备转移cpu()/cuda()本课代码量6 段完整可运行示例常见报错覆盖维度不匹配、广播不兼容、CPU/GPU 设备不一致适合读者刚装好 PyTorch、想真正搞懂张量运算规则的同学从实际教学经验看学张量运算最忌讳的就是只背 API 名字不理解形状变化。你可以跟着本课把每一段代码的输入shape和输出shape自己打印一遍先确认形状规律再去看内部数值。2. 逐元素计算加减乘除与比较运算2.1 基础算术运算逐元素计算指的是两个张量相同位置上的元素独立参与运算结果张量形状与输入张量形状保持一致。最典型的就是加、减、乘、除以及幂运算、取模、整除。import torch # 两个相同形状的张量 a torch.tensor([1.0, 2.0, 3.0]) b torch.tensor([4.0, 5.0, 6.0]) # 加减乘除 print(a b) # tensor([5., 7., 9.]) print(a - b) # tensor([-3., -3., -3.]) print(a * b) # tensor([ 4., 10., 18.]) print(a / b) # tensor([0.2500, 0.4000, 0.5000]) # 幂运算和取余 print(a ** 2) # tensor([1., 4., 9.]) print(b % 2) # tensor([0., 1., 0.]) print(torch.pow(a, 2)) # 等价于 a ** 2注意一个细节a * b是逐元素乘法不是矩阵乘法。矩阵乘法要使用torch.mm、torch.matmul或这一点初学者最容易混淆。实际阅读 PyTorch 源码时经常会看到张量前面有*运算符如果参与运算的两个张量形状是[batch, seq, hidden]那这个*就是在做逐元素乘常用来实现注意力权重的掩码乘法而或torch.matmul才是在做特征维度的线性变换。再看除法/和torch.div等价。PyTorch 中的浮点数除法不会因为除零直接崩溃但会出现inf或nan这在后续计算 loss 时会造成梯度异常。实际工程中用torch.where或者加一个极小值eps来避免除零是常见操作思路本课先掌握基本规则训练场景下的数值稳定性后续再单独讨论。2.2 幂运算、指数、对数与常用激活函数除了基础的加减乘除模型代码里出现频率更高的是torch.exp、torch.log、torch.sqrt、torch.square这类逐元素数学函数。softmax 和 sigmoid 的实现本质就是这些逐元素函数的组合。import torch x torch.tensor([0.0, 1.0, 2.0]) print(torch.exp(x)) # e^x输出 tensor([1.0000, 2.7183, 7.3891]) print(torch.log(torch.exp(x))) # x print(torch.square(x)) # 平方 print(torch.sqrt(torch.tensor([4.0, 9.0, 16.0]))) # 平方根 print(torch.clamp(x, min1.0, max2.0)) # 裁剪到 [1,2]torch.clamp在梯度裁剪和值域限制中非常常见它也是逐元素操作。PyTorch 的逐元素运算都支持out参数可以把结果写入预设张量减少内存拷贝但是日常脚本里用得不多大多数情况下直接写z a b即可。在神经网络中F.relu(x)本质上就是逐元素计算。你可以用 PyTorch 直接验证import torch import torch.nn.functional as F x torch.tensor([-2.0, -1.0, 0.0, 1.0, 2.0]) print(F.relu(x)) # tensor([0., 0., 0., 1., 2.]) # 手动实现 print(torch.maximum(x, torch.tensor(0.0)))手动实现结果与F.relu一致。理解这一点后面看 Transformer 里的FFN、注意力权重计算就不会觉得公式是魔法。2.3 比较运算与逻辑运算逐元素比较运算返回的是布尔张量形状与输入一致。比较运算在掩码构造、条件筛选、评估指标计算里很有用。import torch a torch.tensor([1.0, 2.0, 3.0]) b torch.tensor([3.0, 2.0, 1.0]) print(a b) # tensor([False, False, True]) print(a b) # tensor([False, True, False]) print(a b) # tensor([False, True, True]) # 布尔张量参与运算 mask a 2 print(a[mask]) # tensor([3.])这里就引出一个关键概念布尔掩码索引。a[mask]只会返回满足条件的元素常用于过滤无用数据比如把样本中标签为-1的无效结果剔除。实际使用中还要注意如果要在梯度计算中保留位置信息一般用torch.where(mask, a, b)而不是直接索引c torch.where(a 2, a, torch.zeros_like(a)) print(c) # tensor([0., 0., 3.])torch.where的三个参数都支持广播所以它也是理解接下来广播机制的绝佳例子。3. 矩阵乘法mm、matmul 与 运算符3.1 二维矩阵乘法 torch.mm矩阵乘法是深度学习中最重要的运算。torch.mm只接受二维张量输入[m, k]和[k, n]输出[m, n]。中间维度必须相同。import torch m1 torch.randn(3, 4) m2 torch.randn(4, 5) out torch.mm(m1, m2) print(out.shape) # torch.Size([3, 5])这与线性层nn.Linear的原理一致。一个nn.Linear(in_features4, out_features5)在做前向传播时本质上就是把输入[..., 4]与权重矩阵[4, 5]做矩阵乘法再加上偏置。理解了torch.mm线性层内部机制就破案了。手动实现一个简化版线性层import torch x torch.randn(2, 4) # 2个样本每个样本4个特征 weight torch.randn(4, 5) # 输出5个维度 bias torch.randn(5) y torch.mm(x, weight) bias # 先乘后加偏置 print(y.shape) # torch.Size([2, 5])这里用的是广播加法torch.mm(x, weight)形状为[2, 5]bias形状为[5]两者相加会自动把 bias 广播到每一行。广播机制让代码干净了很多但如果你不理解广播规则看到[2, 5] [5]不报错反而会觉得莫名奇妙。后面第 4 节会专门讲透。3.2 高维张量的批量矩阵乘法 torch.matmul实际模型训练中很少只有二维矩阵。一个 batch 的输入通常是[batch_size, seq_len, hidden_dim]比如一个包含 8 个句子、每个句子 10 个 token、每个 token 用 768 维向量表示的文本 batch形状就是[8, 10, 768]。torch.matmul是torch.mm的推广它支持二维、高维和混合维度的矩阵乘法。规则可以简单理解为把最后两维当作真正的矩阵做乘法前面的维度按广播规则处理。import torch # 批量矩阵乘法 b1 torch.randn(2, 3, 4) b2 torch.randn(2, 4, 5) out torch.matmul(b1, b2) print(out.shape) # torch.Size([2, 3, 5]) # 高维张量与二维矩阵相乘 b3 torch.randn(2, 3, 4) m torch.randn(4, 5) out2 torch.matmul(b3, m) print(out2.shape) # torch.Size([2, 3, 5])第二个例子中二维矩阵m被广播到最后两维上相当于对 batch 中每一个[3, 4]矩阵都执行torch.mm。Transformer 的Q K^T就是这种批量矩阵运算的典型场景其中Q的形状是[batch, heads, seq, head_dim]K转置后是[batch, heads, head_dim, seq]两者matmul得到[batch, heads, seq, seq]的注意力分数。运算符是torch.matmul的语法糖。以下两种写法完全等价out_a b1 b2 out_b torch.matmul(b1, b2)在实际项目源码里写法比torch.matmul更常见。你只需要注意和*在语义上的天壤之别。3.3 矩阵乘法维度不匹配的常见报错矩阵乘法最常见的错误就是内部维度不一致。当你执行torch.mm(A, B)时如果A.shape[-1] ! B.shape[-2]PyTorch 会直接抛出类似mat1 and mat2 shapes cannot be multiplied的报错。这个报错的解法只有一个方向确认谁是矩阵乘法中的k维度。排查顺序建议如下打印两个张量的shape检查A.shape[-1]和B.shape[-2]是否相等如果要相乘的语义是特征变换通常需要转置其中一个张量用A.T、A.transpose(-1, -2)或A.permute如果涉及 batch 维还需要确认两个张量的 batch 维数是否满足广播条件。下面展示一个三维张量转置配合矩阵乘法的场景import torch q torch.randn(2, 4, 8) # batch2, seq4, dim8 k torch.randn(2, 6, 8) # batch2, seq6, dim8 # 想把 q 和 k 的最后一维做内积需要把 k 转成 [2, 8, 6] k_t k.transpose(-1, -2) scores torch.matmul(q, k_t) print(scores.shape) # torch.Size([2, 4, 6])这个例子是所有自注意力机制的第一步。你把维度想清楚了再看scaled_dot_product_attention的官方文档就会非常顺。4. 广播机制规则、原理与实战4.1 广播的规则广播机制是 PyTorch 张量运算里最容易忽略、也最容易出错的规则。简单说是允许形状不完全相同的张量进行逐元素运算但必须满足两个条件从最后一个维度开始向前比较两个维度要么相等要么其中一个是 1要么其中一个不存在。如果一个维度既不相等也不为 1就报错。这句话听起来简单实际操作时很多初学者还是会懵。我们看一个直接例子import torch a torch.randn(3, 1) b torch.randn(1, 4) c a b print(a.shape, b.shape, c.shape) # torch.Size([3, 1]) torch.Size([1, 4]) torch.Size([3, 4])这里a的形状是[3, 1]b的形状是[1, 4]。从最后一个维度比较最后一个维1和4不相等但一个为 1可以广播倒数第二个维3和1不相等但一个为 1可以广播。所以结果是[3, 4]。如果a是[3, 5]b是[1, 4]最后一个维度5和4既不相同也不为 1直接报错。4.2 标量与张量的广播标量是最典型的广播场景。比如对一个张量做标准化减去均值再除以标准差均值是标量张量是任意形状PyTorch 会自动把标量广播到每一个元素上。import torch x torch.tensor([[1.0, 2.0], [3.0, 4.0]]) mean x.mean() std x.std() # 标量广播到整个张量 normalized (x - mean) / std print(normalized)这里的mean和std都是零维张量与二维张量运算时自动广播。从数学上等价于对每个元素执行相同的变换。更常见的逐行标准化场景中均值张量是[3]而数据是[4, 3]从最后维比较3等于3可以广播数据的前导维4在均值张量里不存在相当于自动补齐为 1也满足广播条件。import torch data torch.randn(4, 3) # 4个样本3个特征 mean data.mean(dim0) # 形状 [3] centered data - mean # 每行都减去对应的均值 print(centered.shape)这个操作在特征标准化中极其常见。mean的形状是[3]data的形状是[4, 3]减号执行逐元素计算核心是“形状不相同但满足广播条件”。如果写成data.mean(dim0, keepdimTrue)得到的mean形状就是[1, 3]结果不变但语义上更明确标明这是一个“行向量”。在更复杂的模型中keepdimTrue可以帮助保持维度结构避免后续矩阵乘法时少了一个维度产生误操作。4.3 广播不兼容的常见错误形如RuntimeError: The size of tensor a (3) must match the size of tensor b (4) at non-singleton dimension 1的报错是广播不兼容的典型现象。出现这种错误时不要看整段代码先打印两个张量的shape然后从最后维开始逐个对比。import torch a torch.randn(3, 4) b torch.randn(3, 3) # 下面这行会报错 # c a ba的形状是[3, 4]b的形状是[3, 3]。最后维4和3不相等且都不为 1所以无法广播。解决方案要么是转置b要么是重采样使特征维一致取决于你的业务语义。还有一种容易误判的情况是只多了一个前导维a torch.randn(3, 4) b torch.randn(2, 3, 4) c a b print(c.shape) # torch.Size([2, 3, 4])这里a是[3, 4]b是[2, 3, 4]a会自动补齐为[1, 3, 4]再广播到[2, 3, 4]。这解决了给 batch 数据统一加偏置的问题。实际模型中这种广播很常用但也要注意别把语义理解错a会被复制 2 份参与每个 batch 的计算。4.4 用 torch.where 和 unsqueeze 控制广播位置有时你希望把形状[4]的向量广播到[3, 4]的每一行可以直接加。但如果你希望把形状[3]的向量广播到[3, 4]的每一列就必须先unsqueeze变成[3, 1]。import torch col_means torch.tensor([1.0, 2.0, 3.0]) # 长度是 3 matrix torch.randn(3, 4) # 错误思路直接加 col_means形状 [3] 和 [3, 4] 无法匹配 # corrected matrix col_means # 会报错 # 正确做法把 col_means 变成 [3, 1] col_means_col col_means.unsqueeze(1) # [3, 1] corrected matrix col_means_col print(corrected.shape)unsqueeze(1)在模型量化、特征归一化、位置编码加法中非常常用。位置编码[seq_len, hidden]要加到输入[batch, seq_len, hidden]上时也需要借助广播如果代码里没有unsqueeze或expand那大概率是形状已经天然匹配。理解广播维度对齐规则比死记 API 更重要。5. 在 GPU 上验证张量运算与资源观察这一节是很多人跳过但实际很有用的内容。PyTorch 在 CPU 上跑起来方便但真实训练和推理大多依赖 GPU。本课不涉及具体训练只讲怎么把张量移到 GPU 上执行运算并观察显存变化。5.1 设备转移与设备一致性import torch # 判断 CUDA 是否可用 print(torch.cuda.is_available()) # 创建张量并转移到 GPU x torch.randn(1000, 1000) if torch.cuda.is_available(): x_cuda x.cuda() print(x_cuda.device) # cuda:0一旦张量在 GPU 上参与运算的所有张量必须也在同一设备上否则会报Expected all tensors to be on the same device。实际排查时最简单的方式是逐个打印.device。在 PyTorch 2.x 中更推荐使用to(device)配合统一变量管理device torch.device(cuda if torch.cuda.is_available() else cpu) x torch.randn(1000, 1000, devicedevice)这比到处写.cuda()更适合工程化。分布式训练和 AMP 混合精度训练中设备管理会更复杂但本课先养成统一写device的代码习惯。5.2 通过 nvidia-smi 观察显存在终端执行nvidia-smi重点关注Memory-Usage。PyTorch 在 GPU 上创建张量时显存占用随张量规模和数量变化。实际观测方法很简单先跑一次小张量运算记录显存基线再创建大张量记录峰值两者差值就是当前张量占用的显存。由于每个人的 GPU 型号、驱动、CUDA 版本不同本课不给出固定数字。你可以用一个统一的测试脚本自行探测import torch import os device torch.device(cuda if torch.cuda.is_available() else cpu) print(fusing device: {device}) x torch.randn(4000, 4000, devicedevice) y torch.randn(4000, 4000, devicedevice) z torch.mm(x, y) print(z.shape) print(done)运行这段代码后再用nvidia-smi对比前后显存变化。要注意的是PyTorch 显存分配带有缓存机制进程结束后显存才会完全释放。如果你在 Jupyter 里反复执行显存可能不会立刻下降这时候需要重启 kernel。5.3 CPU 与 GPU 运算结果一致性张量运算在 CPU 和 GPU 上数学上等价但浮点计算顺序不同可能产生极小数值差异。对大多数场景可以忽略。验证时不要直接用判断 CUDA 与 CPU 的结果而是用torch.allcloseimport torch cpu_result torch.mm(torch.randn(100, 100), torch.randn(100, 100)) gpu_result None if torch.cuda.is_available(): a torch.randn(100, 100, devicecuda) b torch.randn(100, 100, devicecuda) gpu_result torch.mm(a, b).cpu() if gpu_result is not None: print(torch.allclose(cpu_result, gpu_result, atol1e-4))torch.allclose是后续写单元测试和检查模型输出一致性的常用函数。建议从本课开始养成用它判断张量近似相等的习惯不要依赖。6. 张量运算的步数、形状与内存新手常见问题6.1 为什么*和结果完全不同*是逐元素乘法要求两个张量形状可广播是矩阵乘法要求最后两维满足矩阵乘法条件。下面的例子可以直观看出差异import torch a torch.tensor([[1., 2.], [3., 4.]]) b torch.tensor([[5., 6.], [7., 8.]]) print(a * b) # tensor([[ 5., 12.], # [21., 32.]]) print(a b) # tensor([[19., 22.], # [43., 50.]])从数值就可以看出*只做了对应位置相乘做了完整的矩阵乘法。如果模型代码里把这两个运算符混用大概率会出现形状对不上或者结果数值明显异常。排查时优先逐层打印中间结果的shape。6.2 广播导致的内存暴增广播机制在写代码时很舒服但它可能隐式复制张量。比如[1000, 1]和[1, 1000]相加实际结果[1000, 1000]会占用 8MB 左右float32 下这还只是一个小例子。如果第一个维度是几万广播后的内存会非常可观。所以在特征处理时如果某个张量本来就要重复使用可以用expand显式表示import torch base torch.tensor([[1.0], [2.0], [3.0]]) # [3, 1] expanded base.expand(3, 4) print(expanded.shape) # [3, 4]expand不会立即复制数据而是创建视图在后续参与计算时才产生实际数据。它比直接用repeat复制更节省内存。但是要注意expand只能作用于维度为 1 的轴。这又回到了广播规则本身。6.3 形状打印习惯查错时建议打印print(a.shape, b.shape, a.device, b.device)把 shape 和 device 一起打印一次解决两类问题。在 Jupyter 或脚本里都可以这样做。很多时候你花半小时找 bug最后发现只是某个中间张量多了一个维度或者少了一个unsqueeze。7. 本课完整测试脚本建议把下面这段代码保存为一个 Python 文件比如tensor_ops_demo.py在终端里运行python tensor_ops_demo.pyimport torch def main(): # 1. 逐元素计算 a torch.tensor([1.0, 2.0, 3.0]) b torch.tensor([4.0, 5.0, 6.0]) print(a b:, a b) print(a * b:, a * b) # 2. 矩阵乘法 m1 torch.randn(3, 4) m2 torch.randn(4, 5) out torch.mm(m1, m2) print(mm out shape:, out.shape) # 3. matmul 高维 b1 torch.randn(2, 3, 4) b2 torch.randn(2, 4, 5) out2 torch.matmul(b1, b2) print(matmul out shape:, out2.shape) # 4. 广播 x torch.randn(3, 1) y torch.randn(1, 4) z x y print(broadcast out shape:, z.shape) # 5. 设备判断 device torch.device(cuda if torch.cuda.is_available() else cpu) gpu_tensor torch.randn(2, 2, devicedevice) print(gpu_tensor device:, gpu_tensor.device) if __name__ __main__: main()我建议你实际运行后把所有shape打印和注释里的输出对照一遍。如果想进一步巩固可以把b1、b2的形状分别改成[3, 4]、[2, 4, 5]和[4, 9]等不同组合观察什么情况下报错、什么情况下广播成功。8. 常见问题与排查方法问题现象可能原因排查方式解决方案mat1 and mat2 shapes cannot be multiplied矩阵乘法的内部维度k不一致打印两个张量的 shape转置A或修改B的最后一维The size of tensor a (3) must match ...广播不兼容从最后维逐个对比 shape用unsqueeze或reshape对齐维度Expected all tensors to be on the same deviceCPU 与 GPU 张量混合参与运算打印每个张量的.device全部to(device)CUDA out of memory广播后张量过大或缓存未释放观察nvidia-smi显存占用缩小 batch 或重启 kernel结果出现nan或inf除零、log(0)、指数溢出检查输入数据范围使用torch.clamp或加epsCant call numpy() on Tensor that requires grad在需要梯度的张量上调用numpy()检查requires_grad和是否在with torch.no_grad()中使用.detach().cpu().numpy()PyTorch 2.6 加载旧模型报WeightsUnpickler errortorch.load默认weights_onlyTrue查看报错堆栈是否在torch.load设置weights_onlyFalse只对可信模型关于最后一条补充几句。PyTorch 2.6 开始torch.load默认把weights_only设为True这是为了安全性防止加载恶意 pickle 数据。如果你加载的是自己训练的模型遇到报错时可以显式传weights_onlyFalse如果你参与的项目代码还在用旧写法请在升级 PyTorch 时留意这个行为变化。9. 最佳实践与工程建议以下几条不仅适用于本课也适用于后续所有 PyTorch 项目统一用torch.device管理设备不要混用.cuda()和.cpu()在数据加载、模型初始化、前向传播三个环节都统一to(device)避免设备不一致报错。每次写完张量运算先打印shape再继续写下一行。尤其是矩阵乘法前后把中间结果形状记录下来能减少大量“玄学 bug”。对逐元素运算和矩阵运算要区分清楚。需要特征变换用nn.Linear或需要掩码过滤、逐元素加权用*。广播机制能省很多代码但不要把维度隐式地“变没”。如果某个中间张量形状不够直观用unsqueeze或reshape显式声明。涉及 GPU 显存时先用nvidia-smi记录基线再跑大张量。PyTorch 的缓存机制会导致显存不立即释放这不是内存泄漏只是缓存策略必要时候可以用torch.cuda.empty_cache()手动释放。如果本地训练环境下载慢安装 PyTorch 可以优先考虑国内镜像源。这里不展开教程但搜索“PyTorch 国内镜像安装”能很快找到方案。加载旧模型权重时注意 PyTorch 2.6 的weights_only默认值变化。升级版本前先看官方 release notes再批量检查代码里的torch.load调用。10. 下一步学习建议本课完成之后你可以自己设计一个小实验来巩固创建形状为[8, 3, 32]的随机张量模拟 8 个样本、每个样本 3 个时间步、每个时间步 32 维特征。尝试对它分别做逐行均值标准化、批量矩阵乘法、与形状为[32, 16]的线性变换矩阵相乘把所有中间结果的shape打印出来。这个实验能同时覆盖逐元素计算、矩阵乘法和广播机制。做完后你会发现Transformer、TCN、股票预测模型里的很多张量流动本质上就是在重复这三种操作的不同组合。下一课可以继续深入autograd和反向传播理解requires_grad、backward()与计算图的关系。在那之前先把本课的运算规则本地跑熟遇到问题翻一下第 8 节的排查表基本够用。 PyTorch张量运算详解：逐元素计算、矩阵乘法与广播机制
这篇教程本来是 PyTorch 系列的第 2 课上一课讲完环境搭建和 Tensor 的创建、索引、变形、类型转换之后接下来自然要进入运算部分。张量运算直接决定了你后面能不能看懂模型源码里的attention、loss、embedding这些实现。很多人卡在张量形状不对、维度对不上、矩阵乘法报错、广播机制看不懂上本质原因就是没有把逐元素计算、矩阵乘法、广播机制这三块规则一次理清。这次我们把这三件事放在一起讲配可运行的代码你把代码复制到自己的环境里跑一遍应该比只看不练效果好得多。本课核心内容逐元素计算的算术运算符和函数torch.mm、torch.matmul、三种矩阵乘法的区别广播机制的规则和常见报错GPU 上的张量运算验证方式以及 PyTorch 2.6 之后加载模型时weights_only参数变化对张量操作的影响。老规矩文章里所有代码都默认你已经完成了 PyTorch 环境安装CPU 版本就能跑通不需要 GPU代码里涉及 GPU 的部分我会单独标注。1. 本课核心能力速览能力项说明涉及版本PyTorch 2.x 通用示例可在 CPU 环境下直接运行硬件要求CPU 即可运行GPU 可选核心内容逐元素计算、矩阵乘法、广播机制需要掌握的前置知识Tensor 创建、shape 查看、设备转移cpu()/cuda()本课代码量6 段完整可运行示例常见报错覆盖维度不匹配、广播不兼容、CPU/GPU 设备不一致适合读者刚装好 PyTorch、想真正搞懂张量运算规则的同学从实际教学经验看学张量运算最忌讳的就是只背 API 名字不理解形状变化。你可以跟着本课把每一段代码的输入shape和输出shape自己打印一遍先确认形状规律再去看内部数值。2. 逐元素计算加减乘除与比较运算2.1 基础算术运算逐元素计算指的是两个张量相同位置上的元素独立参与运算结果张量形状与输入张量形状保持一致。最典型的就是加、减、乘、除以及幂运算、取模、整除。import torch # 两个相同形状的张量 a torch.tensor([1.0, 2.0, 3.0]) b torch.tensor([4.0, 5.0, 6.0]) # 加减乘除 print(a b) # tensor([5., 7., 9.]) print(a - b) # tensor([-3., -3., -3.]) print(a * b) # tensor([ 4., 10., 18.]) print(a / b) # tensor([0.2500, 0.4000, 0.5000]) # 幂运算和取余 print(a ** 2) # tensor([1., 4., 9.]) print(b % 2) # tensor([0., 1., 0.]) print(torch.pow(a, 2)) # 等价于 a ** 2注意一个细节a * b是逐元素乘法不是矩阵乘法。矩阵乘法要使用torch.mm、torch.matmul或这一点初学者最容易混淆。实际阅读 PyTorch 源码时经常会看到张量前面有*运算符如果参与运算的两个张量形状是[batch, seq, hidden]那这个*就是在做逐元素乘常用来实现注意力权重的掩码乘法而或torch.matmul才是在做特征维度的线性变换。再看除法/和torch.div等价。PyTorch 中的浮点数除法不会因为除零直接崩溃但会出现inf或nan这在后续计算 loss 时会造成梯度异常。实际工程中用torch.where或者加一个极小值eps来避免除零是常见操作思路本课先掌握基本规则训练场景下的数值稳定性后续再单独讨论。2.2 幂运算、指数、对数与常用激活函数除了基础的加减乘除模型代码里出现频率更高的是torch.exp、torch.log、torch.sqrt、torch.square这类逐元素数学函数。softmax 和 sigmoid 的实现本质就是这些逐元素函数的组合。import torch x torch.tensor([0.0, 1.0, 2.0]) print(torch.exp(x)) # e^x输出 tensor([1.0000, 2.7183, 7.3891]) print(torch.log(torch.exp(x))) # x print(torch.square(x)) # 平方 print(torch.sqrt(torch.tensor([4.0, 9.0, 16.0]))) # 平方根 print(torch.clamp(x, min1.0, max2.0)) # 裁剪到 [1,2]torch.clamp在梯度裁剪和值域限制中非常常见它也是逐元素操作。PyTorch 的逐元素运算都支持out参数可以把结果写入预设张量减少内存拷贝但是日常脚本里用得不多大多数情况下直接写z a b即可。在神经网络中F.relu(x)本质上就是逐元素计算。你可以用 PyTorch 直接验证import torch import torch.nn.functional as F x torch.tensor([-2.0, -1.0, 0.0, 1.0, 2.0]) print(F.relu(x)) # tensor([0., 0., 0., 1., 2.]) # 手动实现 print(torch.maximum(x, torch.tensor(0.0)))手动实现结果与F.relu一致。理解这一点后面看 Transformer 里的FFN、注意力权重计算就不会觉得公式是魔法。2.3 比较运算与逻辑运算逐元素比较运算返回的是布尔张量形状与输入一致。比较运算在掩码构造、条件筛选、评估指标计算里很有用。import torch a torch.tensor([1.0, 2.0, 3.0]) b torch.tensor([3.0, 2.0, 1.0]) print(a b) # tensor([False, False, True]) print(a b) # tensor([False, True, False]) print(a b) # tensor([False, True, True]) # 布尔张量参与运算 mask a 2 print(a[mask]) # tensor([3.])这里就引出一个关键概念布尔掩码索引。a[mask]只会返回满足条件的元素常用于过滤无用数据比如把样本中标签为-1的无效结果剔除。实际使用中还要注意如果要在梯度计算中保留位置信息一般用torch.where(mask, a, b)而不是直接索引c torch.where(a 2, a, torch.zeros_like(a)) print(c) # tensor([0., 0., 3.])torch.where的三个参数都支持广播所以它也是理解接下来广播机制的绝佳例子。3. 矩阵乘法mm、matmul 与 运算符3.1 二维矩阵乘法 torch.mm矩阵乘法是深度学习中最重要的运算。torch.mm只接受二维张量输入[m, k]和[k, n]输出[m, n]。中间维度必须相同。import torch m1 torch.randn(3, 4) m2 torch.randn(4, 5) out torch.mm(m1, m2) print(out.shape) # torch.Size([3, 5])这与线性层nn.Linear的原理一致。一个nn.Linear(in_features4, out_features5)在做前向传播时本质上就是把输入[..., 4]与权重矩阵[4, 5]做矩阵乘法再加上偏置。理解了torch.mm线性层内部机制就破案了。手动实现一个简化版线性层import torch x torch.randn(2, 4) # 2个样本每个样本4个特征 weight torch.randn(4, 5) # 输出5个维度 bias torch.randn(5) y torch.mm(x, weight) bias # 先乘后加偏置 print(y.shape) # torch.Size([2, 5])这里用的是广播加法torch.mm(x, weight)形状为[2, 5]bias形状为[5]两者相加会自动把 bias 广播到每一行。广播机制让代码干净了很多但如果你不理解广播规则看到[2, 5] [5]不报错反而会觉得莫名奇妙。后面第 4 节会专门讲透。3.2 高维张量的批量矩阵乘法 torch.matmul实际模型训练中很少只有二维矩阵。一个 batch 的输入通常是[batch_size, seq_len, hidden_dim]比如一个包含 8 个句子、每个句子 10 个 token、每个 token 用 768 维向量表示的文本 batch形状就是[8, 10, 768]。torch.matmul是torch.mm的推广它支持二维、高维和混合维度的矩阵乘法。规则可以简单理解为把最后两维当作真正的矩阵做乘法前面的维度按广播规则处理。import torch # 批量矩阵乘法 b1 torch.randn(2, 3, 4) b2 torch.randn(2, 4, 5) out torch.matmul(b1, b2) print(out.shape) # torch.Size([2, 3, 5]) # 高维张量与二维矩阵相乘 b3 torch.randn(2, 3, 4) m torch.randn(4, 5) out2 torch.matmul(b3, m) print(out2.shape) # torch.Size([2, 3, 5])第二个例子中二维矩阵m被广播到最后两维上相当于对 batch 中每一个[3, 4]矩阵都执行torch.mm。Transformer 的Q K^T就是这种批量矩阵运算的典型场景其中Q的形状是[batch, heads, seq, head_dim]K转置后是[batch, heads, head_dim, seq]两者matmul得到[batch, heads, seq, seq]的注意力分数。运算符是torch.matmul的语法糖。以下两种写法完全等价out_a b1 b2 out_b torch.matmul(b1, b2)在实际项目源码里写法比torch.matmul更常见。你只需要注意和*在语义上的天壤之别。3.3 矩阵乘法维度不匹配的常见报错矩阵乘法最常见的错误就是内部维度不一致。当你执行torch.mm(A, B)时如果A.shape[-1] ! B.shape[-2]PyTorch 会直接抛出类似mat1 and mat2 shapes cannot be multiplied的报错。这个报错的解法只有一个方向确认谁是矩阵乘法中的k维度。排查顺序建议如下打印两个张量的shape检查A.shape[-1]和B.shape[-2]是否相等如果要相乘的语义是特征变换通常需要转置其中一个张量用A.T、A.transpose(-1, -2)或A.permute如果涉及 batch 维还需要确认两个张量的 batch 维数是否满足广播条件。下面展示一个三维张量转置配合矩阵乘法的场景import torch q torch.randn(2, 4, 8) # batch2, seq4, dim8 k torch.randn(2, 6, 8) # batch2, seq6, dim8 # 想把 q 和 k 的最后一维做内积需要把 k 转成 [2, 8, 6] k_t k.transpose(-1, -2) scores torch.matmul(q, k_t) print(scores.shape) # torch.Size([2, 4, 6])这个例子是所有自注意力机制的第一步。你把维度想清楚了再看scaled_dot_product_attention的官方文档就会非常顺。4. 广播机制规则、原理与实战4.1 广播的规则广播机制是 PyTorch 张量运算里最容易忽略、也最容易出错的规则。简单说是允许形状不完全相同的张量进行逐元素运算但必须满足两个条件从最后一个维度开始向前比较两个维度要么相等要么其中一个是 1要么其中一个不存在。如果一个维度既不相等也不为 1就报错。这句话听起来简单实际操作时很多初学者还是会懵。我们看一个直接例子import torch a torch.randn(3, 1) b torch.randn(1, 4) c a b print(a.shape, b.shape, c.shape) # torch.Size([3, 1]) torch.Size([1, 4]) torch.Size([3, 4])这里a的形状是[3, 1]b的形状是[1, 4]。从最后一个维度比较最后一个维1和4不相等但一个为 1可以广播倒数第二个维3和1不相等但一个为 1可以广播。所以结果是[3, 4]。如果a是[3, 5]b是[1, 4]最后一个维度5和4既不相同也不为 1直接报错。4.2 标量与张量的广播标量是最典型的广播场景。比如对一个张量做标准化减去均值再除以标准差均值是标量张量是任意形状PyTorch 会自动把标量广播到每一个元素上。import torch x torch.tensor([[1.0, 2.0], [3.0, 4.0]]) mean x.mean() std x.std() # 标量广播到整个张量 normalized (x - mean) / std print(normalized)这里的mean和std都是零维张量与二维张量运算时自动广播。从数学上等价于对每个元素执行相同的变换。更常见的逐行标准化场景中均值张量是[3]而数据是[4, 3]从最后维比较3等于3可以广播数据的前导维4在均值张量里不存在相当于自动补齐为 1也满足广播条件。import torch data torch.randn(4, 3) # 4个样本3个特征 mean data.mean(dim0) # 形状 [3] centered data - mean # 每行都减去对应的均值 print(centered.shape)这个操作在特征标准化中极其常见。mean的形状是[3]data的形状是[4, 3]减号执行逐元素计算核心是“形状不相同但满足广播条件”。如果写成data.mean(dim0, keepdimTrue)得到的mean形状就是[1, 3]结果不变但语义上更明确标明这是一个“行向量”。在更复杂的模型中keepdimTrue可以帮助保持维度结构避免后续矩阵乘法时少了一个维度产生误操作。4.3 广播不兼容的常见错误形如RuntimeError: The size of tensor a (3) must match the size of tensor b (4) at non-singleton dimension 1的报错是广播不兼容的典型现象。出现这种错误时不要看整段代码先打印两个张量的shape然后从最后维开始逐个对比。import torch a torch.randn(3, 4) b torch.randn(3, 3) # 下面这行会报错 # c a ba的形状是[3, 4]b的形状是[3, 3]。最后维4和3不相等且都不为 1所以无法广播。解决方案要么是转置b要么是重采样使特征维一致取决于你的业务语义。还有一种容易误判的情况是只多了一个前导维a torch.randn(3, 4) b torch.randn(2, 3, 4) c a b print(c.shape) # torch.Size([2, 3, 4])这里a是[3, 4]b是[2, 3, 4]a会自动补齐为[1, 3, 4]再广播到[2, 3, 4]。这解决了给 batch 数据统一加偏置的问题。实际模型中这种广播很常用但也要注意别把语义理解错a会被复制 2 份参与每个 batch 的计算。4.4 用 torch.where 和 unsqueeze 控制广播位置有时你希望把形状[4]的向量广播到[3, 4]的每一行可以直接加。但如果你希望把形状[3]的向量广播到[3, 4]的每一列就必须先unsqueeze变成[3, 1]。import torch col_means torch.tensor([1.0, 2.0, 3.0]) # 长度是 3 matrix torch.randn(3, 4) # 错误思路直接加 col_means形状 [3] 和 [3, 4] 无法匹配 # corrected matrix col_means # 会报错 # 正确做法把 col_means 变成 [3, 1] col_means_col col_means.unsqueeze(1) # [3, 1] corrected matrix col_means_col print(corrected.shape)unsqueeze(1)在模型量化、特征归一化、位置编码加法中非常常用。位置编码[seq_len, hidden]要加到输入[batch, seq_len, hidden]上时也需要借助广播如果代码里没有unsqueeze或expand那大概率是形状已经天然匹配。理解广播维度对齐规则比死记 API 更重要。5. 在 GPU 上验证张量运算与资源观察这一节是很多人跳过但实际很有用的内容。PyTorch 在 CPU 上跑起来方便但真实训练和推理大多依赖 GPU。本课不涉及具体训练只讲怎么把张量移到 GPU 上执行运算并观察显存变化。5.1 设备转移与设备一致性import torch # 判断 CUDA 是否可用 print(torch.cuda.is_available()) # 创建张量并转移到 GPU x torch.randn(1000, 1000) if torch.cuda.is_available(): x_cuda x.cuda() print(x_cuda.device) # cuda:0一旦张量在 GPU 上参与运算的所有张量必须也在同一设备上否则会报Expected all tensors to be on the same device。实际排查时最简单的方式是逐个打印.device。在 PyTorch 2.x 中更推荐使用to(device)配合统一变量管理device torch.device(cuda if torch.cuda.is_available() else cpu) x torch.randn(1000, 1000, devicedevice)这比到处写.cuda()更适合工程化。分布式训练和 AMP 混合精度训练中设备管理会更复杂但本课先养成统一写device的代码习惯。5.2 通过 nvidia-smi 观察显存在终端执行nvidia-smi重点关注Memory-Usage。PyTorch 在 GPU 上创建张量时显存占用随张量规模和数量变化。实际观测方法很简单先跑一次小张量运算记录显存基线再创建大张量记录峰值两者差值就是当前张量占用的显存。由于每个人的 GPU 型号、驱动、CUDA 版本不同本课不给出固定数字。你可以用一个统一的测试脚本自行探测import torch import os device torch.device(cuda if torch.cuda.is_available() else cpu) print(fusing device: {device}) x torch.randn(4000, 4000, devicedevice) y torch.randn(4000, 4000, devicedevice) z torch.mm(x, y) print(z.shape) print(done)运行这段代码后再用nvidia-smi对比前后显存变化。要注意的是PyTorch 显存分配带有缓存机制进程结束后显存才会完全释放。如果你在 Jupyter 里反复执行显存可能不会立刻下降这时候需要重启 kernel。5.3 CPU 与 GPU 运算结果一致性张量运算在 CPU 和 GPU 上数学上等价但浮点计算顺序不同可能产生极小数值差异。对大多数场景可以忽略。验证时不要直接用判断 CUDA 与 CPU 的结果而是用torch.allcloseimport torch cpu_result torch.mm(torch.randn(100, 100), torch.randn(100, 100)) gpu_result None if torch.cuda.is_available(): a torch.randn(100, 100, devicecuda) b torch.randn(100, 100, devicecuda) gpu_result torch.mm(a, b).cpu() if gpu_result is not None: print(torch.allclose(cpu_result, gpu_result, atol1e-4))torch.allclose是后续写单元测试和检查模型输出一致性的常用函数。建议从本课开始养成用它判断张量近似相等的习惯不要依赖。6. 张量运算的步数、形状与内存新手常见问题6.1 为什么*和结果完全不同*是逐元素乘法要求两个张量形状可广播是矩阵乘法要求最后两维满足矩阵乘法条件。下面的例子可以直观看出差异import torch a torch.tensor([[1., 2.], [3., 4.]]) b torch.tensor([[5., 6.], [7., 8.]]) print(a * b) # tensor([[ 5., 12.], # [21., 32.]]) print(a b) # tensor([[19., 22.], # [43., 50.]])从数值就可以看出*只做了对应位置相乘做了完整的矩阵乘法。如果模型代码里把这两个运算符混用大概率会出现形状对不上或者结果数值明显异常。排查时优先逐层打印中间结果的shape。6.2 广播导致的内存暴增广播机制在写代码时很舒服但它可能隐式复制张量。比如[1000, 1]和[1, 1000]相加实际结果[1000, 1000]会占用 8MB 左右float32 下这还只是一个小例子。如果第一个维度是几万广播后的内存会非常可观。所以在特征处理时如果某个张量本来就要重复使用可以用expand显式表示import torch base torch.tensor([[1.0], [2.0], [3.0]]) # [3, 1] expanded base.expand(3, 4) print(expanded.shape) # [3, 4]expand不会立即复制数据而是创建视图在后续参与计算时才产生实际数据。它比直接用repeat复制更节省内存。但是要注意expand只能作用于维度为 1 的轴。这又回到了广播规则本身。6.3 形状打印习惯查错时建议打印print(a.shape, b.shape, a.device, b.device)把 shape 和 device 一起打印一次解决两类问题。在 Jupyter 或脚本里都可以这样做。很多时候你花半小时找 bug最后发现只是某个中间张量多了一个维度或者少了一个unsqueeze。7. 本课完整测试脚本建议把下面这段代码保存为一个 Python 文件比如tensor_ops_demo.py在终端里运行python tensor_ops_demo.pyimport torch def main(): # 1. 逐元素计算 a torch.tensor([1.0, 2.0, 3.0]) b torch.tensor([4.0, 5.0, 6.0]) print(a b:, a b) print(a * b:, a * b) # 2. 矩阵乘法 m1 torch.randn(3, 4) m2 torch.randn(4, 5) out torch.mm(m1, m2) print(mm out shape:, out.shape) # 3. matmul 高维 b1 torch.randn(2, 3, 4) b2 torch.randn(2, 4, 5) out2 torch.matmul(b1, b2) print(matmul out shape:, out2.shape) # 4. 广播 x torch.randn(3, 1) y torch.randn(1, 4) z x y print(broadcast out shape:, z.shape) # 5. 设备判断 device torch.device(cuda if torch.cuda.is_available() else cpu) gpu_tensor torch.randn(2, 2, devicedevice) print(gpu_tensor device:, gpu_tensor.device) if __name__ __main__: main()我建议你实际运行后把所有shape打印和注释里的输出对照一遍。如果想进一步巩固可以把b1、b2的形状分别改成[3, 4]、[2, 4, 5]和[4, 9]等不同组合观察什么情况下报错、什么情况下广播成功。8. 常见问题与排查方法问题现象可能原因排查方式解决方案mat1 and mat2 shapes cannot be multiplied矩阵乘法的内部维度k不一致打印两个张量的 shape转置A或修改B的最后一维The size of tensor a (3) must match ...广播不兼容从最后维逐个对比 shape用unsqueeze或reshape对齐维度Expected all tensors to be on the same deviceCPU 与 GPU 张量混合参与运算打印每个张量的.device全部to(device)CUDA out of memory广播后张量过大或缓存未释放观察nvidia-smi显存占用缩小 batch 或重启 kernel结果出现nan或inf除零、log(0)、指数溢出检查输入数据范围使用torch.clamp或加epsCant call numpy() on Tensor that requires grad在需要梯度的张量上调用numpy()检查requires_grad和是否在with torch.no_grad()中使用.detach().cpu().numpy()PyTorch 2.6 加载旧模型报WeightsUnpickler errortorch.load默认weights_onlyTrue查看报错堆栈是否在torch.load设置weights_onlyFalse只对可信模型关于最后一条补充几句。PyTorch 2.6 开始torch.load默认把weights_only设为True这是为了安全性防止加载恶意 pickle 数据。如果你加载的是自己训练的模型遇到报错时可以显式传weights_onlyFalse如果你参与的项目代码还在用旧写法请在升级 PyTorch 时留意这个行为变化。9. 最佳实践与工程建议以下几条不仅适用于本课也适用于后续所有 PyTorch 项目统一用torch.device管理设备不要混用.cuda()和.cpu()在数据加载、模型初始化、前向传播三个环节都统一to(device)避免设备不一致报错。每次写完张量运算先打印shape再继续写下一行。尤其是矩阵乘法前后把中间结果形状记录下来能减少大量“玄学 bug”。对逐元素运算和矩阵运算要区分清楚。需要特征变换用nn.Linear或需要掩码过滤、逐元素加权用*。广播机制能省很多代码但不要把维度隐式地“变没”。如果某个中间张量形状不够直观用unsqueeze或reshape显式声明。涉及 GPU 显存时先用nvidia-smi记录基线再跑大张量。PyTorch 的缓存机制会导致显存不立即释放这不是内存泄漏只是缓存策略必要时候可以用torch.cuda.empty_cache()手动释放。如果本地训练环境下载慢安装 PyTorch 可以优先考虑国内镜像源。这里不展开教程但搜索“PyTorch 国内镜像安装”能很快找到方案。加载旧模型权重时注意 PyTorch 2.6 的weights_only默认值变化。升级版本前先看官方 release notes再批量检查代码里的torch.load调用。10. 下一步学习建议本课完成之后你可以自己设计一个小实验来巩固创建形状为[8, 3, 32]的随机张量模拟 8 个样本、每个样本 3 个时间步、每个时间步 32 维特征。尝试对它分别做逐行均值标准化、批量矩阵乘法、与形状为[32, 16]的线性变换矩阵相乘把所有中间结果的shape打印出来。这个实验能同时覆盖逐元素计算、矩阵乘法和广播机制。做完后你会发现Transformer、TCN、股票预测模型里的很多张量流动本质上就是在重复这三种操作的不同组合。下一课可以继续深入autograd和反向传播理解requires_grad、backward()与计算图的关系。在那之前先把本课的运算规则本地跑熟遇到问题翻一下第 8 节的排查表基本够用。