从 2017 年第一次接触 PyTorch 到现在我已经用它完成了十几个项目从图像分类到目标检测从单机训练到多卡分布式可以说踩过的坑比写的模型还多。如果你正准备选一个深度学习框架入门或者已经写了一些代码但总感觉知其然不知其所以然这篇文章就是写给你的。我会从框架选型的底层逻辑讲起把 Tensor、autograd、nn.Module 这些核心概念掰开揉碎再带你完整走一遍训练图像分类模型的实操流程最后把那些文档里不会写的显存崩溃、不收敛、过拟合问题一次性聊透。1. 为什么我最终选了 PyTorch框架选型的真实考量1.1 动态图给我的工作效率提升先聊一个很多人忽视的点PyTorch 最大的杀手锏不是好用的 API而是动态计算图。什么意思就是你每次 forward 的时候模型会实时构建一张计算图下一次输入不同形状的数据图就跟着变。这听起来抽象但实际体验差别巨大。我之前用静态图框架写一个带条件分支的网络时得先把所有可能的路径都声明出来还得用专门的tf.cond或者tf.while_loop去控制流程代码写起来绕得不行。而 PyTorch 里就是普通的if、for跟写 Python 一样直接。举个例子你想根据输入的长度决定是否跳过某个模块直接这么写def forward(self, x): if x.size(0) 32: x self.big_branch(x) else: x self.small_branch(x) return x这在 PyTorch 里是合法且完全正常的因为计算图是运行时动态构建的。动态图让调试变得非常直观你可以在任意一行打断点打印中间 tensor 的值不需要像静态图那样先编译再执行。对于做研究或者快速验证想法的人这几乎就是效率翻倍。1.2 生态和调试体验从玄学调参到所见即所得第二个让我坚定选择它的理由是调试体验。说实话刚转过来的时候我都惊了——原来训练循环是可以这么写清楚的。PyTorch 的代码风格极度 Pythonic训练一个模型就是标准的五步走前向传播、算 loss、反向传播、更新参数、清零梯度。每一步都是显式调用没有任何魔法隐藏在里面。optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, labels) loss.backward() optimizer.step()这种显式的流程对新手极其友好你想看哪一步的中间结果直接print或者用调试器进去看就行。反观某些高级封装框架一句model.fit()就把所有细节吞进去了遇到问题根本不知道从哪里排查。PyTorch 的生态也成熟得可怕HuggingFace 的 Transformers、Ultralytics 的 YOLO、各种 SOTA 模型的官方实现基本都是 PyTorch 优先这意味着你几乎不需要自己从零复现别人论文里的结构直接拿来改就行。1.3 和另外两个主流方案的对比如果用一句话总结我的选型结论追求灵活性和快速迭代就选 PyTorch追求极致的部署优化和自动化再考虑另外两个方案。我之前也认真试过其他框架。静态图框架在服务部署上有优势导出的模型经过优化后推理效率高但它那套先定义后运行的模式对复杂模型来说维护成本太高。另一个老牌框架虽然上手快、文档全但模型定义和训练的底层逻辑有时候黑箱感很强出了问题难以深入定位。PyTorch 刚好站在一个很好的平衡点动态图的灵活 足够丰富的官方工具链TorchScript、ONNX、量化、分布式既能满足快速实验也能在生产环境中硬碰硬。提示框架没有绝对的好坏关键看团队技术储备和项目需求。但如果你的目标是快速验证想法、灵活改进网络结构、把深度学习原理搞清楚PyTorch 是目前最优解。2. 动手之前必须吃透的 5 个核心概念2.1 Tensor不只是多维数组Tensor 是 PyTorch 的基本数据结构你可以把它理解成能自动求导的 N 维数组。它和 NumPy 的 ndarray 长得几乎一样但多了两个关键能力可以搬到 GPU 上计算可以参与自动求导。建议新手先记住几个操作规律。张量创建常用torch.tensor()、torch.zeros()、torch.randn()形状变换推荐优先用view()而不是reshape()区别在于view()要求张量在内存中是连续的如果不连续就会报错而reshape()会自动处理拷贝。实际写代码时发现view()报错换成reshape()就能解决原因就是源张量经过转置或切片后内存布局不连续了。还有一个高频踩坑点是torch.Tensor()和torch.tensor()的差异。torch.Tensor()是类的构造器默认生成 float32 的未初始化数据torch.tensor(data)会主动拷贝数据并推断 dtype。如果传入的是整数列表torch.tensor([1, 2])得到 int64而torch.Tensor([1, 2])得到 float32。这个细节在拼接张量时偶尔会引发 dtype 不一致的报错我一度被坑得不轻。import torch a torch.tensor([1, 2, 3]) print(a.dtype) # torch.int64 b torch.Tensor([1, 2, 3]) print(b.dtype) # torch.float322.2 autograd自动求导到底怎么工作autograd 是 PyTorch 的灵魂。你只需要把网络中参与梯度计算的张量设置requires_gradTrue然后每次前向传播时PyTorch 会自动记录所有涉及这个张量的运算构建一张反向传播图。当你调用loss.backward()时它按照链式法则从 loss 开始反向计算每个叶子节点的梯度并把结果存到对应张量的.grad属性上。这里有一个重要的认知只有叶子节点才会保存梯度。中间运算得到的张量虽然也有.grad但默认是None。如果想拿到中间变量的梯度要主动调用register_hook这个功能在分析模型内部状态时非常有用。x torch.tensor(2.0, requires_gradTrue) y x ** 2 3 * x 1 y.backward() print(x.grad) # 2*x 3 7很多新手犯的错误是在训练循环里反复调用backward()但忘了把上一轮存的梯度清零。梯度是累加的如果不清零两次 batch 的梯度会叠加在一起参数更新方向完全错乱。正确做法是在每个 batch 开始时调用optimizer.zero_grad()这一步不能省。2.3 nn.Module把网络写成乐高积木nn.Module是所有网络层的基类。你用nn.Linear、nn.Conv2d搭积木一样组装网络这些层都是nn.Module的子类。自定义模型时有两个必须记住的约定class MyNet(nn.Module): def __init__(self): super().__init__() self.conv1 nn.Conv2d(3, 16, kernel_size3) self.fc nn.Linear(16 * 6 * 6, 10) def forward(self, x): x self.conv1(x) x torch.flatten(x, start_dim1) return self.fc(x)第一在__init__里定义子层在forward里定义计算逻辑。你只需要保证调用模型前参数已经初始化好不需要自己手动绑定权重。第二模型默认处于训练模式要记住切换——这个我在下一节单独说。nn.Module还有一个很有用的机制调用.to(device)时它会递归地把所有子模块的参数和缓冲区都搬到指定设备上。所以不要手动遍历参数一个个搬直接整个模型model.to(cuda)就行省事且不容易漏。2.4 Dataset 与 DataLoader数据管线的正确姿势很多人拿到数据就直接for batch in data循环读取数据量小看不出问题一旦数据量大GPU 会时刻处于等待状态训练效率极低。正确姿势是自定义Dataset类然后交给DataLoader做多进程加载。class MyDataset(torch.utils.data.Dataset): def __init__(self, file_list): self.file_list file_list def __len__(self): return len(self.file_list) def __getitem__(self, idx): img read_image(self.file_list[idx]) label get_label(self.file_list[idx]) return img, labelDataLoader里我最常用的两个参数是shuffleTrue和num_workers4。shuffle在每个 epoch 开始时打乱数据顺序避免模型学到样本顺序的规律num_workers控制加载数据的子进程数量通常设为 CPU 核心数的一半左右。注意注意在 Windows 系统上num_workers设成 0改为多进程时在主进程里加载反而最稳否则可能出现子进程报错这是踩过无数次才得出的经验。2.5 训练模式与评估模式一个容易忽略的开关PyTorch 模型有两种模式model.train()和model.eval()切换它会影响 BatchNorm 和 Dropout 等层的行为。BatchNorm 在训练模式下使用当前 batch 的均值和方差做归一化并更新全局统计量而评估模式下直接使用训练阶段累计的全局统计量。Dropout 在训练时会随机丢弃部分神经元评估时则保留所有神经元。如果不切换模式最典型的症状是训练 loss 正常下降但验证集上指标完全不对或者相同数据在训练和验证时输出不一样。我在第一次用 ResNet 做迁移学习时就栽过这个跟头整整查了两天最后发现是忘了写model.eval()。记住这个最小流程model.eval() with torch.no_grad(): total_correct 0 for data, label in val_loader: outputs model(data) total_correct (outputs.argmax(1) label).sum()torch.no_grad()会关闭梯度计算和追踪既省显存又能让推理速度更快。推理时千万别让网络计算梯度这在批量大的时候甚至能快一倍以上。3. 从零训练一个图像分类模型完整实操流程3.1 准备数据以猫狗分类为例我拿一个经典的猫狗二分类任务演示。假设你本地有一个data/目录下面分成train/cat、train/dog、val/cat、val/dog四个子目录。我们先构建一个简洁的 Datasetimport torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader, Dataset from torchvision import transforms from PIL import Image class CatDogDataset(Dataset): def __init__(self, folder, transformNone): self.samples [] self.transform transform for label, name in enumerate([cat, dog]): for img_name in os.listdir(os.path.join(folder, name)): self.samples.append((os.path.join(folder, name, img_name), label)) def __len__(self): return len(self.samples) def __getitem__(self, idx): path, label self.samples[idx] image Image.open(path).convert(RGB) if self.transform: image self.transform(image) return image, label这里有个细节Image.open(...).convert(RGB)是为了统一通道数。灰度图如果不转成 RGB 会导致后面卷积层输入通道不匹配PNG 图如果没有convert(RGB)可能会出现带透明通道RGBA的输入同样会报错。预处理 transform 我建议至少包含Resize((224, 224))、ToTensor()和归一化。transform_train transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])数据增强里RandomHorizontalFlip是免费午餐对一般图像分类任务不会降低指标反而能增加数据多样性缓解过拟合。归一化的 mean/std 如果不想自己算可以直接用 ImageNet 统计出来的这个经典值通用性很强。3.2 定义网络结构从简单 CNN 开始不要一开始就上 ResNet、ViT 这种大模型先用一个小型 CNN 把整个流程跑通逻辑没问题之后再换大模型。我写了一个只有三层卷积和一层全连接的迷你网络参数量很小CPU 上也能跑class SimpleCNN(nn.Module): def __init__(self, num_classes2): super().__init__() self.features nn.Sequential( nn.Conv2d(3, 16, kernel_size3, padding1), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(16, 32, kernel_size3, padding1), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(32, 64, kernel_size3, padding1), nn.ReLU(), nn.MaxPool2d(2), ) self.classifier nn.Sequential( nn.Flatten(), nn.Linear(64 * 28 * 28, 128), nn.ReLU(), nn.Linear(128, num_classes) ) def forward(self, x): return self.classifier(self.features(x))注意全连接层输入维度的计算输入图像是 224×224经过三次 MaxPool2d(2)空间尺寸变成 224 ÷ 2 ÷ 2 ÷ 2 28所以64 * 28 * 28。这个数字如果不对跑前向传播时会直接报矩阵乘法维度错误。我的习惯是先把输入 shape 打印出来确认慢一点但稳妥。3.3 配置训练超参数学习率、批次大小、优化器选择超参数的选择直接决定模型能不能收敛。学习率是最关键的。我一般先用1e-3作为初始值如果是微调预训练模型就降到1e-4甚至1e-5。批次大小受显存限制常规范围是 16~128但批次太小容易导致梯度噪声大、震荡批次太大则收敛变慢且考验显存。优化器我日常用 Adam 最多因为它自带自适应学习率对调参新手很友好。但 Adam 也有坑——在使用预训练模型微调时如果学习率偏高容易突然梯度爆炸导致 loss 变成 NaN我后面会讲如何规避。device torch.device(cuda if torch.cuda.is_available() else cpu) model SimpleCNN().to(device) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr1e-3)选择CrossEntropyLoss时有个隐含条件模型输出的是原始 logits不需要先经过 softmax因为该 loss 内部会做log_softmax和NLLLoss。很多新手在模型后面多加一个nn.Softmax()再接CrossEntropyLoss会得到很奇怪的结果——不能算错但已经冗余且数值上可能不稳正确做法是都用nn.CrossEntropyLoss处理 logits。3.4 训练循环框架代码与解释下面是一段通用训练循环代码我在很多项目中反复使用唯一需要改动的是train_one_epoch里的数据加载方式def train_one_epoch(model, loader, criterion, optimizer): model.train() total_loss 0 correct 0 total 0 for inputs, labels in loader: inputs, labels inputs.to(device), labels.to(device) optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, labels) loss.backward() optimizer.step() total_loss loss.item() * inputs.size(0) _, preds torch.max(outputs, 1) correct (preds labels).sum().item() total inputs.size(0) return total_loss / total, correct / total核心步骤就是刚才说的五步。注意loss.item()会把 GPU 上的标量同步到 CPU 并转成 Python float这个过程是必要的否则 loss 会随着计算图的累积越攒越大最终导致显存溢出。我见过新手直接把loss累积到 list 里最后再算均值结果一个 epoch 没跑完就 OOM 了。正确做法是每步就.item()取出数值断开计算图。训练主循环epochs 10 for epoch in range(epochs): train_loss, train_acc train_one_epoch(model, train_loader, criterion, optimizer) val_loss, val_acc evaluate(model, val_loader, criterion) print(fEpoch {epoch1}/{epochs} | train_loss {train_loss:.4f} | train_acc {train_acc:.4f} | val_acc {val_acc:.4f}) if val_acc best_acc: torch.save(model.state_dict(), best_model.pth) best_acc val_acc我习惯每个 epoch 都做一次验证并把验证集上表现最好的模型保存下来而不是保留最后一个 epoch 的权重。原因是训练后期模型可能在训练集上继续提升但验证集已经过拟合保存验证集最优的那一版更实用。3.5 评估与保存模型别只在训练集上自嗨评估函数和训练函数几乎相同但必须切换成 eval 模式并关闭梯度追踪def evaluate(model, loader, criterion): model.eval() total_loss 0 correct 0 total 0 with torch.no_grad(): for inputs, labels in loader: inputs, labels inputs.to(device), labels.to(device) outputs model(inputs) loss criterion(outputs, labels) total_loss loss.item() * inputs.size(0) _, preds torch.max(outputs, 1) correct (preds labels).sum().item() total inputs.size(0) return total_loss / total, correct / total保存模型有两种情况如果只是做推理直接保存state_dict也就是参数权重字典就够了如果还要在另一个脚本里加载继续训练或者做完整的部署那就保存整个模型结构定义。大多数情况下torch.save(model.state_dict(), model.pth)最轻量加载时用model.load_state_dict(torch.load(model.pth))。提示加载权重前必须保证当前模型的state_dict的 key 和保存的文件一致。如果网络结构改了名字加载会报错所以不要随意更改nn.Sequential内部的层名。4. 训练时最容易踩的 6 个坑含速查表4.1 显存不够怎么办GPU 显存溢出OOM是最常见的坑。报错信息里一般会有CUDA out of memory然后告诉你尝试分配了多少、还剩多少。我的排查顺序是先看是不是真的显存不够还是代码泄漏。用nvidia-smi看占用率如果程序结束显存还没释放多半是某些张量被留在计算图里。一个常见原因是把每个 batch 的 loss 累积到 Python list 而不是标量导致一整轮 epoch 的计算图全部缓存。解决办法确保每次反向传播后计算图被释放通常调用optimizer.zero_grad()后旧的图就会被清掉。如果确实是 batch 太大就减小batch_size到当前显存能容纳的最大的值如果数据本身不小还可以尝试用torch.utils.checkpoint做激活值重计算以时间换空间。但新手优先用最简单的方法batch_size8或16重新试先能跑起来再说。4.2 模型不收敛学习率与初始化loss 一直不降或者降到某个平台就卡住了。有经验的玩家会先怀疑学习率。我觉得 90% 的不收敛问题都是学习率设置不合理。不同优化器对学习率的容忍度差别很大。SGD 用0.01可能没问题Adam 用0.01大概率炸。我的建议是Adam 系列从1e-3起步SGD 从1e-2起步。如果两个 epoch loss 纹丝不动先打印一下梯度的范数print(torch.nn.utils.clip_grad_norm_(model.parameters(), 100))看看梯度是不是小到消失或大到爆炸。梯度很小就调大学习率梯度大得离谱就调小学习率。还有一种情况是模型本身的初始化就有问题。PyTorch 的nn.Conv2d和nn.Linear都自带合理的初始化策略正常不会出问题。但如果你的forward里做了奇怪的权重暴露操作比如直接用nn.Parameter(torch.zeros(...))定义偏置却忘了初始化代价就是梯度为零、参数长期不更新。建议自定义参数时检查初始化的均值方差避免全零或全一。4.3 过拟合的信号与对策如果训练准确率接近 100%但验证集准确率明显低一截这就是过拟合。信号就是训练和验证的差距越来越大。对策按优先级排序增加训练数据量。最朴素但最有效如果没有新数据就做更强的数据增强随机裁剪、色彩抖动、旋转。增加 Dropout 层或权重衰减Adam 里的weight_decay1e-4就是 L2 正则。减小模型参数量。我的经验是先用小模型跑通再决定是否加宽加深模型提容不是玄学是在过拟合和欠拟合之间找平衡点。早停Early Stopping每轮验证保存最优模型当验证指标连续 N 个 epoch 没有提升就停止训练。这个操作能节省大量时间。过拟合不丢人深度学习模型本来就是大容量函数逼近器数据少就会过拟合。关键是你要能识别它并知道去哪里调。4.4 数据加载慢导致 GPU 空转最常见的现象GPU 利用率在 30%~50% 震荡而不是接近 100%。你打开资源管理器发现 CPU 核占用高但 GPU 在等待。这种Eating time等待时间通常是数据预处理瓶颈。解决方案很直接调高DataLoader的num_workers。从 4 调到 8 往往立竿见影。但如果你的数据是从机械硬盘或网络盘读取IO 本身就慢加再多 worker 也有限建议先把图像预处理成更小的尺寸或直接缓存成内存映射格式。另一个容易忽略的点pin_memoryTrue可以让数据从 CPU 到 GPU 的拷贝更快对训练吞吐有可感知的提升。train_loader DataLoader(dataset, batch_size32, shuffleTrue, num_workers8, pin_memoryTrue)4.5 梯度爆炸与梯度消失训练过程中出现 loss 变成 NaN十有八九是梯度爆炸。常见场景是网络层数较深时连续乘法导致梯度指数级增长或者交叉熵 loss 中间出现 log(0) 这样的数值问题。快速解决是用梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm5.0)。在backward()之后、step()之前插入把梯度范数上限限制在 5能阻止爆炸。但这只是治标根本原因是学习率太高或者网络设计不合理。把学习率降到1e-4再观察多半能缓解。梯度消失则相反训练几乎没进展梯度范数趋近于 0。常见于深层网络 Sigmoid 激活函数Sigmoid 在饱和区梯度很小容易死在局部。换成 ReLU 家族ReLU、LeakyReLU、SiLU能显著改善。我现在默认用 ReLU很少直接用 Sigmoid 做中间层。4.6 多 GPU 训练与随机种子如果你有一张以上的 GPU想把训练加速最简单的方案是torch.nn.DataParallel不过它在新版本里越来越边缘化了。我更推荐用torch.nn.parallel.DistributedDataParallel虽然配置复杂一点但性能更好、扩展性更强。import torch.distributed as dist dist.init_process_group(nccl, init_methodtcp://localhost:23456, rank0, world_size1) model torch.nn.parallel.DistributedDataParallel(model.to(device))但如果你只是单卡使用者这部分可以跳过。一个比多卡更值得注意的问题是随机种子。为了让实验结果可复现我总是在程序最前面设置random.seed(0) np.random.seed(0) torch.manual_seed(0) torch.cuda.manual_seed_all(0) torch.backends.cudnn.deterministic True注意cudnn.deterministic True会让部分卷积计算变慢因为它关闭了非确定性算法。平时做实验可以不设但当你需要严格对比模型结构时这套方案必不可少。问题典型症状排查 / 解决方法显存溢出CUDA OOM 或程序崩溃batch_size 减小、释放中间张量、用 checkpoint不收敛loss 不降或 NaN调低学习率、检查梯度范数、检查数据预处理过拟合训练指标远高于验证指标增加数据增强、加 Dropout、早停数据瓶颈GPU 利用率低num_workers 调大、pin_memoryTrue梯度爆炸loss 变成 NaN梯度裁剪、降低学习率梯度消失loss 卡住不动换 ReLU 系激活、检查初始化模式未切换验证结果异常记得model.eval()torch.no_grad()梯度未清零训练震荡每个 batch 调optimizer.zero_grad()5. 模型训练完还没结束部署与加速要点5.1 用 torch.jit 做模型序列化训练完成后模型要离开 Python 训练环境去线上推理。PyTorch 提供了 TorchScript可以把模型转换为一种可被 C 调用的静态表示。最常用的方式是torch.jit.script或torch.jit.trace。trace是追踪执行轨迹生成图适合模型中不含数据依赖分支的情况script是真正的编译器支持if、for等控制流。我个人的经验首先试trace如果转换报错或计算结果不对再换script。model.eval() scripted_model torch.jit.script(model.cpu()) scripted_model.save(model_scripted.pt)注意转换前一定要在 CPU 上做因为 GPU 模型在部署时可能不保证环境里有 GPU。而且script过程对nn.Module的 forward 代码有一定要求不要使用过多 Python 内置的列表推导式、字典操作这些未必能被完全追踪。如果想保留部分灵活性就培训时尽量写纯 tensor 运算逻辑。5.2 ONNX 导出的注意事项如果线上服务不是 C 或 Python 环境而是需要介入其他推理框架ONNX 是通用中间格式。导出其实一行代码torch.onnx.export(model, dummy_input, model.onnx, input_names[input], output_names[output])但你大概率会遇到几个问题。第一个是dummy_input的尺寸必须固定因为 ONNX 默认会把张量形状当作模型输入 shape 的一部分。如果希望输入尺寸可变要在 export 时指定dynamic_axestorch.onnx.export(model, dummy_input, model.onnx, dynamic_axes{input: {0: batch_size}, output: {0: batch_size}})第二个问题是某些算子导出后可能不支持比如一些特殊的激活函数、Resize 模式在旧版本 ONNX 里容易出幺蛾子。遇到算子导出失败我的建议是换等价算子组合或者改用onnx-simplifier做简化但简化前记得验证输出一致性。导出前后用同一个输入跑一遍比对结果误差在 1e-5 以内就还算安全。5.3 量化和剪枝让模型跑得更快如果模型直接部署在资源受限的设备比如手机上那么量化比剪枝更实用。PyTorch 的量化分为动态量化、静态量化和量化感知训练三种。动态量化最简单训练后直接把权重转成 int8推理时再反量化回来适合 LSTM 或 Transformer 这类以线性层为主的模型几乎没有精度损失。静态量化需要校准数据通过收集激活值的分布来优化缩放系数精度损失稍大但推理更快。量化感知训练是在训练过程中就模拟量化误差模型精度最理想但改造成本高。我的建议是如果你只是想把模型体积压一压在 CPU 上用torch.quantization.quantize_dynamic最省事quantized_model torch.quantization.quantize_dynamic( model, {nn.Linear, nn.Conv2d}, dtypetorch.qint8 ) torch.save(quantized_model.state_dict(), model_quantized.pth)量化后模型大小通常能减到原模型的 1/4推理速度也明显提升。但注意量化不是银弹当模型结构里有大量对数值敏感的操作时比如 BatchNorm 在量化时表现不佳需要谨慎验证精度。我做过一个简单的 CNN 分类器量化后准确率掉了 0.3%可以接受但换到一个带注意力机制的小模型精度直接掉了 2%最后只能在速度和精度之间妥协没有做彻底量化。剪枝这里简单说一句训练后剪枝跑在 CPU 上通常能带来 30%~50% 的加速但需要配合稀疏矩阵存储才能发挥最大优势。如果你的部署环境没有特殊优化剪枝收益不如量化直接所以我一般先量化再考虑剪枝。最后再分享一个个人技巧无论最终用哪种部署方案都建议在导出前后做一次黄金样本回归测试。准备 20~50 张有代表性的图片记录原模型和导出模型的输出 logits计算平均绝对误差。这不是为了精益求精而是为了在改动任何一个参数后能快速发现变了没变、变了多少。我靠这套方法救回过好几次上线前的危机——你以为只是加了一行torch.jit.script实际某个算子已经悄悄变了行为。从第一次接触到今天我用 PyTorch 完成的项目大概可以分成三个阶段第一阶段是照猫画虎地跑通教程第二阶段是无脑换模型调参第三阶段是把更多时间花在理解机制和数据上。如果你现在正处于第一阶段不要焦虑按我上面列的步骤先跑一个完整的图像分类流程然后把每个 API 的文档翻一遍你很快会发现 PyTorch 的 API 设计极其一致需要 GPU 就是.to(device)需要梯度就设置requires_grad需要保存就state_dict。这些都掌握之后剩下的就是在具体业务场景里灵活组合的问题了。希望这篇分享能帮你少踩一些我当年踩过的坑。
从 2017 年第一次接触 PyTorch 到现在我已经用它完成了十几个项目从图像分类到目标检测从单机训练到多卡分布式可以说踩过的坑比写的模型还多。如果你正准备选一个深度学习框架入门或者已经写了一些代码但总感觉知其然不知其所以然这篇文章就是写给你的。我会从框架选型的底层逻辑讲起把 Tensor、autograd、nn.Module 这些核心概念掰开揉碎再带你完整走一遍训练图像分类模型的实操流程最后把那些文档里不会写的显存崩溃、不收敛、过拟合问题一次性聊透。1. 为什么我最终选了 PyTorch框架选型的真实考量1.1 动态图给我的工作效率提升先聊一个很多人忽视的点PyTorch 最大的杀手锏不是好用的 API而是动态计算图。什么意思就是你每次 forward 的时候模型会实时构建一张计算图下一次输入不同形状的数据图就跟着变。这听起来抽象但实际体验差别巨大。我之前用静态图框架写一个带条件分支的网络时得先把所有可能的路径都声明出来还得用专门的tf.cond或者tf.while_loop去控制流程代码写起来绕得不行。而 PyTorch 里就是普通的if、for跟写 Python 一样直接。举个例子你想根据输入的长度决定是否跳过某个模块直接这么写def forward(self, x): if x.size(0) 32: x self.big_branch(x) else: x self.small_branch(x) return x这在 PyTorch 里是合法且完全正常的因为计算图是运行时动态构建的。动态图让调试变得非常直观你可以在任意一行打断点打印中间 tensor 的值不需要像静态图那样先编译再执行。对于做研究或者快速验证想法的人这几乎就是效率翻倍。1.2 生态和调试体验从玄学调参到所见即所得第二个让我坚定选择它的理由是调试体验。说实话刚转过来的时候我都惊了——原来训练循环是可以这么写清楚的。PyTorch 的代码风格极度 Pythonic训练一个模型就是标准的五步走前向传播、算 loss、反向传播、更新参数、清零梯度。每一步都是显式调用没有任何魔法隐藏在里面。optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, labels) loss.backward() optimizer.step()这种显式的流程对新手极其友好你想看哪一步的中间结果直接print或者用调试器进去看就行。反观某些高级封装框架一句model.fit()就把所有细节吞进去了遇到问题根本不知道从哪里排查。PyTorch 的生态也成熟得可怕HuggingFace 的 Transformers、Ultralytics 的 YOLO、各种 SOTA 模型的官方实现基本都是 PyTorch 优先这意味着你几乎不需要自己从零复现别人论文里的结构直接拿来改就行。1.3 和另外两个主流方案的对比如果用一句话总结我的选型结论追求灵活性和快速迭代就选 PyTorch追求极致的部署优化和自动化再考虑另外两个方案。我之前也认真试过其他框架。静态图框架在服务部署上有优势导出的模型经过优化后推理效率高但它那套先定义后运行的模式对复杂模型来说维护成本太高。另一个老牌框架虽然上手快、文档全但模型定义和训练的底层逻辑有时候黑箱感很强出了问题难以深入定位。PyTorch 刚好站在一个很好的平衡点动态图的灵活 足够丰富的官方工具链TorchScript、ONNX、量化、分布式既能满足快速实验也能在生产环境中硬碰硬。提示框架没有绝对的好坏关键看团队技术储备和项目需求。但如果你的目标是快速验证想法、灵活改进网络结构、把深度学习原理搞清楚PyTorch 是目前最优解。2. 动手之前必须吃透的 5 个核心概念2.1 Tensor不只是多维数组Tensor 是 PyTorch 的基本数据结构你可以把它理解成能自动求导的 N 维数组。它和 NumPy 的 ndarray 长得几乎一样但多了两个关键能力可以搬到 GPU 上计算可以参与自动求导。建议新手先记住几个操作规律。张量创建常用torch.tensor()、torch.zeros()、torch.randn()形状变换推荐优先用view()而不是reshape()区别在于view()要求张量在内存中是连续的如果不连续就会报错而reshape()会自动处理拷贝。实际写代码时发现view()报错换成reshape()就能解决原因就是源张量经过转置或切片后内存布局不连续了。还有一个高频踩坑点是torch.Tensor()和torch.tensor()的差异。torch.Tensor()是类的构造器默认生成 float32 的未初始化数据torch.tensor(data)会主动拷贝数据并推断 dtype。如果传入的是整数列表torch.tensor([1, 2])得到 int64而torch.Tensor([1, 2])得到 float32。这个细节在拼接张量时偶尔会引发 dtype 不一致的报错我一度被坑得不轻。import torch a torch.tensor([1, 2, 3]) print(a.dtype) # torch.int64 b torch.Tensor([1, 2, 3]) print(b.dtype) # torch.float322.2 autograd自动求导到底怎么工作autograd 是 PyTorch 的灵魂。你只需要把网络中参与梯度计算的张量设置requires_gradTrue然后每次前向传播时PyTorch 会自动记录所有涉及这个张量的运算构建一张反向传播图。当你调用loss.backward()时它按照链式法则从 loss 开始反向计算每个叶子节点的梯度并把结果存到对应张量的.grad属性上。这里有一个重要的认知只有叶子节点才会保存梯度。中间运算得到的张量虽然也有.grad但默认是None。如果想拿到中间变量的梯度要主动调用register_hook这个功能在分析模型内部状态时非常有用。x torch.tensor(2.0, requires_gradTrue) y x ** 2 3 * x 1 y.backward() print(x.grad) # 2*x 3 7很多新手犯的错误是在训练循环里反复调用backward()但忘了把上一轮存的梯度清零。梯度是累加的如果不清零两次 batch 的梯度会叠加在一起参数更新方向完全错乱。正确做法是在每个 batch 开始时调用optimizer.zero_grad()这一步不能省。2.3 nn.Module把网络写成乐高积木nn.Module是所有网络层的基类。你用nn.Linear、nn.Conv2d搭积木一样组装网络这些层都是nn.Module的子类。自定义模型时有两个必须记住的约定class MyNet(nn.Module): def __init__(self): super().__init__() self.conv1 nn.Conv2d(3, 16, kernel_size3) self.fc nn.Linear(16 * 6 * 6, 10) def forward(self, x): x self.conv1(x) x torch.flatten(x, start_dim1) return self.fc(x)第一在__init__里定义子层在forward里定义计算逻辑。你只需要保证调用模型前参数已经初始化好不需要自己手动绑定权重。第二模型默认处于训练模式要记住切换——这个我在下一节单独说。nn.Module还有一个很有用的机制调用.to(device)时它会递归地把所有子模块的参数和缓冲区都搬到指定设备上。所以不要手动遍历参数一个个搬直接整个模型model.to(cuda)就行省事且不容易漏。2.4 Dataset 与 DataLoader数据管线的正确姿势很多人拿到数据就直接for batch in data循环读取数据量小看不出问题一旦数据量大GPU 会时刻处于等待状态训练效率极低。正确姿势是自定义Dataset类然后交给DataLoader做多进程加载。class MyDataset(torch.utils.data.Dataset): def __init__(self, file_list): self.file_list file_list def __len__(self): return len(self.file_list) def __getitem__(self, idx): img read_image(self.file_list[idx]) label get_label(self.file_list[idx]) return img, labelDataLoader里我最常用的两个参数是shuffleTrue和num_workers4。shuffle在每个 epoch 开始时打乱数据顺序避免模型学到样本顺序的规律num_workers控制加载数据的子进程数量通常设为 CPU 核心数的一半左右。注意注意在 Windows 系统上num_workers设成 0改为多进程时在主进程里加载反而最稳否则可能出现子进程报错这是踩过无数次才得出的经验。2.5 训练模式与评估模式一个容易忽略的开关PyTorch 模型有两种模式model.train()和model.eval()切换它会影响 BatchNorm 和 Dropout 等层的行为。BatchNorm 在训练模式下使用当前 batch 的均值和方差做归一化并更新全局统计量而评估模式下直接使用训练阶段累计的全局统计量。Dropout 在训练时会随机丢弃部分神经元评估时则保留所有神经元。如果不切换模式最典型的症状是训练 loss 正常下降但验证集上指标完全不对或者相同数据在训练和验证时输出不一样。我在第一次用 ResNet 做迁移学习时就栽过这个跟头整整查了两天最后发现是忘了写model.eval()。记住这个最小流程model.eval() with torch.no_grad(): total_correct 0 for data, label in val_loader: outputs model(data) total_correct (outputs.argmax(1) label).sum()torch.no_grad()会关闭梯度计算和追踪既省显存又能让推理速度更快。推理时千万别让网络计算梯度这在批量大的时候甚至能快一倍以上。3. 从零训练一个图像分类模型完整实操流程3.1 准备数据以猫狗分类为例我拿一个经典的猫狗二分类任务演示。假设你本地有一个data/目录下面分成train/cat、train/dog、val/cat、val/dog四个子目录。我们先构建一个简洁的 Datasetimport torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader, Dataset from torchvision import transforms from PIL import Image class CatDogDataset(Dataset): def __init__(self, folder, transformNone): self.samples [] self.transform transform for label, name in enumerate([cat, dog]): for img_name in os.listdir(os.path.join(folder, name)): self.samples.append((os.path.join(folder, name, img_name), label)) def __len__(self): return len(self.samples) def __getitem__(self, idx): path, label self.samples[idx] image Image.open(path).convert(RGB) if self.transform: image self.transform(image) return image, label这里有个细节Image.open(...).convert(RGB)是为了统一通道数。灰度图如果不转成 RGB 会导致后面卷积层输入通道不匹配PNG 图如果没有convert(RGB)可能会出现带透明通道RGBA的输入同样会报错。预处理 transform 我建议至少包含Resize((224, 224))、ToTensor()和归一化。transform_train transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])数据增强里RandomHorizontalFlip是免费午餐对一般图像分类任务不会降低指标反而能增加数据多样性缓解过拟合。归一化的 mean/std 如果不想自己算可以直接用 ImageNet 统计出来的这个经典值通用性很强。3.2 定义网络结构从简单 CNN 开始不要一开始就上 ResNet、ViT 这种大模型先用一个小型 CNN 把整个流程跑通逻辑没问题之后再换大模型。我写了一个只有三层卷积和一层全连接的迷你网络参数量很小CPU 上也能跑class SimpleCNN(nn.Module): def __init__(self, num_classes2): super().__init__() self.features nn.Sequential( nn.Conv2d(3, 16, kernel_size3, padding1), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(16, 32, kernel_size3, padding1), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(32, 64, kernel_size3, padding1), nn.ReLU(), nn.MaxPool2d(2), ) self.classifier nn.Sequential( nn.Flatten(), nn.Linear(64 * 28 * 28, 128), nn.ReLU(), nn.Linear(128, num_classes) ) def forward(self, x): return self.classifier(self.features(x))注意全连接层输入维度的计算输入图像是 224×224经过三次 MaxPool2d(2)空间尺寸变成 224 ÷ 2 ÷ 2 ÷ 2 28所以64 * 28 * 28。这个数字如果不对跑前向传播时会直接报矩阵乘法维度错误。我的习惯是先把输入 shape 打印出来确认慢一点但稳妥。3.3 配置训练超参数学习率、批次大小、优化器选择超参数的选择直接决定模型能不能收敛。学习率是最关键的。我一般先用1e-3作为初始值如果是微调预训练模型就降到1e-4甚至1e-5。批次大小受显存限制常规范围是 16~128但批次太小容易导致梯度噪声大、震荡批次太大则收敛变慢且考验显存。优化器我日常用 Adam 最多因为它自带自适应学习率对调参新手很友好。但 Adam 也有坑——在使用预训练模型微调时如果学习率偏高容易突然梯度爆炸导致 loss 变成 NaN我后面会讲如何规避。device torch.device(cuda if torch.cuda.is_available() else cpu) model SimpleCNN().to(device) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr1e-3)选择CrossEntropyLoss时有个隐含条件模型输出的是原始 logits不需要先经过 softmax因为该 loss 内部会做log_softmax和NLLLoss。很多新手在模型后面多加一个nn.Softmax()再接CrossEntropyLoss会得到很奇怪的结果——不能算错但已经冗余且数值上可能不稳正确做法是都用nn.CrossEntropyLoss处理 logits。3.4 训练循环框架代码与解释下面是一段通用训练循环代码我在很多项目中反复使用唯一需要改动的是train_one_epoch里的数据加载方式def train_one_epoch(model, loader, criterion, optimizer): model.train() total_loss 0 correct 0 total 0 for inputs, labels in loader: inputs, labels inputs.to(device), labels.to(device) optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, labels) loss.backward() optimizer.step() total_loss loss.item() * inputs.size(0) _, preds torch.max(outputs, 1) correct (preds labels).sum().item() total inputs.size(0) return total_loss / total, correct / total核心步骤就是刚才说的五步。注意loss.item()会把 GPU 上的标量同步到 CPU 并转成 Python float这个过程是必要的否则 loss 会随着计算图的累积越攒越大最终导致显存溢出。我见过新手直接把loss累积到 list 里最后再算均值结果一个 epoch 没跑完就 OOM 了。正确做法是每步就.item()取出数值断开计算图。训练主循环epochs 10 for epoch in range(epochs): train_loss, train_acc train_one_epoch(model, train_loader, criterion, optimizer) val_loss, val_acc evaluate(model, val_loader, criterion) print(fEpoch {epoch1}/{epochs} | train_loss {train_loss:.4f} | train_acc {train_acc:.4f} | val_acc {val_acc:.4f}) if val_acc best_acc: torch.save(model.state_dict(), best_model.pth) best_acc val_acc我习惯每个 epoch 都做一次验证并把验证集上表现最好的模型保存下来而不是保留最后一个 epoch 的权重。原因是训练后期模型可能在训练集上继续提升但验证集已经过拟合保存验证集最优的那一版更实用。3.5 评估与保存模型别只在训练集上自嗨评估函数和训练函数几乎相同但必须切换成 eval 模式并关闭梯度追踪def evaluate(model, loader, criterion): model.eval() total_loss 0 correct 0 total 0 with torch.no_grad(): for inputs, labels in loader: inputs, labels inputs.to(device), labels.to(device) outputs model(inputs) loss criterion(outputs, labels) total_loss loss.item() * inputs.size(0) _, preds torch.max(outputs, 1) correct (preds labels).sum().item() total inputs.size(0) return total_loss / total, correct / total保存模型有两种情况如果只是做推理直接保存state_dict也就是参数权重字典就够了如果还要在另一个脚本里加载继续训练或者做完整的部署那就保存整个模型结构定义。大多数情况下torch.save(model.state_dict(), model.pth)最轻量加载时用model.load_state_dict(torch.load(model.pth))。提示加载权重前必须保证当前模型的state_dict的 key 和保存的文件一致。如果网络结构改了名字加载会报错所以不要随意更改nn.Sequential内部的层名。4. 训练时最容易踩的 6 个坑含速查表4.1 显存不够怎么办GPU 显存溢出OOM是最常见的坑。报错信息里一般会有CUDA out of memory然后告诉你尝试分配了多少、还剩多少。我的排查顺序是先看是不是真的显存不够还是代码泄漏。用nvidia-smi看占用率如果程序结束显存还没释放多半是某些张量被留在计算图里。一个常见原因是把每个 batch 的 loss 累积到 Python list 而不是标量导致一整轮 epoch 的计算图全部缓存。解决办法确保每次反向传播后计算图被释放通常调用optimizer.zero_grad()后旧的图就会被清掉。如果确实是 batch 太大就减小batch_size到当前显存能容纳的最大的值如果数据本身不小还可以尝试用torch.utils.checkpoint做激活值重计算以时间换空间。但新手优先用最简单的方法batch_size8或16重新试先能跑起来再说。4.2 模型不收敛学习率与初始化loss 一直不降或者降到某个平台就卡住了。有经验的玩家会先怀疑学习率。我觉得 90% 的不收敛问题都是学习率设置不合理。不同优化器对学习率的容忍度差别很大。SGD 用0.01可能没问题Adam 用0.01大概率炸。我的建议是Adam 系列从1e-3起步SGD 从1e-2起步。如果两个 epoch loss 纹丝不动先打印一下梯度的范数print(torch.nn.utils.clip_grad_norm_(model.parameters(), 100))看看梯度是不是小到消失或大到爆炸。梯度很小就调大学习率梯度大得离谱就调小学习率。还有一种情况是模型本身的初始化就有问题。PyTorch 的nn.Conv2d和nn.Linear都自带合理的初始化策略正常不会出问题。但如果你的forward里做了奇怪的权重暴露操作比如直接用nn.Parameter(torch.zeros(...))定义偏置却忘了初始化代价就是梯度为零、参数长期不更新。建议自定义参数时检查初始化的均值方差避免全零或全一。4.3 过拟合的信号与对策如果训练准确率接近 100%但验证集准确率明显低一截这就是过拟合。信号就是训练和验证的差距越来越大。对策按优先级排序增加训练数据量。最朴素但最有效如果没有新数据就做更强的数据增强随机裁剪、色彩抖动、旋转。增加 Dropout 层或权重衰减Adam 里的weight_decay1e-4就是 L2 正则。减小模型参数量。我的经验是先用小模型跑通再决定是否加宽加深模型提容不是玄学是在过拟合和欠拟合之间找平衡点。早停Early Stopping每轮验证保存最优模型当验证指标连续 N 个 epoch 没有提升就停止训练。这个操作能节省大量时间。过拟合不丢人深度学习模型本来就是大容量函数逼近器数据少就会过拟合。关键是你要能识别它并知道去哪里调。4.4 数据加载慢导致 GPU 空转最常见的现象GPU 利用率在 30%~50% 震荡而不是接近 100%。你打开资源管理器发现 CPU 核占用高但 GPU 在等待。这种Eating time等待时间通常是数据预处理瓶颈。解决方案很直接调高DataLoader的num_workers。从 4 调到 8 往往立竿见影。但如果你的数据是从机械硬盘或网络盘读取IO 本身就慢加再多 worker 也有限建议先把图像预处理成更小的尺寸或直接缓存成内存映射格式。另一个容易忽略的点pin_memoryTrue可以让数据从 CPU 到 GPU 的拷贝更快对训练吞吐有可感知的提升。train_loader DataLoader(dataset, batch_size32, shuffleTrue, num_workers8, pin_memoryTrue)4.5 梯度爆炸与梯度消失训练过程中出现 loss 变成 NaN十有八九是梯度爆炸。常见场景是网络层数较深时连续乘法导致梯度指数级增长或者交叉熵 loss 中间出现 log(0) 这样的数值问题。快速解决是用梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm5.0)。在backward()之后、step()之前插入把梯度范数上限限制在 5能阻止爆炸。但这只是治标根本原因是学习率太高或者网络设计不合理。把学习率降到1e-4再观察多半能缓解。梯度消失则相反训练几乎没进展梯度范数趋近于 0。常见于深层网络 Sigmoid 激活函数Sigmoid 在饱和区梯度很小容易死在局部。换成 ReLU 家族ReLU、LeakyReLU、SiLU能显著改善。我现在默认用 ReLU很少直接用 Sigmoid 做中间层。4.6 多 GPU 训练与随机种子如果你有一张以上的 GPU想把训练加速最简单的方案是torch.nn.DataParallel不过它在新版本里越来越边缘化了。我更推荐用torch.nn.parallel.DistributedDataParallel虽然配置复杂一点但性能更好、扩展性更强。import torch.distributed as dist dist.init_process_group(nccl, init_methodtcp://localhost:23456, rank0, world_size1) model torch.nn.parallel.DistributedDataParallel(model.to(device))但如果你只是单卡使用者这部分可以跳过。一个比多卡更值得注意的问题是随机种子。为了让实验结果可复现我总是在程序最前面设置random.seed(0) np.random.seed(0) torch.manual_seed(0) torch.cuda.manual_seed_all(0) torch.backends.cudnn.deterministic True注意cudnn.deterministic True会让部分卷积计算变慢因为它关闭了非确定性算法。平时做实验可以不设但当你需要严格对比模型结构时这套方案必不可少。问题典型症状排查 / 解决方法显存溢出CUDA OOM 或程序崩溃batch_size 减小、释放中间张量、用 checkpoint不收敛loss 不降或 NaN调低学习率、检查梯度范数、检查数据预处理过拟合训练指标远高于验证指标增加数据增强、加 Dropout、早停数据瓶颈GPU 利用率低num_workers 调大、pin_memoryTrue梯度爆炸loss 变成 NaN梯度裁剪、降低学习率梯度消失loss 卡住不动换 ReLU 系激活、检查初始化模式未切换验证结果异常记得model.eval()torch.no_grad()梯度未清零训练震荡每个 batch 调optimizer.zero_grad()5. 模型训练完还没结束部署与加速要点5.1 用 torch.jit 做模型序列化训练完成后模型要离开 Python 训练环境去线上推理。PyTorch 提供了 TorchScript可以把模型转换为一种可被 C 调用的静态表示。最常用的方式是torch.jit.script或torch.jit.trace。trace是追踪执行轨迹生成图适合模型中不含数据依赖分支的情况script是真正的编译器支持if、for等控制流。我个人的经验首先试trace如果转换报错或计算结果不对再换script。model.eval() scripted_model torch.jit.script(model.cpu()) scripted_model.save(model_scripted.pt)注意转换前一定要在 CPU 上做因为 GPU 模型在部署时可能不保证环境里有 GPU。而且script过程对nn.Module的 forward 代码有一定要求不要使用过多 Python 内置的列表推导式、字典操作这些未必能被完全追踪。如果想保留部分灵活性就培训时尽量写纯 tensor 运算逻辑。5.2 ONNX 导出的注意事项如果线上服务不是 C 或 Python 环境而是需要介入其他推理框架ONNX 是通用中间格式。导出其实一行代码torch.onnx.export(model, dummy_input, model.onnx, input_names[input], output_names[output])但你大概率会遇到几个问题。第一个是dummy_input的尺寸必须固定因为 ONNX 默认会把张量形状当作模型输入 shape 的一部分。如果希望输入尺寸可变要在 export 时指定dynamic_axestorch.onnx.export(model, dummy_input, model.onnx, dynamic_axes{input: {0: batch_size}, output: {0: batch_size}})第二个问题是某些算子导出后可能不支持比如一些特殊的激活函数、Resize 模式在旧版本 ONNX 里容易出幺蛾子。遇到算子导出失败我的建议是换等价算子组合或者改用onnx-simplifier做简化但简化前记得验证输出一致性。导出前后用同一个输入跑一遍比对结果误差在 1e-5 以内就还算安全。5.3 量化和剪枝让模型跑得更快如果模型直接部署在资源受限的设备比如手机上那么量化比剪枝更实用。PyTorch 的量化分为动态量化、静态量化和量化感知训练三种。动态量化最简单训练后直接把权重转成 int8推理时再反量化回来适合 LSTM 或 Transformer 这类以线性层为主的模型几乎没有精度损失。静态量化需要校准数据通过收集激活值的分布来优化缩放系数精度损失稍大但推理更快。量化感知训练是在训练过程中就模拟量化误差模型精度最理想但改造成本高。我的建议是如果你只是想把模型体积压一压在 CPU 上用torch.quantization.quantize_dynamic最省事quantized_model torch.quantization.quantize_dynamic( model, {nn.Linear, nn.Conv2d}, dtypetorch.qint8 ) torch.save(quantized_model.state_dict(), model_quantized.pth)量化后模型大小通常能减到原模型的 1/4推理速度也明显提升。但注意量化不是银弹当模型结构里有大量对数值敏感的操作时比如 BatchNorm 在量化时表现不佳需要谨慎验证精度。我做过一个简单的 CNN 分类器量化后准确率掉了 0.3%可以接受但换到一个带注意力机制的小模型精度直接掉了 2%最后只能在速度和精度之间妥协没有做彻底量化。剪枝这里简单说一句训练后剪枝跑在 CPU 上通常能带来 30%~50% 的加速但需要配合稀疏矩阵存储才能发挥最大优势。如果你的部署环境没有特殊优化剪枝收益不如量化直接所以我一般先量化再考虑剪枝。最后再分享一个个人技巧无论最终用哪种部署方案都建议在导出前后做一次黄金样本回归测试。准备 20~50 张有代表性的图片记录原模型和导出模型的输出 logits计算平均绝对误差。这不是为了精益求精而是为了在改动任何一个参数后能快速发现变了没变、变了多少。我靠这套方法救回过好几次上线前的危机——你以为只是加了一行torch.jit.script实际某个算子已经悄悄变了行为。从第一次接触到今天我用 PyTorch 完成的项目大概可以分成三个阶段第一阶段是照猫画虎地跑通教程第二阶段是无脑换模型调参第三阶段是把更多时间花在理解机制和数据上。如果你现在正处于第一阶段不要焦虑按我上面列的步骤先跑一个完整的图像分类流程然后把每个 API 的文档翻一遍你很快会发现 PyTorch 的 API 设计极其一致需要 GPU 就是.to(device)需要梯度就设置requires_grad需要保存就state_dict。这些都掌握之后剩下的就是在具体业务场景里灵活组合的问题了。希望这篇分享能帮你少踩一些我当年踩过的坑。
