自己动手训一个语义分割模型最绕不开的坎就是数据集和结构实现之间的那道缝。市面上的教程不少但多数只给你一段能跑的代码换到自己的数据上立刻露馅标签不对、尺寸对不上、显存爆了、迭代几轮后发现loss是负的。这篇记录是我个人从零开始用PyTorch搭DeepLabv3并完整训练自定义数据集的全程复盘覆盖数据准备、网络搭建、训练配置和常见坑排查每个环节都给出可复现的细节。照着走一遍你应该能少走掉我当初踩的大半弯路。1. 整体设计与思路拆解1.1 为什么用DeepLabv3作为分割模型基线做语义分割模型的选型会直接决定你在数据处理、训练时间、调优策略上投入多少精力。早期尝试过FCN和U-Net的一个直接感受是FCN的上下文信息不足对多尺度目标容易出现破碎预测U-Net在小数据集上表现不错但一旦图像语义复杂度上升它的编码器主干显得不够强。相比之下DeepLabv3在编码器-解码器框架上做了很好的折中编码器用空洞卷积扩大感受野配合ASPP模块并行采样多尺度特征解码器把浅层空间信息与深层语义信息融合修整分割边界。这个结构对大目标和小目标都能照顾到而且训练难度相对可控。当时也考虑过PSPNet金字塔池化模块对场景级分割很有效但要处理大量自定义类别的细节纹理时ASPP的扩张卷积采样比池化更灵活。DeepLabv3另一个优势是PyTorch社区里先例很多遇到问题容易找到参考自己从零写也更有把握。如果你手里的是一个几百张到几千张的中小型自定义数据集DeepLabv3是一个投入产出比非常高的选择。1.2 从零搭建的总体路线整个流程我拆成四件事数据集准备与增强、骨干网络与ASPP/Decoder实现、训练策略与评估脚本、推理与调优。顺序不能乱因为每一步的坑会直接影响下一步。先准备好干净有效的标签再搭建网络然后用合适的损失与学习率调度去训练最后用mIoU指标衡量效果并回炉调整。工具链方面PyTorch版本用的1.9以上CUDA环境自行匹配单卡训练就够了。图像处理用OpenCV和Pillow搭配数据加载用PyTorch原生Dataset和DataLoader。如果你之前只做过分类任务这里需要建立起一个转变分割模型的输出不是单标签而是一个与输入尺寸相关的类别索引图所以后续所有数据操作都得围绕“像素级对应关系”展开。2. 数据集准备从原始图像到训练加载器2.1 标注与目录组织我这次用的是一套自采集的道路场景数据内容包含五个类别背景、路面、车辆、行人和标志物。标注工具用的是某开源图像标注软件导出时选择VOC格式的PNG索引图每个像素值直接对应类别编号。注意如果你导出的是带调色板的彩色PNG训练前必须转成灰度索引图读取不然模型会把RGB三通道当作三个独立分割目标。目录结构我做得和PASCAL VOC保持一致这样后面很多现成脚本可以直接复用dataset/ ├── JPEGImages/ ├── SegmentationClass/ └── ImageSets/ └── Segmentation/ ├── train.txt └── val.txtJPEGImages放原图SegmentationClass放标注图train.txt和val.txt分别记录不带扩展名的文件名。文件划分上我用了一个随机种子做8:2切分保证每次运行实验的可重复性。有个细节原图和标注图的文件名必须严格一致否则加载时配对会失败我一开始就是吃了命名不统一的亏。2.2 Dataset类的关键实现数据加载这一步核心是保证原图和标签图在经历同样的预处理后仍然对齐坐标。这里不能简单地把标签图当成普通彩色图片来做归一化而是要先读取成像素索引格式再在之后转成long tensor作为交叉熵损失的目标。一个相对稳妥的写法是import cv2 import torch from torch.utils.data import Dataset from PIL import Image import numpy as np class SegDataset(Dataset): def __init__(self, file_list, root, crop_size513, modetrain): self.file_list [line.strip() for line in open(file_list)] self.root root self.crop_size crop_size self.mode mode def __len__(self): return len(self.file_list) def __getitem__(self, idx): name self.file_list[idx] img cv2.imread(f{self.root}/JPEGImages/{name}.jpg) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) mask cv2.imread(f{self.root}/SegmentationClass/{name}.png, cv2.IMREAD_GRAYSCALE) if self.mode train: img, mask self.random_crop(img, mask) if np.random.rand() 0.5: img img[:, ::-1].copy() mask mask[:, ::-1].copy() else: img cv2.resize(img, (self.crop_size, self.crop_size)) mask cv2.resize(mask, (self.crop_size, self.crop_size), interpolationcv2.INTER_NEAREST) img img.astype(np.float32) / 255.0 img (img - np.array([0.485, 0.456, 0.406])) / np.array([0.229, 0.224, 0.225]) img torch.from_numpy(img.transpose(2, 0, 1)).float() mask torch.from_numpy(mask.astype(np.int64)).long() return img, mask def random_crop(self, img, mask): h, w img.shape[:2] ch min(self.crop_size, h) cw min(self.crop_size, w) top np.random.randint(0, h - ch 1) if h ch else 0 left np.random.randint(0, w - cw 1) if w cw else 0 img img[top:topch, left:leftcw] mask mask[top:topch, left:leftcw] if h self.crop_size or w self.crop_size: img cv2.resize(img, (self.crop_size, self.crop_size)) mask cv2.resize(mask, (self.crop_size, self.crop_size), interpolationcv2.INTER_NEAREST) return img, mask这里最容易被忽略的一点是mask的resize插值方式。任何时候缩放标签图都必须用INTER_NEAREST绝对不能用线性插值或立方插值否则边界会产生不属于任何类别的中间值。随机裁剪我设置了如果图像比裁剪尺寸小就先resize放大的兜底逻辑防止因为个别小尺寸图片导致程序中断。2.3 数据增强策略怎么选才不翻车数据增强的目标是让模型对光照、视角、遮挡更鲁棒但分割任务里增强方式必须谨慎。水平翻转、随机缩放、随机旋转小角度、色彩抖动都是常用手段。我实际测试下来随机缩放配合固定crop_size效果最好因为多尺度训练能帮助模型适应不同大小的目标。还有一个很容易踩的坑颜色归一化的mean和std必须使用骨干网络预训练时的统计值。如果用自己的数据重新计算mean和std网络加载预训练权重后会出现严重的分布偏移收敛速度会慢一大截。ResNet系列的统计值是[0.485, 0.456, 0.406]和[0.229, 0.224, 0.225]直接沿用即可。3. 模型搭建骨架、ASPP与Decoder的PyTorch实现3.1 骨干网络与空洞卷积的配合DeepLabv3的骨干我用的是ResNet101。这里的关键在于把标准ResNet的最后两个stage即stage3和stage4的stride改成1然后引入空洞卷积来补偿下采样带来的分辨率损失而不是直接把最后的全局平均池化接上输出。核心逻辑是模型输出特征图尺寸为输入尺寸的1/8然后经过ASPP采样和Decoder恢复空间细节。具体来说ResNet101的输出步长默认是32做语义分割时通常要降到8或16。常见的做法是将stage3的stride从2改成1同时用dilation2stage4的stride从2改成1dilation4。这个过程在torchvision提供的现成ResNet101中需要手动调整replace_stride_with_dilation参数或者在自定义ResNet layer时传入不同的dilation配置。import torchvision.models as models时可以直接用models.resnet101(pretrainedTrue, replace_stride_with_dilation[False, True, True])这样stage3和stage4的stride会被替换掉。这里需要留意的是预训练权重是基于ImageNet的图像分类训练的它的stage层结构和空洞卷积并不冲突因为空洞卷积只改变了采样间隔不改变卷积核参数个数。所以加载官方预训练参数是安全的文件大小大约170MB如果下载速度不理想可以先离线下载后再加载。3.2 ASPP模块的完整代码与设计思路ASPPAtrous Spatial Pyramid Pooling的作用是在同一个特征图上并行用不同采样率的空洞卷积捕捉不同尺度的上下文。编码器输出的特征图经过这几个分支后拼接在一起最后用一个1x1卷积压缩通道数。各分支的设计如下import torch import torch.nn as nn import torch.nn.functional as F class ASPP(nn.Module): def __init__(self, in_channels2048, out_channels256): super(ASPP, self).__init__() self.conv1 nn.Conv2d(in_channels, out_channels, 1, biasFalse) self.bn1 nn.BatchNorm2d(out_channels) self.conv2 nn.Conv2d(in_channels, out_channels, 3, dilation6, padding6, biasFalse) self.bn2 nn.BatchNorm2d(out_channels) self.conv3 nn.Conv2d(in_channels, out_channels, 3, dilation12, padding12, biasFalse) self.bn3 nn.BatchNorm2d(out_channels) self.conv4 nn.Conv2d(in_channels, out_channels, 3, dilation18, padding18, biasFalse) self.bn4 nn.BatchNorm2d(out_channels) self.global_avg nn.AdaptiveAvgPool2d(1) self.conv5 nn.Conv2d(in_channels, out_channels, 1, biasFalse) self.bn5 nn.BatchNorm2d(out_channels) self.project nn.Sequential( nn.Conv2d(out_channels * 5, out_channels, 1, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue) ) def forward(self, x): b, c, h, w x.size() x1 F.relu(self.bn1(self.conv1(x))) x2 F.relu(self.bn2(self.conv2(x))) x3 F.relu(self.bn3(self.conv3(x))) x4 F.relu(self.bn4(self.conv4(x))) x5 self.global_avg(x) x5 F.relu(self.bn5(self.conv5(x5))) x5 F.interpolate(x5, size(h, w), modebilinear, align_cornersTrue) x torch.cat([x1, x2, x3, x4, x5], dim1) return self.project(x)dilation取值6、12、18是我测试后觉得比较稳定的组合。如果输入图像分辨率比较大dilation可以适当上调例如8、16、24这取决于目标尺度分布。global_avg分支相当于一个全局上下文池化它能把整幅图像的语义信息压缩成一小块再上采样回原尺寸与局部分支融合。首次写这块时容易漏掉padding参数结果就是输出特征图尺寸比输入小拼起来维度对不上。3.3 Decoder让边界回来编码器的输出分辨率是输入的1/8如果直接上采样回原尺寸边界会非常糊。DeepLabv3的做法是引入一个轻量级Decoder取骨干网络stage2输出的低层特征分辨率是输入的1/4经过1x1卷积压缩到48通道再与编码器输出上采样4倍后的特征在通道维度上拼接随后经过3x3卷积和两次上采样恢复分辨率。这里为什么低层特征要压到48通道而不是保留原始256或512通道原因是低层特征虽然含有丰富的空间细节但语义信息弱通道数过多会让模型训练时低层信息占据主导同时增加显存消耗。48通道是一个折中值原论文实验也验证了在48附近效果最好。class Decoder(nn.Module): def __init__(self, low_level_channels256, num_classes5): super(Decoder, self).__init__() self.low_level_conv nn.Sequential( nn.Conv2d(low_level_channels, 48, 1, biasFalse), nn.BatchNorm2d(48), nn.ReLU(inplaceTrue) ) self.conv3x3 nn.Sequential( nn.Conv2d(256 48, 256, 3, padding1, biasFalse), nn.BatchNorm2d(256), nn.ReLU(inplaceTrue), nn.Conv2d(256, 256, 3, padding1, biasFalse), nn.BatchNorm2d(256), nn.ReLU(inplaceTrue) ) self.classifier nn.Conv2d(256, num_classes, 1) def forward(self, encoder_feat, low_level_feat): low_level_feat self.low_level_conv(low_level_feat) encoder_feat F.interpolate(encoder_feat, sizelow_level_feat.size()[2:], modebilinear, align_cornersTrue) x torch.cat([encoder_feat, low_level_feat], dim1) x self.conv3x3(x) x self.classifier(x) return F.interpolate(x, scale_factor4, modebilinear, align_cornersTrue)有个细节值得强调Decoder中的上采样建议使用align_cornersTrue这能保证与标注图的对齐方式一致虽然从数值上不会产生巨大差异但在严格评估mIoU时会避免少数像素偏移导致的精度波动。4. 训练配置与优化策略4.1 训练主循环的搭建细节训练脚本的整体结构其实不复杂但每个环节的细节都需要对齐。我采用的是VOC标准训练方式一个epoch内对每张图随机裁剪与增强然后批量送入模型。下面是我用来稳定收敛的训练循环骨架dataloader DataLoader(train_set, batch_size8, shuffleTrue, num_workers4, drop_lastTrue) model DeepLabV3Plus(num_classes5, backboneresnet101) model.cuda() model.train() criterion nn.CrossEntropyLoss(ignore_index255) optimizer torch.optim.SGD(model.parameters(), lr0.007, momentum0.9, weight_decay0.0001) total_iters len(dataloader) * epochs current_iter 0 for epoch in range(epochs): for batch_idx, (images, labels) in enumerate(dataloader): images, labels images.cuda(), labels.cuda() outputs model(images) loss criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() lr base_lr * (1 - current_iter / total_iters) ** 0.9 for param_group in optimizer.param_groups: param_group[lr] lr current_iter 1batch_size这里写的是8如果你的是16或32请按比例调低初始学习率否则容易在训练前几个step就发散。我在单卡16GB显存下crop_size513batch_size8ResNet101骨干刚好能跑。4.2 损失函数与类别不均衡的处理分割任务的损失函数首选还是交叉熵。但自定义数据集里类别分布经常非常不均衡。在我这份道路数据里背景像素数量远超车辆和行人如果不做任何处理模型会把所有区域都预测为背景最终mIoU却显得不算太差这是因为背景占了大多数像素。解决思路有两层。第一层是ignore_index只在标签里把不确定区域设为255并忽略它们。第二层是给每个类别的交叉熵加权权重可以按类别像素占比的倒数来算也可以用中值频率平衡。我这里使用的方案是先统计每个类别的像素比例然后构造权重向量import torch.nn.functional as F def compute_class_weights(mask_paths, num_classes): counts [1] * num_classes for p in mask_paths: m cv2.imread(p, cv2.IMREAD_GRAYSCALE) for c in range(num_classes): counts[c] int((m c).sum()) total sum(counts) weights [total / (num_classes * c) for c in counts] return torch.tensor(weights, dtypetorch.float32)这类权重能让少数类别在损失中的贡献变大但前提是标注数据足够可靠。如果某些类别标签噪声很大过高的权重反而会放大错误监督信号这时需要结合人工抽检结果来平衡。4.3 mIoU评估脚本不能只看loss曲线训练过程中loss下降未必代表分割质量好尤其是类别不平衡严重时。最好每个epoch结束时在验证集上算一次mIoU。IoU的定义是某一类别的交集和并集之比IoU TP / (TP FP FN)mIoU是所有类别IoU的算术平均。具体实现时我维护一个混淆矩阵每批次累加预测结果和真实标签的统计关系def compute_miou(confusion_matrix): iou_list [] for i in range(confusion_matrix.shape[0]): tp confusion_matrix[i, i] fp confusion_matrix[i, :].sum() - tp fn confusion_matrix[:, i].sum() - tp union tp fp fn if union 0: iou_list.append(float(nan)) else: iou_list.append(tp / union) return np.mean([x if not np.isnan(x) else 0 for x in iou_list]) def update_confusion_matrix(pred, label, num_classes): mask (label 0) (label num_classes) hist np.zeros((num_classes, num_classes)) p pred[mask] t label[mask] hist np.bincount(t * num_classes p, minlengthnum_classes ** 2).reshape(num_classes, num_classes) return hist这里很容易犯的错是把预测结果直接传成概率图必须用torch.argmax(output, dim1)先把logits转成类别索引再移到CPU转NumPy计算。另一个细节是评估时的图像尺寸验证集一般也需要先缩放到固定尺寸再送入模型不能直接使用原图尺寸否则batch内形状不一致。4.4 学习率调度和训练轮数经验值我用的学习率调度是poly策略即lr base_lr * (1 - iter / total_iter) ** 0.9比固定学习率或StepLR收敛得更平滑因为分割模型一般需要尽量充分地探索后期的小学习率区间。base_lr这个值的选取和batch_size直接相关原论文在batch16的情况下用0.007但那是基于多卡和较深骨干。单卡batch8时base_lr设成0.001到0.005之间比较稳我最终用0.003。训练轮数常规设置是60到80个epoch左右。如果你的数据集规模只有几百张更早的监控频率会更有用例如每三到五个epoch验证一次即可。太频繁的验证会打断训练节奏而且验证集样本变化带来的波动会误导调参判断。5. 常见问题与排查技巧实录5.1 训练后全黑或全同一色输出这个情况出现在我第一次跑完20个epoch去可视化预测图的时候模型几乎把整幅图都预测成背景。排查思路从几个方向入手先看单batch的loss是否正常下降再看验证集的mIoU曲线。如果loss正常但mIoU几乎为零多半是标签映射出错了比如标注图的像素值范围是0-4但normalize时不小心除以了255导致所有标签都变成了0或1。还有可能是背景类别在数据集中占比太高模型陷入局部最优。解决办法是给背景类降权或者对背景类做随机下采样让训练时每次batch中出现的前景像素比例相对提高。另外模型最后一层默认输出形状是[N, num_classes, H, W]如果num_classes和实际类别数不一致可视化时自然对不上。5.2 loss变成NaN的典型原因NaN大多出现在网络刚初始化以后也可能出现在训练中途。一种常见触发点是学习率过大尤其是SGD配合poly策略在初期迭代还没衰减时如果设置的base_lr超过0.01很容易爆炸。另一种可能是数据里面有异常标注图比如某个文件是灰度模式但像素值是负值或者读入时channel顺序不对导致输入包含异常值。我的排查办法是训练脚本里加一个小开关每隔一定迭代数打印输入图像的mean/std和当前loss值。如果输入图片出现全零或全255的数据直接跳过该样本。另外如果使用混合精度训练也需要检查梯度缩放器是否正常工作不过我这次用的是单精度训练稳定性和可复现性都更好些。5.3 如何判断是欠拟合还是数据噪声大训练到中后期经常出现验证集mIoU波动很大停在某个数值上不去。这时候我会先减少数据增强强度只保留翻转和随机缩放看看mIoU是否反而上升。如果是说明增强带来的样本分布扰动过大模型还没学到足够稳健的特征。再一个常用手段是把某些置信度极低的样本挑出来在训练结束后用模型预测一遍训练集把预测错误且IoU低的单张图拿出来看标签和原图。很多时候会发现标签根本没对齐比如标注时把车辆和行人画进了同一个轮廓里这时单纯调参没有意义得回标注端修改。这类脏数据对分割模型的影响远大于分类模型。6. 推理可视化与后续扩展建议6.1 推理脚本要点训练完成后的推理阶段把模型设为eval模式关闭梯度计算输入图像同样做resize和归一化然后取argmax得到类别索引图。这里要注意的是推理时的resize必须和训练时保持一致的目标尺寸不能随意改变否则模型性能会下降。推理脚本里我还会做一步简单的后处理去除过大过小的独立区域使用连通域分析过滤掉面积小于某个阈值的噪声预测块。这个对道路场景尤其有效因为单个像素被误判为车辆或行人的碎片会被清掉。阈值选多少要看具体类别我这边设成20像素既能滤掉噪声又不影响真实小目标。6.2 模型导出与部署方向如果后续要部署到服务端或者移动端可以把训练好的权重转成ONNX。转换时需要注意PyTorch中F.interpolate的align_corners参数和输入动态尺寸ONNX导出时要固定输入分辨率或者设置动态轴。我实际试过用动态轴导出推理速度和易用性整体都不错但某些版本的ONNXRuntime对AdaptiveAvgPool2d的导出支持有限需要先调用一次model.eval()并用假数据跑一次前向来固化图结构。如果想要进一步提升分割效果可以尝试把骨干从ResNet101换成MobileNetV3-Large精度降低不多但推理速度会快很多适合边缘设备。另一个成本更低的方向是接入数据增强中的多尺度预测推理时对同一张图使用三个不同尺度输入把概率图平均后再做argmax通常能带来1到2个mIoU点的提升。动手实操中的几点真实体会反复调试这套流程之后我最大的感受是分割模型的代码量并不大真正的工程量在数据与训练细节。网上很多现成项目仓库虽然能跑通VOC但换到自定义数据上总会出现各种兼容问题。写这套实现的时候我尽量把网络结构写清楚了Dataset部分也保留了足够的可修改空间方便你替换成自己的标注格式。如果你准备动手先把整个流程跑通一次哪怕类别只有两三类也比一开始就挑战多类别要稳。确认训练曲线和mIoU正常后再逐步增加类别数量和数据增强强度。这套DeepLabv3的样板之后还能继续迁移到遥感分割、医疗影像等多个方向希望你也能在迁移中把参数调出自己的手感。
