先聊个比较普遍的现象很多人拿到开源仓库的deeplabv3代码跑通一个预训练模型就觉得自己“会了”一旦换成自己的数据集立刻开始怀疑人生。要么是训练跑到一半loss飞了要么是验证出来mIoU低得没法看要么干脆被一堆看不懂的配置参数劝退。我最初做语义分割项目时也是这样——想直接白嫖别人的训练脚本结果光是搞懂数据加载和checkpoint恢复就花了好几天。后来我干脆抛弃“先clone后改”的思路从零手写了一套基于PyTorch的deeplabv3训练流程反而顺了很多。这篇文章就是我当时整个过程的完整复盘从模型结构怎么拆解、数据集怎么准备到训练细节怎么调、踩过的坑有哪些全部记录下来。适合的人群很明确已经会基本的PyTorch训练流程但对分割模型只有一个模糊概念想真正吃透deeplabv3并用在自己数据集上的人。即便你是第一次接触语义分割照着走也能把流程跑通。1. 先想清楚deeplabv3到底改了什么再动手写代码很多教程上来就贴模型代码但我觉得不先理解结构就直接抄代码后面调参时会非常被动。deeplabv3这个模型说白了是在解决语义分割最核心的两个矛盾一是要保住图像的空间细节不然物体边界会糊二是要扩大感受野不然大物体内部的语义会丢失。它给出的方案就是Encoder-Decoder结构加一个叫ASPP的特殊模块。1.1 空洞卷积的本质分辨率与感受野的折中普通卷积在做下采样时比如stride2特征图尺寸减半感受野变大但位置信息也在不断丢失。这在分类任务里无所谓到了像素级的分割任务里就是灾难——经过5次下采样后原图32x的细节早就没了。deeplabv3的做法是在backbone的最后两个stage里不进行真正的下采样而是使用空洞卷积来保持感受野的同时维持较高的分辨率。空洞卷积的核心概念可以用一句话概括在卷积核元素之间插入空洞让同样大小的卷积核“看”到的范围更大但不增加参数量。比如一个3x3卷积rate2时感知范围等效于5x5rate4时等效于9x9。早期backbone比如ResNet默认是连续下采样得到最终输出stride32而deeplab系列把这个过程改成最终输出stride16或stride8。输出stride越小feature map的空间尺寸越大细节信息保留越多代价是计算量和显存占用上升。我最终选择了output_stride16因为它在精度和显存消耗之间最均衡。如果要换到8ASPP部分的输入分辨率翻倍显存直接多出一大截收益却很有限。1.2 ASPP和Decoder各司其职ASPPAtrous Spatial Pyramid Pooling的思路很朴素也很有效用不同空洞率的卷积并行地去“看”同一个特征图然后把结果堆在一起。为什么需要多尺度因为图像里的物体会出现在不同尺寸上。一个road sign可能在近距离时占满半个画面在远距离时只有几十个像素。如果只用同一个感受野去处理模型很难同时兼顾这两种情况。deeplabv3的ASPP由四路构成我在代码里也原样复刻了一个1x1普通卷积负责提取通道维度的全局语义三个3x3空洞卷积rate分别取6、12、18覆盖从小到大三个尺度的感受野一个图像级平均池化分支把整张feature map压成1x1再恢复相当于给模型提供一个“全局上下文”这四路输出concat之后经过一个1x1卷积融合通道得到Encoder的最终输出。到这里模型其实只完成了一半——它输出的特征图分辨率仍然只是原图的1/16直接上采样回去做预测边界一定糊成一团。Decoder的工作就是补救这个边界问题。实现细节是把backbone中较浅层的特征通常是stage3的输出也叫low-level feature用1x1卷积压缩到48个通道然后跟经过4倍上采样的Encoder输出concat在一起再经过两个3x3卷积把通道融合、上采样4倍得到与原图尺寸一致的预测。浅层特征分辨率高但语义弱深层特征语义强但分辨率低两者叠加边界质量肉眼可见地提升。2. 把自己的数据集变成模型能吃的格式我第一次准备数据集时犯过一个很蠢的错误标注图保存成png但背景区域直接是黑色像素值0前景区域是白色像素值255。模型训练倒是没报错就是mIoU永远卡在低位最后排查发现是255这个值被当成了一个正常类别ID而不是忽略区域。所以数据准备这一节值得多花一点时间讲清楚。2.1 标注工具与格式选择数据格式严格来说是按类别ID的索引图每个像素的取值应该是一个整数范围在0到N-1之间N是类别数。如果你的任务有两个类别背景和裂缝那么背景像素应该是0裂缝像素应该是1而不是255。千万不要直接用标注软件导出的彩色PNG当训练标签除非你在代码里做颜色到索引的映射。我这里用一个路面裂缝分割的小数据集来演示整个流程类别只有背景和裂缝两类。标注使用普通的图像标注工具把裂缝区域用多边形标出来导出JSON格式后我写了一个小脚本把JSON转成索引掩膜图。转换逻辑很简单初始化一张全0灰度图按多边形坐标填充成对应类别ID。如果标注时存在多个类别注意填充时每个类别用独立ID并且把“不确定区域”统一设为255代码里会忽略这些像素。文件目录结构我直接仿照VOC分割数据集的样式这套结构简单清晰几乎所有分割训练代码都能通用dataset/ ├── JPEGImages/ │ ├── 001.jpg │ ├── 002.jpg │ └── ... ├── SegmentationClass/ │ ├── 001.png │ ├── 002.png │ └── ... └── train.txttrain.txt里每行写一个不带扩展名的文件名例如001。测试集按同样方式组织一份val.txt。用txt记录数据集划分的好处是想调整训练/验证比例时只改文件列表就行不需要移动任何图片文件。2.2 数据增强配置语义分割任务的数据增强比分类要谨慎原因很简单过强的颜色扰动会让分割边界失去参考过强的几何变形会让标注掩膜和图像对不上。但我实践下来下面这组增强组合基本是安全的效果提升也很明显随机水平翻转概率0.5随机缩放范围0.5到2.0图像和掩膜同步缩放随机裁剪到固定尺寸520x520颜色抖动亮度、对比度、饱和度幅度轻微归一化均值0.485, 0.456, 0.406方差0.229, 0.224, 0.225特别注意随机缩放和随机裁剪的顺序。我习惯先缩放再随机裁剪因为如果先裁剪再缩放每个batch的输入尺寸会不一致后续没法统一成一个tensor。分割任务的标注掩膜在缩放时要用最近邻插值而不是双线性插值不然0和1之间会出现模糊的中间值直接污染标签。2.3 类别不均衡问题路面裂缝数据集有个典型bug裂缝区域可能只占整个画面的不到1%背景占了99%以上。如果直接拿交叉熵训练模型大概率学出一个“什么都预测为背景”的平庸模型mIoU看似还行因为背景占大头实际毫无用处。处理方式我用了两层第一是计算每类的像素占比给损失函数加类别权重裂缝类的权重会显著提高第二是在数据加载器里加一个“裁剪后过滤”的逻辑如果裁剪出的patch里目标类别像素数低于阈值比如50像素就重新采一个patch保证每个训练样本里裂缝区域不会完全消失。这个trick很小但效果非常直接。如果你想追求更高级的做法可以试试OHEM在线难例挖掘损失它会把loss最高的20%像素单独拿出来反传本质上是在告诉模型别老盯着那些简单的大块背景了去啃难分类的边界和细小目标。我个人经验的排序是常规CE 类别权重是最稳妥的起步方案OHEM适合在基础方案已经跑通后进一步提点。3. 手写模型代码Encoder、ASPP与Decoder逐段过从零搭建的最大好处是每一行代码你都知道为什么存在。model类的完整代码这里不逐行贴但几个关键模块值得单独拆开讲因为这些地方最容易写错。3.1 ResNet主干改造把stride换成空洞卷积的switchdeeplabv3不是把ResNet原封不动拿过来用的它要修改ResNet最后两个stage的下采样行为。标准ResNet从stage2到stage5的stride分别是2、2、2最终输出是原图的1/32。如果我们想要output_stride16需要把stage4和stage5的stride都替换成1改为在3x3卷积层里使用空洞率2和4来补偿感受野的损失。这个改造我用了一个非常直观的方式实现继承标准resnet50的类然后重写forward时对特殊层单独处理。核心逻辑是下面这段import torch import torch.nn as nn from torchvision.models import resnet50 class ResNetBackbone(nn.Module): def __init__(self, output_stride16): super().__init__() self.backbone resnet50(pretrainedTrue) self.output_stride output_stride if output_stride 8: # 把最后两个stage都改成空洞模式 self._modify_layer(self.backbone.layer3, 1, 2) self._modify_layer(self.backbone.layer4, 2, 4) elif output_stride 16: # 只改最后一个stage self._modify_layer(self.backbone.layer4, 2, 4) def _modify_layer(self, layer, stride, rate): # 把该层所有残差块中3x3卷积的stride改为1并设置空洞率 for n, module in layer.named_modules(): if isinstance(module, nn.Conv2d) and module.kernel_size (3, 3): module.stride (stride, stride) module.dilation (rate, rate) module.padding (rate, rate) def forward(self, x): x self.backbone.conv1(x) x self.backbone.bn1(x) x self.backbone.relu(x) x self.backbone.maxpool(x) x self.backbone.layer1(x) low_level self.backbone.layer2(x) x self.backbone.layer3(low_level) x self.backbone.layer4(x) return x, low_level这里有个非常关键的细节low-level特征不能取layer1的输出我最终验证下来取layer2的输出效果最稳。原版论文描述的是“backbone的early feature map”官方实现用的是layer2的倒数第一层输出。channel数是512相比layer1的256能保留更多信息又不至于像layer3那样因为下采样次数太多丢失边界细节。3.2 ASPP模块的代码实现ASPP的实现不复杂每个分支的卷积我都用一个Sequential包成完整的Conv-BN-ReLU结构注意BN放在Conv之后、ReLU之前。这点看起来不起眼但对最终收敛速度影响很大如果你的代码里顺序错了训练时梯度很容易不稳定。class ASPP(nn.Module): def __init__(self, in_channels, out_channels256): super().__init__() self.branch1 nn.Sequential( nn.Conv2d(in_channels, out_channels, 1, 1, 0, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue) ) self.branch2 nn.Sequential( nn.Conv2d(in_channels, out_channels, 3, 1, padding6, dilation6, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue) ) self.branch3 nn.Sequential( nn.Conv2d(in_channels, out_channels, 3, 1, padding12, dilation12, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue) ) self.branch4 nn.Sequential( nn.Conv2d(in_channels, out_channels, 3, 1, padding18, dilation18, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue) ) self.branch5 nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Conv2d(in_channels, out_channels, 1, 1, 0, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue) ) def forward(self, x): size x.shape[-2:] b5 self.branch5(x) # 全局池化分支需要再上采样回输入尺寸 b5 F.interpolate(b5, sizesize, modebilinear, align_cornersFalse) return torch.cat([self.branch1(x), self.branch2(x), self.branch3(x), self.branch4(x), b5], dim1)最后可以把5路输出再过一个1x1卷积把通道压回256也可以不压直接进入decoder的cat阶段。从我的经验看压回256能让decoder侧的参数量明显下降而且训练更稳所以我每次都加这一步。3.3 Decoder的low-level融合decoder部分有个特别容易踩的坑low-level特征通道数太多直接concat会让decoder里的3x3卷积变成一个大计算块而且冗余信息严重。所以必须先做一次1x1卷积降维官方实现中降到48。我当时也试过降到32和64最终结果差距非常小48是一个足够稳的默认值。另一个细节是上采样的倍率。encoder输出的是1/16特征图decoder的流程是先4倍上采到1/4concat low-level此时也是1/4分辨率经过两个3x3卷积后再做一次4倍上采样回到原图。两次插值都建议使用bilinear模式并设置align_cornersFalse与PyTorch官方预训练模型行为保持一致否则在部分任务上会导致边界出现轻微偏移。4. 训练脚本设计影响最终精度的四个关键开关模型结构搭好之后训练脚本就是决定结果上限的关键。很多人盲抄repo中的超参数却不知道每个参数存在的理由。这一节我把训练阶段最重要的几个开关列出来每个都说明为什么这么设。4.1 优化器、学习率策略与warmup分割任务我推荐SGD而不是Adam原因很直接Adam虽然前期收敛快但容易在后期陷入尖锐局部极小值SGD配多项式学习率衰减在语义分割上表现更稳、最终精度更高。我的参数设定如下优化器SGD动量0.9权重衰减1e-4初始学习率0.007batch size为16时学习率策略poly指数为0.9即 ( lr lr_{init} \times (1 - \frac{iter}{max_iter})^{0.9} )前3个epoch用warmup学习率线性从0.1倍涨到初始学习率Poly策略的意思是训练越到后期学习率下降越平缓有利于在最后阶段稳定挖掘细节。有人会问为什么poly比cosine好我用同样的数据集对比过poly在分割任务上的收敛曲线更平滑尤其在训练轮次较短100轮以内时优势更明显。4.2 损失函数的取舍基础损失我用的就是带类别权重的CrossEntropyLoss。PyTorch里的CrossEntropyLoss自带ignore_index参数记得设置成255这样标注中不确定区域和padding区域都不会参与梯度计算。如果数据严重不均衡再叠加一个Dice Loss分支但注意Dice Loss在小目标上的梯度波动非常大需要在总损失里把它的权重压低到0.1左右否则训练早期会剧烈震荡。辅助损失这一块deeplabv3原生实现里没有但我在实践中会在ASPP之后加一个侧边输出把1/8分辨率的特征图接一个3x3卷积和分类器单独算一次损失与主损失按0.4:1的比例加权。这个设计能加速backbone前中期的收敛最终几乎不影响推理速度。我的经验是数据量越少这个aux loss带来的提升越明显。4.3 batch size、BN与num_workers分割训练很吃显存。以520x520的输入、resnet50的backbone为例单卡batch size到8就需要大约16GB显存。如果你只有一张12GB的卡有两个选择一是把输入裁剪降到448x448牺牲一点精度二是保持batch size4同时使用SyncBN的等效替代——GroupNorm。这里要特别提一下BN的坑当batch size很小的时候BatchNorm统计的均值和方差波动非常大模型经常在验证集上忽高忽低。网上有个常见方案是把BN换成GroupNorm但GroupNorm的分组数需要针对通道数调整而且要修改所有backbone的BN层工程量大且推理速度略有下降。我的折中方案是尽量保证batch size≥8如果实在不够就先冻结backbone的BN参数BN的running统计用预训练值只训练后面新加层的BN。4.4 一个可直接改的train.py骨架核心训练循环我就不贴完整代码了只把最容易写错的部分单独拎出来每轮的迭代逻辑和loss回传。提示一个容易被忽略的点在计算预测结果时模型输出的是逻辑值需要先经过argmax转成类别索引再与目标做比较。而CrossEntropyLoss内部会自动做softmax切勿在损失函数前手动加softmax否则会双重softmax导致梯度消失。for images, masks in train_loader: images images.cuda() masks masks.long().cuda() outputs model(images) # 若开了aux loss此处是tuple loss criterion(outputs[0], masks) if aux_enabled: loss 0.4 * criterion(outputs[1], masks) optimizer.zero_grad() loss.backward() optimizer.step() # torch.cuda.synchronize() 可以真实统计单step耗时但别在训练循环里一直开着我还有一个自查习惯训练每跑完一个epoch手动保存一份当前模型在验证集某几张图上的可视化对比图人眼扫一眼比看几个数值指标更能发现问题。比如如果预测结果全部是背景说明损失函数或标签映射有严重bug这时候继续训练只会浪费时间。5. 训练过程中的监控与排坑记录我把训练中常见的几类问题集中放在这里每一类都是我自己踩过并且最终解决的记录的思路比结论更重要因为类似问题的排查套路是可以迁移的。5.1 mIoU怎么算最靠谱测试时经常有个误区直接拿整个验证集的预测结果做混淆矩阵统计。这种做法对大图不友好因为显存限制会强制把大图切块预测切块边界可能有重叠统计会混乱。我的做法是把验证集每张图拉伸到固定尺寸比如520x520预测得到全分辨率的概率图后先做argmax得到类别索引再与原图标签逐像素比对累计混淆矩阵最后按标准公式算mIoU。mIoU的计算公式是每个类别IoU的平均值其中IoU 该类别的交并比。对于一个两类问题如果裂缝类没预测出任何一个正样本该类的IoU就是0mIoU直接掉一半。所以验证的时候不要只看总mIoU一定要看每个类别的IoU尤其是小目标那个类别的单独IoU那才是模型真正有没有学会检测难类别的信号。5.2 显存不够的三种对策显存OOM在训练大尺寸输入时几乎是必经之路。我的处理优先级是先降batch size再用梯度累积模拟大batch最后才考虑降输入分辨率。梯度累积的写法是每跑几次小batch才更新一次参数比如原本batch size16现在batch size4每4步累积一次梯度效果上非常接近16的大batch训练。但注意BatchNorm仍然会按照实际batch size4计算所以如果积攒的步数比较多BN统计还是不准确需要配合冻结BN或使用SyncBN才能完全等价。另一个小的显存优化点是把一次性把整张图放大到520x520训练改成随机裁剪到448x448。实际上对于大多数物体分割448的分辨率已经能覆盖绝大部分特征520反而会让模型的感受野只覆盖局部。对裂缝这类细长目标我自己试过448训练、520推理效果反而更好。5.3 loss异常、全黑或全白输出的排查链路训练不一会儿loss变成nan十有八九是学习率太大尤其初始学习率直接沿用分类任务常用的0.1这类值时几乎必爆。从0.01、0.007往下压一般能解决。如果还没解决就要检查标签中是否出现了超出类别数的值比如忘了把255以外的标注清理干净。预测结果全白或全黑常见原因有两个。一个是标签和预测的类别ID错位背景是0、目标是1但模型输出的argmax也是0或1如果加载模型权重时忘记处理类别数就会整个偏移。另一个原因是模型训练完后被误用在了未归一化的输入上预训练模型默认输入为ImageNet归一化区间直接喂0-255的图片会得到离谱输出。排查方法就一条找一张验证集图片打印预测的类别分布直方图看是不是单类占绝对多数如果是检查这两点基本就能定位。6. 测试推理与模型落地训练结束后模型要从“能跑”变成“好用”中间还隔着推理和部署这一层。很多人在这一步才发现训练时没留后路导致验证指标很高但实际应用时惨不忍睹。6.1 推理时的两件小事第一测试阶段务必开启model.eval()模式这不仅仅是切换BN层用running统计量还同时关闭了dropout等训练行为。很多人忘了这一步导致每次预测结果都不一样还以为模型没收敛。第二预测原图尺寸输出时不要在模型内部做上采样到原图的操作建议在推理脚本里先缩放到训练尺寸预测完之后再用bilinear插值还原到原始尺寸。这样能避免模型在任意scale下插值带来的不可控误差。大图推理时如果单张图分辨率超过15兆像素直接整张送入模型会导致显存爆炸。我的方案是滑窗预测窗口大小取训练尺寸重叠区域设25%重叠部分在拼接时取预测概率的平均而不是简单用argmax。这样窗口边界的分割质量不会明显劣化。6.2 ONNX导出与转换训练好的模型要落地到C推理或移动端最好导出ONNX。这里有一个容易踩的坑PyTorch的bilinear上采样算子在ONNX里可能被导出为不同的算子集合导致在不同推理框架下速度差异巨大。解决方法是把F.interpolate显式替换成nn.Upsample并且固定输出尺寸避免动态shape导致后续优化困难。导出的一个基本样例model.eval() dummy_input torch.randn(1, 3, 520, 520).cuda() torch.onnx.export( model, dummy_input, deeplabv3p.onnx, input_names[input], output_names[output], opset_version12, dynamic_axes{input: {0: batch, 2: height, 3: width}, output: {0: batch, 2: height, 3: width}} )如果你只用了固定的520x520推理建议把dynamic_axes去掉这样转出来的模型在部分推理引擎里性能更好。动态shape看着灵活但很多优化器会对固定shape做算子融合这一步差出来的推理延迟能到2-3倍。6.3 模型轻量化思路如果只是自己实验resnet50的deeplabv3完全够用。但如果要跑在嵌入式设备上可以换用MobileNetV2作backbone官方提供对应预训练权重通道数少了一大半参数量降到三分之一左右mIoU通常只有3-5个点的损失。此外训练时如果用EMA指数移动平均维护一份权重副本推理用EMA权重在大多数分割数据集上能额外涨1-2个点几乎免费开销只是多一份模型权重在内存里推荐一试。这段时间把整套流程从零搭完我最大的感受是以前以为deeplabv3的难点在ASPP和decoder的结构真正动手才发现数据集处理、训练超参数、推理细节这些“外围”环节才是决定项目能不能落地的关键。如果你也在搭建自己的分割训练流程建议按这里面的顺序一步步走每个模块单独验证通过后再拼装。多花半天把每一行代码弄懂比盲目复制开源代码再用十倍时间去debug要划算得多。