PyTorch复现经典CV算法:分类、检测、分割全流程实战指南
做计算机视觉算法调研最扎实的方式不是只看论文而是把图像识别、目标检测、语义分割三类任务的经典算法用同一套工程框架在 PyTorch 里完整复现一遍。这样每个模型算出来的指标、每个训练阶段的实际损耗、每个损失函数的作用机制都不是道听途说而是亲手验证过的。这篇文章是我近期一次“经典算法 PyTorch 复现计划”的完整技术整理覆盖三类任务中必须掌握的骨干网络、检测头和分割头、训练配置、评估脚本以及我从零跑通几十个模型后沉淀下来的坑与经验。先说结论复现经典算法不一定要追求超高精度但一定要把数据加载、模型前向、损失计算、训练循环、评估指标这五件事全部打通。我建议刚入门的同学不要只跑 torchvision 里现成的权重而是从结构最简单的 LeNet、VGG、ResNet 起步再逐步深入 Faster R-CNN、YOLO、DeepLabV3。下面我会按照“统一工程框架 → 分类 → 检测 → 分割 → 通用工程细节 → 踩坑清单”的顺序把整套复现过程拆开讲。1. 为什么把三类经典算法放到同一个框架里复现1.1 一次整理解决三个核心问题图像识别分类、目标检测、语义分割本质上是同一个视觉理解问题的三个层次。分类回答“图里有什么”检测回答“图里的东西在哪里”分割回答“哪些像素属于哪个东西”。三者的模型设计在很多层面是共享的主干网络可以直接复用特征金字塔可以搭在任意主干上数据增强、学习率策略、评估思路也高度通用。因此分开单独学习容易重复踩坑放在一起整理反而能一眼看出各个算法之间为什么长这样。我在复现时的第一个决定就是统一环境与工程目录。整个项目只用一个虚拟环境Python 3.10 PyTorch 2.1 CUDA 11.8所有脚本共用一套 config 加载方式。模型虽然不同但 train.py 的核心逻辑不做大改只切换 model、loss、dataset 三个入口。这个决定在后面几十次训练循环里非常省事每次实验只需要复制一份配置文件改几个路径和参数剩下的全部走同一套代码。1.2 复现的边界到什么程度算“复现成功”很多人以为把论文里的结构代码抄一遍、把训练跑起来就算复现了。我的标准是三件事同时满足第一模型前向输出能跑通能正常反传第二在公开小数据集上损失能稳定下降验证集指标能明显高于随机第三训练好的权重可以导出并完成独立的推理脚本不只是训练代码里的一个附属品。这三个标准里第二条最重要。复现算法的核心目标不是达到论文里的 90% 准确率而是理解“为什么它能学出来”。所以我尽量不直接加载官方预训练权重而是从随机初始化开始训练哪怕在小数据集上先跑通流程。这个过程中的梯度传播、BN 统计量更新、学习率衰减节奏才是动手复现真正能学到的东西。2. 图像识别分类经典算法的复现要点2.1 分类任务工程骨架不要重复发明轮子图像分类是所有视觉任务的基础。我在项目里先搭了一个极简的训练循环之后所有模型都复用这套骨架。核心逻辑只有五步取一批数据前向得到 logits计算交叉熵损失反向传播更新参数定期记录 top-1 与 top-5。用 PyTorch 写起来很直接for images, labels in train_loader: images, labels images.cuda(), labels.cuda() outputs model(images) loss nn.CrossEntropyLoss()(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step()这段代码虽短却包含了复现任何分类算法的完整闭环。我强烈建议在这个阶段就把分布式、混合精度、学习率调度都先放一边只跑最朴素的流程确认模型和数据没问题后再逐步加入优化。真正常用的训练配置其实也简单SGD 优化器、momentum 0.9、weight_decay 5e-4、初始学习率 0.1CIFAR 系或 0.01小数据微调配合 CosineAnnealingLR 或 MultiStepLR。2.2 从 LeNet 到 ResNet结构复现的取舍逻辑图像分类的经典网络里我最推荐复现四条主线。LeNet 是最早的 CNN结构极其简单适合用最小的计算量把卷积、池化、全连接这条链路跑通AlexNet 引入了 ReLU、Dropout 和双 GPU 训练让当时的研究方向转向了深度激活函数VGG 的核心贡献是证明小卷积核堆叠能替代大卷积核用 3x3 卷积反复叠加就能大幅加深网络ResNet 通过残差连接解决了深层网络的退化问题是真正意义上让“深度”变得可行的关键设计。ResNet 残差块是必须亲手实现的部分它的核心代码非常短class BasicBlock(nn.Module): def __init__(self, in_ch, out_ch, stride1): super().__init__() self.conv1 nn.Conv2d(in_ch, out_ch, 3, stride, 1, biasFalse) self.bn1 nn.BatchNorm2d(out_ch) self.conv2 nn.Conv2d(out_ch, out_ch, 3, 1, 1, biasFalse) self.bn2 nn.BatchNorm2d(out_ch) self.shortcut nn.Identity() if stride ! 1 or in_ch ! out_ch: self.shortcut nn.Sequential( nn.Conv2d(in_ch, out_ch, 1, stride, biasFalse), nn.BatchNorm2d(out_ch) ) def forward(self, x): out torch.relu(self.bn1(self.conv1(x))) out self.bn2(self.conv2(out)) out self.shortcut(x) return torch.relu(out)这里有个关键细节残差连接分支里的 1x1 卷积必须保持输出通道与主分支一致否则通道数对不上相加时会报错。我当时在复现 ResNet50 的 Bottleneck 结构时因为 1x1 卷积的通道写错前向检查一直报 shape mismatch。后来养成了一个习惯写完每个模块先用随机输入跑一次 shape 检查再放进训练循环。2.3 分类训练的细节预处理、优化器、学习率和正则复现分类算法真正拉开精度差距的不是网络结构本身而是训练技巧。我常用的数据增强组合是 RandomResizedCrop、RandomHorizontalFlip然后做标准化。CIFAR 这类小图要先把图像补零到合适尺寸ImageNet 级大图则要注意 resize 与 crop 的比例。标准化的 mean 和 std 直接用数据集官方统计值不要自己随便编一组否则收敛速度肉眼可见地变慢。优化器方面SGD 搭配 momentum 在普通分类任务上表现很稳AdamW 在需要快速验证模型时更省心。学习率调度我用 CosineAnnealingLR 比较多因为它在 epoch 数充足时能平滑下降不需要手动设置几个里程碑。需要注意配置了 weight_decay 之后BatchNorm 层里的 bias 和 scale 参数通常可以不衰减这一步操作对一些轻量网络的精度影响很明显。还有一个提高分类鲁棒性但经常被忽略的小技巧标签平滑。把 hard label 从 0/1 改成 [0.1, 0.9] 这类软标签能有效减少过拟合尤其在数据量不足时top-1 指标往往能高出零点几个点。MixUp 或 CutMix 这类增强在复现不用重新实现PyTorch 里都有现成工具但要注意它改变了标签分布损失计算前的 label 也需要一并做线性组合。2.4 分类评估top-1 和 top-5 怎么算才规范分类评估不需要引入额外依赖用一行 PyTorch 就能完成 top-k 统计。我的评估脚本在验证阶段记录每个 batch 的 top-1 预测对错和 top-5 预测对错最后除以总样本数。代码实现可以这样写def accuracy(output, target, topk(1, 5)): maxk max(topk) batch_size target.size(0) _, pred output.topk(maxk, 1, True, True) pred pred.t() correct pred.eq(target.view(1, -1).expand_as(pred)) res [] for k in topk: correct_k correct[:k].reshape(-1).float().sum(0) res.append(correct_k.item() * 100.0 / batch_size) return res这里有个容易踩的细节统计 top-5 时如果类别总数少于 5比如某些小数据集只有 3 类top-5 就失去意义得先判断一下 maxk 与类别数的关系。另一个建议是评估时一定要关闭梯度、切到 eval 模式否则 BatchNorm 的 running_mean 会被当前 batch 影响验证指标会忽高忽低。3. 目标检测经典算法复现核心3.1 检测任务的三种主流路线图像分类解决的是“整张图是什么”目标检测则要在分类基础上增加“目标在哪”的定位输出输出格式是多个边界框 类别 置信度。PyTorch 体系里主流检测算法可以按三种路线区分两阶段、单阶段、无锚框。路线代表算法核心思路精度速度两阶段Faster R-CNNRegion Proposal RoI 分类回归高较慢单阶段SSD、YOLO 系列回归加分类一步到位中高快无锚框FCOS、RetinaNet逐像素预测目标中心与边界高中我在复现时建议按这个顺序推进。Faster R-CNN 的概念虽然繁琐但它是理解 anchor、RoI、proposal 这些核心概念的根SSD 和 YOLO 能让你明白“直接在特征图每个位置预测边界框”为什么更快FCOS 则是近年很多检测器的基础思想值得单独梳理。3.2 两阶段Faster R-CNN 的结构与损失细节Faster R-CNN 的核心由两部分组成。RPN 网络先在主干特征上生成一组候选区域并给出每个候选是前景还是背景的预测RoI 头再对这些候选区域做精细分类和边界框回归。这种“先粗后细”的设计让两阶段检测器在精度上天然占优。复现时最容易出问题的是 anchor 的匹配逻辑。RPN 会给特征图每个位置生成多个不同尺寸和长宽比的 anchor然后与真实标注框计算 IoU。IoU 大于某个阈值通常 0.7的 anchor 作为正样本小于另一个阈值通常 0.3的作为负样本中间部分直接忽略。阈值选得不对会让正负样本极不平衡训练初期 loss 完全压不下去。Faster R-CNN 的损失由四部分相加RPN 的分类损失、RPN 的回归损失、RoI 头的分类损失、RoI 头的回归损失。回归损失常用 SmoothL1它对离群点不像 L2 那样敏感。我用一个小例子说明它的作用当预测框和真实框差距很大时SmoothL1 的梯度不会无限增大训练过程更稳定。def smooth_l1_loss(pred, target, beta1.0): diff torch.abs(pred - target) loss torch.where(diff beta, 0.5 * diff ** 2 / beta, diff - 0.5 * beta) return loss.mean()我把这个函数单独放在一个脚本里因为后面分割和其他任务也可能用到类似形式的鲁棒损失。3.3 单阶段SSD 与 YOLO 系列如何演进SSD 的核心思想是“在不同的特征图尺度上预测不同大小的目标”。浅层特征图分辨率高适合检测小目标深层特征图语义信息丰富适合检测大目标。每个特征图位置预置一组 default box模型预测的是这些预置框相对目标框的偏移量和类别概率。YOLO 系列则是把检测问题彻底变成“一张图的回归问题”。它把图像划分为网格每个网格预测若干个边界框和类别概率。从 YOLOv3 开始引入了多尺度预测和 anchor 机制YOLOv5 在工程上做了大量调优v8 则又回归到了无锚框风格。复现 YOLO 系列时我推荐直接看最新的训练配置重点不在于背数据增强花活而在于理解它的标签分配和损失函数设计。YOLO 系列最容易被忽略的是正负样本匹配。目标框内的位置才被当作正样本参与训练周围背景作为负样本如果直接让所有位置都参与损失计算背景数量会压倒性占优模型会退化成一个只会预测背景的傻子。所以无论你复现哪个 YOLO 版本都要先仔细阅读它匹配正样本的代码。3.4 无锚框与进阶FCOS 和 Focal LossFCOS 的核心思想很直观对特征图上的每个点如果这个点落在某个目标的框内就预测它到目标框四条边的距离同时预测类别。由于一个点可能落在多个重叠目标的框内FCOS 引入了 centerness 分支让模型学习“当前点离目标中心有多近”中心点越近置信度越高同时在测试时抑制低质量的边界框。RetinaNet 与 FCOS 都离不开 Focal Loss专门解决单阶段检测里“简单负样本太多”的问题。标准交叉熵会让大量易分类的背景样本主导梯度Focal Loss 通过一个调制因子降权易分类样本让模型把注意力集中在那些被分错或置信度不高的样本上。它的形式是FL(p_t) -alpha_t * (1 - p_t)^gamma * log(p_t)其中 gamma 通常取 2alpha_t 用于平衡正负样本。训练中我实际观察到的效果非常明显不用 Focal Loss 时模型前 10 个 epoch 的 loss 主要在背景上下降前景目标的召回率一直上不去换成 Focal Loss 之后验证集 mAP 能快速拉开差距。复现时不要把 Focal Loss 想得太复杂它只是在标准 CE 前面乘了一个“置信度越接近 1、权重越低”的系数。3.5 检测评估NMS、IoU 与 mAP 的完整计算目标检测的评估比分类复杂得多。测试时模型会输出大量重叠的候选框需要先用 NMS非极大值抑制去掉重复框。NMS 的流程是把所有候选框按置信度排序保留最高分的框删除与它 IoU 大于阈值的其余框然后重复这个过程。NMS 阈值通常取 0.45 到 0.5阈值取小了会漏掉重叠目标取大了又容易保留太多重复框。mAP 的计算核心是对每个类别分别计算 Precision-Recall 曲线然后求 PR 曲线下的面积。计算时先把所有预测框按置信度从高到低排序逐个与真实框匹配每匹配到一个正确检测就算一个 TP没有匹配上的就算 FP。生成 PR 曲线后AP 可以通过插值法计算。COCO 风格的 mAP 会在多个 IoU 阈值0.5 到 0.95步长 0.05下分别计算 AP再求平均。我在复现 mAP 时的经验是先拿几个已知结果的简单场景做单测比如只有两个框和一个真实框手动算一遍然后对拍脚本输出确保评估代码没问题。否则后续每次模型改进你都无法确定是模型变好了还是评估代码算错了。4. 语义分割经典算法复现要点4.1 分割任务的独特性像素级分类也有自己的难点语义分割要把图像中的每个像素归类到某个语义类别输出是一张和原图分辨率相同、通道数等于类别数的概率图。它和分类共享主干网络但为了让空间信息在深层特征中不丢失分割网络必须带上采样结构。复现分割算法的难点通常不在模型本身而在三个地方类别不平衡、边界不清晰、mIoU 评估容易写错。先说类别不平衡。在自动驾驶场景里道路和背景像素数量可能占 90%行人只占极少一部分。直接用交叉熵会让模型把所有像素都预测成背景类。所以我在分割项目里通常会同时使用加权交叉熵和 Dice Loss。Dice Loss 直接优化的是预测区域与真实区域的交并比对小目标的敏感性明显好于朴素的交叉熵。4.2 FCN 与 U-Net编解码结构的基础FCN全卷积网络是语义分割的奠基之作核心想法是把分类网络最后的全连接层替换成卷积层然后通过转置卷积或上采样恢复像素级的分类结果。它的局限在于上采样过程过于简单小目标的边界信息丢失严重。U-Net 的改进在于引入了“编码器-解码器 跳跃连接”结构。编码器逐层下采样提取语义解码器逐层上采样恢复空间分辨率跳跃连接把编码器各层的细节特征直接拼到解码器对应层让小目标和边缘信息能沿着短连接传递到最终输出。即使今天出了很多新模型U-Net 依然是小数据集分割任务里最稳的起步解尤其适合医学影像、卫星影像这类样本量不大的场景。复现 U-Net 时要注意通道数对齐。跳跃连接需要把编码器某一层的输出通道与解码器对应层拼接起来如果通道不一致要么先 1x1 卷积对齐要么在解码器部分缩减通道。我经历过几次因为通道数对不上而 debug 到深夜强烈建议每个模块输出后打印 shape。4.3 DeepLabV3 与 PSPNet空洞卷积与多尺度上下文分割网络最大的矛盾是下采样能扩大感受野但会损失分辨率。空洞卷积可以在不增加参数量的前提下扩大感受野让特征图保持较密的空间信息。DeepLab 系列利用空洞卷积构造了 ASPP 模块用多个不同 dilation rate 的并行空洞卷积捕获不同尺度的上下文信息再通过 1x1 卷积融合。PSPNet 走的是另一条路用金字塔池化模块把特征图划分为多个区域对每个区域做池化后再上采样由此聚合全局和局部的上下文。复现这两类模型时核心难点都是对齐多分支输出的 shape。ASPP 里的不同空洞率分支并行输出后要 concat 到一起如果输入分辨率设计得不合适某个分支的 padding 算错张量尺寸就会出问题。我在复现中常用的一招是先把输入固定为 512x512定义好所有卷积的 padding 和 dilation然后用一个测试输入检查每一层输出 shape确认无误后再接入训练循环。这种“前向先行”的调试习惯能省掉大量训练中期才发现的结构错误。4.4 分割训练技巧采样、损失与后处理分割训练时我建议把图像 Resize 到固定尺寸或者采用随机裁剪后拼接的采样策略。因为显存有限完整分辨率常常训练不动。RandomCrop 会让不同的 batch 覆盖不同区域相当于一种天然的数据增强但需要保证裁剪尺寸足够大否则小目标可能永远学不到。分割损失方面除了交叉熵和 Dice我经常加一个边界权重项让模型更关注目标边缘像素。实现方法不复杂用拉普拉斯算子或形态学运算提取标签中的边界区域给边界像素的损失乘以一个高于背景的系数。这个技巧在视觉效果上提升很大边界区域的分割不会再出现一条粗粗的“糊带”。后处理方面条件随机场CRF曾经很流行但现在多数场景直接用模型输出的概率图取 argmax 就够用。我的建议是把 CRF 当作可选项不要让它成为整个流程的依赖因为它的运行速度较慢对指标提升也有限。4.5 分割评估mIoU 计算的正确姿势mIoU 的计算比分类 top-1 容易出错。最稳妥的做法是利用混淆矩阵。对每个像素预测类别和真实类别会落在混淆矩阵的一个位置上最后统计所有像素再对每个类别计算 IoU该类别的交集像素数除以并集像素数最后取所有类别的平均值。代码可以这样写def compute_miou(pred, target, num_classes): pred pred.view(-1) target target.view(-1) mask (target 0) (target num_classes) hist torch.bincount( target[mask] * num_classes pred[mask], minlengthnum_classes * num_classes, ).reshape(num_classes, num_classes).float() inter torch.diag(hist) union hist.sum(1) hist.sum(0) - inter iou inter / (union 1e-6) return iou.mean().item()每次评估前先确认目标掩码的取值范围如果有 ignore_index比如 255 表示忽略像素一定要在统计前过滤掉。这个细节我在第一次复现分割模型时忽略了导致 mIoU 比预期低了一大截排查了很久才发现是背景像素的类别编号被多算了。5. 复现过程中的通用工程细节5.1 数据管线Dataset、DataLoader、Collate 与多进程无论是分类、检测还是分割数据管线都值得认真设计。我的统一做法是每个任务写一个 Dataset返回的样本格式保持一致然后用 DataLoader 加载。检测和分割任务里一张图对应的标注数量不等所以要用自定义 collate_fn 把图像和标注分别 pad 成 batch而不是让 DataLoader 自动堆叠。多进程加载的 worker 数量要合理设置。worker 太少会让 GPU 等待数据worker 太多会吃满 CPU 内存。我在服务器上常用 4 到 8 个 worker并用 pin_memoryTrue 加速数据传输。还有个容易忽略的点如果自定义 Dataset 里做了大量随机增强务必在 fit 阶段确保随机性不会导致数据分布偏移过大比如缩放后不标准化。5.2 训练稳定化的关键配置初始化、BN、AMP 与 Checkpoint复现时不要直接跳过参数初始化。PyTorch 自带的默认初始化虽然能跑但很多算法会明确写如何初始化卷积层和 BN 层。常见做法是卷积层用 Kaiming 初始化BN 层 weight 初始化为 1、bias 初始化为 0。不重视初始化的话训练前期 loss 可能一直降不下来或者干脆直接 NaN。BatchNorm 的坑主要表现为 batch size 太小时统计量不稳定。如果显存只允许小 batch就要考虑使用 SyncBN或者手动控制虚拟 batch size。另一个容易被忽略的点是模型在训练模式与评估模式之间切换时BN 的 running_mean 和 running_var 的更新行为完全不同评估前必须 model.eval()否则结果偏差很大。混合精度训练我推荐直接开启。PyTorch 自带 AMP 接口只需要在训练循环里加入 GradScaler既显存减半速度通常也有提升。但要注意AMP 开启时 BatchNorm 的计算精度可能出现细微变化如果发现指标异常波动可以先关掉 AMP 对比一次。Checkpoint 的保存要覆盖两个维度最好每 N 个 epoch 保存一次完整模型和优化器状态同时维护一个 best 模型副本。只保存模型权重不保存优化器状态模型中断后重新训练会丢失优化器动量信息后续恢复的训练效果会变差。5.3 可视化与调试TensorBoard 和推理脚本复现时我强烈建议记录训练损失、验证指标、学习率这三条曲线。TensorBoard 可以直接记录标量与图像还能可视化卷积核和特征图对判断模型是否学到了合理特征帮助很大。我习惯在每个 epoch 结束时把训练 loss 和验证指标写进同一个 log方便对比不同配置的差异。推理脚本应该从 Checkpoint 恢复模型权重后独立完成“载入图片 → 预处理 → 前向 → 后处理 → 保存可视化结果”的完整流程。这个脚本不需要包含训练逻辑它的作用是验证模型的可部署性。我在复现目标检测时专门写了一个脚本把检测框、类别标签和置信度画到原图上保存下来肉眼对比结果比只看 mAP 更有直觉。6. 踩坑清单与实用经验6.1 高频踩坑速查表下面这些坑是我在复现三类任务时遇到频率最高的问题。整理成速查表遇到类似情况可以直接照方抓药。现象可能原因排查与解决损失出现 NaN学习率过高、AMP 溢出、标签越界调低学习率、关闭 AMP、检查标签范围loss 一直不下降正负样本失衡、优化器配置不对、BatchNorm 未设置检查损失函数权重、优化器动量、调大 batch size验证指标波动大评估时没有 model.eval()、数据标准化错误切 eval 模式核对 mean/std检测框全部为空NMS 阈值过低、置信度阈值过高调低置信度阈值先打印出所有预测框信息分割 mIoU 偏低ignore_index 没过滤、类别编号错位检查混淆矩阵统计范围显存溢出特征图分辨率太大、batch size 太大下调输入尺寸、使用 AMP、梯度累积6.2 几条经过实践验证的复现顺序建议先跑通小数据集。我第一次复现 Faster R-CNN 时直接在完整数据集上跑调试周期长到让人崩溃。后来改成在 20 到 50 张小图上跑 10 个 epoch能 10 分钟内验证整个流程是否正确再换到大数据集。这个策略对检测和分割尤其有效。把模型结构写完后先随机初始化跑 10 个 batch观察 loss 是否能下降再开始正常训练。固定随机种子。复现算法的可对比性非常重要我所有脚本开头都会加上固定 seed 的操作包括 Python、NumPy、PyTorch 和 CUDA 随机种子。虽然这样不能完全消除非确定性操作但在大部分实验里能保证结果基本可复现。如果两个实验配置相同但指标差很多先检查是否忘了固定种子。保持代码模块化。把 model、dataset、loss、trainer、evaluate 分成独立的 Python 模块比堆在同一个训练脚本里好维护一万倍。我的目录结构很简单model 文件夹下放每个模型的实现config 文件夹下放不同算法的配置eval 文件夹下放评估脚本。这样切换算法时只需要新建一个 config 文件不用改主流程代码。最后再分享一个我个人的习惯每次跑通一个模型我会在项目根目录建一个 README记录该模型的训练配置、最终指标和踩过的坑。这份记录比代码本身更有价值因为三个月后你回头来看真正能帮你快速找回上下文的不是模型的 forward 函数而是当初写下的那些“为什么这么配”的备注。复现经典算法的意义正在于此它不只是一次跑通实验而是让你把视觉任务的技术演变脉络和工程化细节真正攥在手里。