1. 先说清楚这个项目到底想用PyTorch解决什么问题我最近在做一个三分类的图像识别小项目场景很朴素给一批产品照片自动分到三个类别里。数据量不大每类大概四五百张加起来勉强过千。模型不需要做到工业级精度但至少要能跑通一条完整的链路——从原始图片文件夹到训练出一个能用的模型再到把模型导出成可部署的形式。选PyTorch而不是其他框架首先是因为它的调试体验。训练循环是自己一行行写出来的数据在什么时候进到哪个张量、梯度在哪个环节算出来都是显式的。对一个小项目来说这种看得见摸得着的感觉特别重要。其次是生态HuggingFace那边的预训练模型基本都用PyTorch后面如果想换更好的主干网络操作成本很低。这篇东西不是教科书也不是官方教程的翻译。我会把我搭建这个项目时踩过的坑、验证过的做法、还有排查问题的顺序都记录下来给同样准备用PyTorch做实际任务的读者一个参照。如果你只是想看一个API怎么用官方文档更合适如果你想知道一个完整的项目是怎么从零搭起来的这篇文章应该能帮到你。先交代一下我最终跑通的软硬件环境后面所有操作都是在这个基础上说的显卡是某款8GB显存的消费级卡显存不算大所以Batch Size设置得很谨慎PyTorch版本用的是2.x系列CUDA 11.8操作系统是Windows 11Python环境用虚拟环境管理Python版本3.10这个配置并不特殊你的机器哪怕没有独立显卡把代码里的cuda改成cpu流程也完全成立只是速度会慢不少。再提醒一句我用的是三分类图像任务作为示例但下面的大部分内容——尤其是数据加载、训练循环、排查思路——对回归任务、多标签任务、甚至简单的文本分类任务都是通用的你只需要把Dataset和模型输出层对应改一改就行。2. 环境与数据准备真正耽误时间的往往不是模型代码很多初学者以为搭建PyTorch项目最耗时的是写模型实际上我这次的经验是反过来的环境配置和数据整理占了整个项目一半以上的时间。这一节我详细说一下我做的事情和踩过的坑。2.1 环境配置先用CPU把流程跑通再说GPU我的安装过程其实不复杂但有一个习惯我觉得值得推荐先不装CUDA版用CPU模式把代码跑通再换GPU版本。原因有二。第一CPU版本安装简单几乎不会遇到驱动不匹配、cuDNN缺失这类问题第二代码逻辑的错误在CPU模式下暴露得更直接报错信息不会和CUDA设备相关的问题混在一起。我这里说的跑通指的是从数据加载到模型前向传播再到反向传播完整走一遍哪怕loss值是随机水平也没关系。这一步过了再执行带CUDA的版本安装把.to(cuda)相关的代码加进去通常就顺畅多了。还有一个小细节Windows下安装PyTorch的时候官方命令里给的--index-url参数一定要带上。我见过不少人直接pip install torch结果装成了CPU版本后面训练时才发现用不了显卡再折腾一遍。浪费的时间完全可以从一开始就避免。2.2 文件目录的组织方式数据这块我推荐的最简结构是这样dataset/ ├── train/ │ ├── class_a/ │ ├── class_b/ │ └── class_c/ ├── val/ │ ├── class_a/ │ ├── class_b/ │ └── class_c/ └── test/ ├── class_a/ ├── class_b/ └── class_c/把训练集、验证集、测试集按照类别分文件夹放好然后直接用torchvision.datasets.ImageFolder读取几行代码就能把数据集和标签都搞定。不需要自己写复杂的CSV解析逻辑也不需要手工维护一个images.txt再单独配一个labels.txt。我见过有人把数据和标签分开存成数组训练的时候对齐出了问题最后差了一个样本整整排查了一个下午。文件夹即标签这个方案虽然土但它是视觉项目里最不容易出错的原始数据组织形式。2.3 DataLoader的四个参数每个都有讲究在我这个项目里DataLoader配置是这样的from torch.utils.data import DataLoader train_loader DataLoader( train_dataset, batch_size32, shuffleTrue, num_workers0, pin_memoryTrue ) val_loader DataLoader( val_dataset, batch_size32, shuffleFalse, num_workers0, pin_memoryTrue )有几个参数值得展开说。最大的坑是num_workers。在Windows平台上num_workers如果设成大于0经常会出现子进程反复重启的报错而且训练前期的卡顿非常严重。我在Windows上跑了两次都发现num_workers设成2以上之后数据加载的速度不仅没有变快反而因为进程反复创建让整个程序在启动时卡了好几分钟。最后直接统一设成0让数据加载在主进程里跑反而最稳定。如果你的环境是Linux服务器num_workers可以设为4或者8Windows用户建议就从0开始。pin_memory我建议打开。它的作用是把数据锁在固定内存区域这样CPU向GPU拷贝数据的时候会省一步。实际体感虽然谈不上变快了一个量级但在数据加载频繁的情况下训练过程会更顺畅。唯一的代价是占用的CPU内存会比不开启时高一些8GB内存的机器需要留意。shuffle参数和batch_size同样值得注意。很多人只在训练集上开启shuffle验证集不开这个做法是对的。Batch Size在8GB显存下我只能设为32因为输入图片的尺寸是224×224显存占用比想象中要高。如果你用的是更大的图片输入比如384×384Batch Size可能需要降到16甚至8否则会直接爆显存。这里有一个简单的爆显存判断标准如果程序运行到某一轮batch时弹出一个CUDA out of memory而且任务管理器里显示GPU的专用显存占用到了99%那基本就是Batch Size设太大把它调小一半再试十有八九能解决。2.4 数据增强和归一化宁可少做不要做乱数据增强策略我一开始参考ImageNet的标准方案就是随机裁剪加水平翻转from torchvision import transforms train_transforms transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) val_transforms transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])这里我踩了一个细节坑验证集和测试集不能使用训练集的随机增强方案这会导致同一张图片在不同时刻的评估结果不一样。正确做法是验证集只做固定尺寸的缩放和中心裁剪。Mean和Std我用的是ImageNet预训练模型的默认值因为我用的是在ImageNet上预训练过的模型权重。如果你要完全重新训练一个网络需要自己计算的遍历所有图片统计每个通道的均值和标准差否则用ImageNet的统计值反而会违背你的数据分布。但多数情况下既然要迁移学习直接用预训练模型的统计值是最合理的选择。3. 模型搭建的核心逻辑迁移学习与其把每个API都搞懂不如把结构想通模型部分我只讲我的核心决策——使用预训练权重进行迁移学习以及背后为什么这样做。3.1 为什么不等量齐观地从头训练我这个项目的训练数据只有一千多张从头训练一个ResNet级别的网络在这点数据量下很容易过拟合训练起来也慢。而使用在ImageNet上预训练好的权重我就等于站在了已经见过一千万张图的基础上只需要微调最后的分类层。这个思路适合大多数中小规模数据项目也是实际工程中最常见的做法。选择具体模型时我用的是resnet34。原因很直白它的参数量比ResNet50小不少推理速度快三层分类任务根本不需要五十层那么深的能力。在你自己的项目里如果目标类别特别相似比如质检里区分划痕类型那可以考虑更大的模型否则ResNet18或ResNet34是性价比很高的起点。加载模型的方式如下import torch import torch.nn as nn from torchvision import models model models.resnet34(weightsmodels.ResNet34_Weights.IMAGENET1K_V1) num_classes 3 model.fc nn.Linear(model.fc.in_features, num_classes)在旧版代码里会看到pretrainedTrue这种写法现在PyTorch官方推荐用weights参数因为这样预训练权重的来源会被显式记录在模型的元信息里用起来更规范。model.fc是ResNet最后的全连接层我看过它源码之后知道它的输入维度就是前面主干网络特征输出的维度所以直接把原来的1000类输出替换成3类输出就行。注意力放在这个替换最后一层的动作上比逐行读网络结构源码要省时间得多。3.2 冻结与微调先固定可以固定的我做过两种策略的实验。第一种是完全只训练最后一层全连接层冻结前面所有层第二种是从一开始就解冻所有层按一个很小的学习率全量微调。先说结论在数据量有限的情况下先冻结主干只训练分类头往往更稳训练时间也短。等训练若干轮之后再把主干解冻用低学习率继续微调效果通常会更好一些。但是要先做一步不要直接一上来就全量微调。冻结实现的细节是for param in model.parameters(): param.requires_grad False for param in model.fc.parameters(): param.requires_grad True这一步的本质是利用PyTorch的自动求导机制requires_grad设为False的参数不会计算梯度也就不会在反向传播时被更新。这样显存占用会小一些前向传播也更快。要注意的是我冻结的是整个主干网络但BatchNorm层建议保留训练状态。如果你做了一开始的判别之后发现精度不高可以先尝试只解冻最后一两个残差块再用小学习率训练。整体操作并不复杂关键是顺序别搞反。3.3 损失函数和优化器的选择逻辑这个分类任务用的是交叉熵损失criterion nn.CrossEntropyLoss() optimizer torch.optim.AdamW(model.parameters(), lr1e-4)为什么用AdamW而不是SGD我个人觉得在现代PyTorch使用习惯下AdamW对学习率的敏感度比SGD低调参空间更大。而且AdamW自带权重衰减也就是L2正则的改良版能帮助抑制过拟合。SGD不是不能用只是它需要你更仔细地调学习率和动量参数初学阶段用AdamW不容易把训练带崩。学习率从1e-4开始算是我的个人经验。如果你用的是多个GPU或更大Batch Size可能需要调小到5e-5。Batch Size翻倍时学习率通常也翻倍不过真到了那个阶段我建议跑一个小脚本做对比实验会更好。4. 训练循环的写法与记录它比你想的更值得花时间训练循环是整个项目里看起来最不起眼的部分但也是影响项目可复现性的关键。我来写一下我实际的写法以及为什么要这么写。4.1 训练一个epoch的标准循环def train_one_epoch(model, loader, criterion, optimizer, device): model.train() total_loss 0.0 correct 0 total 0 for images, labels in loader: images images.to(device) labels labels.to(device) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() total_loss loss.item() * images.size(0) _, predicted outputs.max(1) total labels.size(0) correct predicted.eq(labels).sum().item() avg_loss total_loss / total avg_acc correct / total return avg_loss, avg_acc有几个小点值得提。optimizer.zero_grad()必须在loss.backward()之前执行否则梯度会累加。这是PyTorch默认行为如果你不想清空上一次的梯度就是另一套写法。一般来说默认做法是每个batch都要调用。我之前为图省事把它放在loss.backward()之后结果acc一直不涨后来发现梯度跨batch累加数值早就乱套了。.item()方法必须显式调用如果你直接把loss这个张量塞进列表训练过程中会产生一个不断累加的缓慢内存泄漏。这个问题在长训练里很致命在短训练里也不该有。建议每次记录loss时都加.item()。model.train()和model.eval()的切换也重要。train()会让Dropout生效BatchNorm也处于更新统计量的状态eval()则相反。验证阶段忘了切eval()就很可能发现验证准确率不稳定因为模型的行为不一致。4.2 验证循环和训练循环的区别def validate(model, loader, criterion, device): model.eval() total_loss 0.0 correct 0 total 0 with torch.no_grad(): for images, labels in loader: images images.to(device) labels labels.to(device) outputs model(images) loss criterion(outputs, labels) total_loss loss.item() * images.size(0) _, predicted outputs.max(1) total labels.size(0) correct predicted.eq(labels).sum().item() avg_loss total_loss / total avg_acc correct / total return avg_loss, avg_acc验证循环和训练循环最大的区别是with torch.no_grad():。这个上下文管理器会让PyTorch跳过梯度的计算和存储。验证阶段不需要回传更新参数所以开了它之后显存占用明显下降速度也会快很多。如果不开模型照样能算loss但你可能会在某个时刻发现显存被验证阶段无情占满。在我的项目里最终保存的模型不是最后一个epoch的结果而是验证集上准确率最高的那一次。为了做到这一点我在每个epoch结束后比较验证acc只有当它超过历史最优时才保存。这一点对模型质量和最后的测试指标影响很大不要省。4.3 日志记录别只盯着终端刷数字我训练的时候会把每个epoch的train_loss、train_acc、val_loss、val_acc统一打印到控制台同时把人眼方便阅读的内容记录到CSV文件里。为什么CSV就够了因为CSV可以直接用Excel打开也可以被pandas读取画图简单直接。如果你想要更丰富的信息可以用TensorBoard。用法非常小from torch.utils.tensorboard import SummaryWriter writer SummaryWriter(runs/exp1) # 在训练循环里 writer.add_scalar(Loss/train, avg_loss, epoch) writer.add_scalar(Loss/val, val_loss, epoch) writer.add_scalar(Acc/train, train_acc, epoch) writer.add_scalar(Acc/val, val_acc, epoch)训练结束后运行tensorboard --logdirruns即可。我建议至少把以下四个指标记录完整训练loss、验证loss、训练acc、验证acc。如果你做过这两个指标的组合对比你会很快发现一个常见现象训练loss一直在降验证loss却涨了这通常就是过拟合的信号。没有日志记录这种判断只能靠直觉有了日志你就能确切地定位到是从第几个epoch开始的。5. 训练不收敛的排查顺序这不是玄学是一套固定流程训练过程中遇到最大的问题是第9个epoch后模型在验证集上的准确率停在0.85左右不再上升。我排查的过程很繁琐但也很有价值分享出来。5.1 经验排查顺序遇到这类瓶颈我的排查链条是固定的第一步确认代码逻辑没有问题。比如model.train()是否忘了写labels是否和inputs对得上损失函数输入顺序是否反了。这步最基础但出错率最高的人往往就卡在这一步。第二步检查学习率。如果loss呈现剧烈的震荡训练acc也在一个区间来回横跳十有八九是学习率过大。这时把学习率调低为原来的五分之一或十分之一往往训练就稳定了。第三步看是过拟合还是欠拟合。训练acc很高比如90%以上验证acc较低比如80%差值持续拉大这意味着过拟合。处理办法增加数据增强、加权重衰减、使用Dropout。如果训练acc和验证acc都低比如都只有五六十这说明模型欠拟合应当加大模型容量、调低学习率、增加训练epoch。第四步检查数据。有些人会在这里忽略最简单的检查把各个类别在训练集和验证集的分布画出来确认验证集每个类别都够数。我这次遇到的acc停滞最后发现类别比例在验证集里有偏斜某个类别的样本占比太小导致模型在该类别上泛化不好。重新划分数据集后同样的训练配置下验证acc直接提升到了0.92。5.2 我在这个项目里真正解决的两个问题第一个问题是验证集随机划分时用了不固定的随机种子。我的代码这样写的# 第一次划分 train_list, val_list random_split(dataset, [0.8, 0.2]) # 后来又一次跑实验 train_list, val_list random_split(dataset, [0.8, 0.2])看似一样但每次划分结果都不同导致两次实验的验证集不是同一份模型指标不具可比性。排查半天才发现这个巧合。修法也很简单generator torch.Generator().manual_seed(42) train_list, val_list random_split(dataset, [0.8, 0.2], generatorgenerator)这是很多人容易踩的暗坑。某个版本实验跑出0.85代码没改只改了个数据划分从头再跑变成了0.87你真以为模型进步了。其实没有。一切对比实验都必须在完全相同的数据划分上进行。第二个问题是优化器参数解冻时机。我第一阶段冻结主干只训练分类头跑5个epoch后换成了全量微调。但我在换阶段时忘记把主干参数的requires_grad重新设为True导致实际训练的还是只有最后一层。新阶段的学习率调得很低模型收敛自然慢。这个bug直接导致训练曲线出现一段横盘。解决办法是在切换阶段后加一段简单断言或者打印检查一下确认确实有参数处于requires_gradTrue的状态# 在训练开始前检查可训练参数量 trainable_params sum(p.numel() for p in model.parameters() if p.requires_grad) total_params sum(p.numel() for p in model.parameters()) print(fTrainable: {trainable_params} / {total_params})这个检查虽然只花了一行代码但它能让你在一个长训练刚开始时就发现配置错误而不是白跑一晚上。6. 模型保存、加载与推理导出最后一个坑往往是最贵的坑模型训练完之后需要把模型保存下来并在实际业务里使用或者把模型文件交给其他人。这里面也有几个需要注意的细节。6.1 保存整个模型还是保存state_dict我推荐优先保存state_dict原因有两个模型的架构代码会被分开保存所以就算以后PyTorch版本升级也不至于因为官方的类实现变化导致文件无法加载另一方面state_dict的文件大小通常更小更便于分享。torch.save({ model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), epoch: epoch, val_acc: best_val_acc, }, checkpoint.pth.tar)我不建议只保存模型权重推荐至少附带一个字典把epoch、验证acc、甚至学习率都存进去。这样做的好处是恢复训练时可以知道这个权重文件是哪个阶段产生的也便于追踪实验。6.2 推理时如果想直接用模型要怎么做需要做模型推理的时候最容易漏掉的是model.eval()。前面说过这个操作会关掉Dropout、固定BatchNorm的统计量减少推理结果的不确定性。如果有多个输入不能像训练时那样分批喂而是直接创建一个batch。很多人在单张图预测时写成了inputs.unsqueeze(0)这其实就是在构造一个batch很常见的操作。最忌讳的是把单张图的三个维度直接裸着传进模型会直接报维度错误。6.3 导出ONNX时的一个典型问题我的最终目标是能在一个做实时推理的小工具里加载模型所以需要导出为ONNX格式。导出过程本身很简单import torch.onnx model.eval() dummy_input torch.randn(1, 3, 224, 224, devicecuda) torch.onnx.export( model, dummy_input, model.onnx, input_names[input], output_names[output], dynamic_axes{ input: {0: batch_size}, output: {0: batch_size} }, opset_version17 )这里我在dynamic_axes里允许了batch维度是动态的这样推理时一次吃1张、8张、32张图都可以不需要为不同batchsize导出多份模型。但有一个问题在动态batch下非常容易出现模型在导出前必须在eval模式下进行一次真实输入的前向目的是让BatchNorm的running_mean和running_var处于确定状态。如果忘了model.eval()导出后的ONNX模型在某些推理框架下会出现很陡的推理结果波动原因就是BatchNorm统计量还在被更新。7. 经验沉淀下次再搭PyTorch项目时我会提前做好的几件事这一节聊点项目之外的体会都是经过这次跑完整个流程之后沉淀下来的东西希望对你有参考价值。第一项目目录从一开始就按运行方式分层。我会把数据预处理、模型定义、训练循环、验证逻辑、工具函数分别放在不同模块里。这种做法最直接的好处是画实验对比图时不需要在长脚本里来回翻找某个数据处理的代码。即便项目小也值得按这个习惯来。第二所有的超参数包括batch_size、学习率、epoch、图片尺寸、随机种子都集中在一个配置文件里。我之前喜欢把超参数散落在脚本的各个角落结果是每次调参都要全文搜索还容易漏改。这个项目我改成集中定义之后每次实验都清清楚楚。哪怕你只是用Python的argparse或一个简单的字典也比散落强得多。第三每个新实验都记录实验日志。我在跑实验时会按exp_001、exp_002这样给每个实验建目录里面放一份配置副本、一份CSV日志、一张训练曲线截图。这样过了几个星期之后看到一个结果打开文件夹就能知道当时的配置是什么用了什么数据划分、什么学习率完全不需要靠记忆。我自己之前就吃过这个亏跑出一个好结果三天后想复现却忘了当时用的是哪个学习率只能靠猜。第四不要害怕在一个小项目里也使用迁移学习和预训练权重。很多人觉得预训练模型是别人的模型不如自己训练来得干净但工程上能用且好用才是第一原则。只要有符合你任务的预训练权重并且你没有极其特殊的约束用它做起点几乎总是比从头训练要好。第五验证集上的评估必须和最后实际场景尽量一致。这次项目的验证集划分偏斜最后实际部署时准确率下降本质就是验证场景和实际场景不一致。尽可能让验证数据的分布贴近真实使用场景这个原则比任何调参技巧都重要。第六也是我最想强调的一点把PyTorch的训练代码写成可以随时跑完、随时重跑的状态。所谓方便重跑就是固定随机种子、记录所有超参数、保存最优权重并附带验证指标。如果一个训练脚本在调整了一个小参数之后要重新花一整晚才能复现这个脚本本身就值得改进。这个项目整体跑下来我最大的收获不是模型指标本身而是把一套可以复用到多个任务的流程打磨顺了。从一个只有图片文件夹的原始数据集开始到能导出为ONNX并用于一个推理工具这个过程里遇到的几乎所有问题在文档里都有答案但文档不会告诉你你要按这个顺序去排查。希望这篇记录能帮你少走我走过的弯路。