PyTorch实战：DeepLabV3在Cityscapes上的语义分割全流程

简介这份资源面向计算机视觉方向的研究者与开发者提供在Cityscapes数据集上训练DeepLabV3语义分割模型的完整PyTorch实现适合具备一定深度学习基础、希望复现或改进分割模型的中高级学习者。压缩包共18个文件约258.23MB以12个Python脚本为核心涵盖模型定义、训练、评估与数据预处理流程另含4个pth预训练权重、1个md说明文档和1个license许可文件目录结构清晰便于按模块查阅与二次开发。资源围绕DeepLabV3的ASPP结构与全局上下文模块展开配套Cityscapes数据预处理、DataLoader构建、损失函数与优化器设置、mIoU评估及模型保存加载等环节可直接用于训练与验证。目前已有2094人学习下载能帮助读者快速搭建实验环境、理解模型内部原理并在此基础上针对特定场景优化分割效果。1. 从一张街景图说起DeepLabV3 在 Cityscapes 上到底解决了什么问题手里有一批行车记录仪或街拍图像想把里面的人、车、道路、天空、建筑逐像素分开这是语义分割最典型的落地场景。Cityscapes 数据集就是为这个场景准备的5000 张精细标注的城市街景图覆盖 30 类目标其中 19 类用于训练评估。而 DeepLabV3 是这套任务里被反复验证过的骨架之一它用空洞卷积加多尺度 ASPP 模块在保持感受野的同时不丢分辨率。PyTorch 实现意味着你能改结构、换 backbone、调 loss而不是被某个黑盒工具锁死。这篇笔记面向的是想自己跑通训练、拿到可复现指标、并且知道每一步在干什么的工程师。不管你是刚配好 pytorch 环境搭建的新手还是已经用过 pytorch实战 做过检测的老手下面这套流程都能直接抄。2. 先把数据管线搭对Cityscapes 的目录结构与标签映射2.1 为什么 Cityscapes 的标注不能直接读Cityscapes 原始标注是 JSON 格式的 polygon 多边形不是常见的 PNG mask。每张图的标注文件在gtFine/下命名规则是{city}_{seq}_{frame}_gtFine_polygons.json。里面记录了每个多边形对应的类别和顶点坐标。如果直接拿这个 JSON 去训练你需要在 Dataset 的__getitem__里实时做多边形填充速度极慢一个 epoch 可能要多花三到五倍时间。常见做法是先离线把 JSON 转成单通道 PNG 标签图训练时只读图。转换脚本的核心逻辑是用 PIL 的 ImageDraw 把每个 polygon 按类别 ID 填充到一张全零的 uint8 图上。import json import numpy as np from PIL import Image, ImageDraw import os # Cityscapes 19 类训练 ID 映射-1 表示忽略 CLASS_MAP { road: 0, sidewalk: 1, building: 2, wall: 3, fence: 4, pole: 5, traffic light: 6, traffic sign: 7, vegetation: 8, terrain: 9, sky: 10, person: 11, rider: 12, car: 13, truck: 14, bus: 15, train: 16, motorcycle: 17, bicycle: 18 } def json_to_label_png(json_path, out_path): with open(json_path, r) as f: data json.load(f) h, w data[imgHeight], data[imgWidth] label np.zeros((h, w), dtypenp.uint8) img Image.fromarray(label) draw ImageDraw.Draw(img) for obj in data[objects]: label_name obj[label] if label_name in CLASS_MAP: cls_id CLASS_MAP[label_name] polygon obj[polygon] # polygon 是 [[x1,y1],[x2,y2],...] 格式 flat [tuple(p) for p in polygon] draw.polygon(flat, fillcls_id) img.save(out_path)这段代码里CLASS_MAP只保留 19 类其余类别如 ego vehicle、rectification border直接丢弃对应像素保持 0 但后续要在 loss 里忽略。draw.polygon的填充顺序会影响重叠区域的最终类别Cityscapes 官方建议按 polygon 在 JSON 中的顺序绘制后画的覆盖先画的。转换完成后你的目录应该长这样leftImg8bit/train/aachen/aachen_000000_000019_leftImg8bit.png对应gtFine/train/aachen/aachen_000000_000019_gtFine_labelIds.png。注意labelIds和labelTrainIds的区别前者是原始 34 类 ID后者才是 19 类训练 ID用错了会导致类别错位。2.2 Dataset 与 DataLoader 的四个关键参数写 Dataset 类时除了基本的__init__、__len__、__getitem__有几个参数直接决定训练稳定性和显存占用。第一是crop_sizeCityscapes 原图 1024x2048直接整图训练在 11GB 显存上 batch size 只能设 1常见做法是随机裁剪 512x1024 或 769x769。第二是scale_rangeDeepLabV3 原论文用了 0.5 到 2.0 的多尺度增强但实际训练时固定尺度加随机翻转往往更稳。第三是ignore_index必须设为 255 或你在转换时约定的忽略值否则边界像素会污染 loss。第四是num_workers在 Linux 下设 4 到 8Windows 下如果遇到卡死就设 0 并检查if __name__ __main__保护。from torch.utils.data import Dataset, DataLoader import torchvision.transforms as T from PIL import Image import os class CityscapesDataset(Dataset): def __init__(self, root, splittrain, crop_size(512, 1024)): self.img_dir os.path.join(root, leftImg8bit, split) self.label_dir os.path.join(root, gtFine, split) self.crop_size crop_size self.images [] for city in os.listdir(self.img_dir): city_dir os.path.join(self.img_dir, city) for f in os.listdir(city_dir): if f.endswith(_leftImg8bit.png): self.images.append(os.path.join(city_dir, f)) def __len__(self): return len(self.images) def __getitem__(self, idx): img_path self.images[idx] label_path img_path.replace(leftImg8bit, gtFine).replace( _leftImg8bit.png, _gtFine_labelTrainIds.png) img Image.open(img_path).convert(RGB) label Image.open(label_path) # 同步随机裁剪 i, j, h, w T.RandomCrop.get_params(img, self.crop_size) img T.functional.crop(img, i, j, h, w) label T.functional.crop(label, i, j, h, w) img T.ToTensor()(img) img T.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225])(img) label torch.from_numpy(np.array(label)).long() return img, labelRandomCrop.get_params保证图像和标签用同一组裁剪参数这是分割任务里最容易翻车的地方之一。归一化用的 ImageNet 均值方差因为 backbone 通常加载 ImageNet 预训练权重。labelTrainIds文件里忽略区域的值是 255在 CrossEntropyLoss 里设ignore_index255即可。DataLoader 的shuffleTrue训练时开启验证时关闭pin_memoryTrue在 GPU 训练时能减少数据搬运开销。3. DeepLabV3 的 PyTorch 实现ASPP 模块与 backbone 选型3.1 ASPP 到底在做什么为什么 dilation rate 是 6、12、18DeepLabV3 的核心是 Atrous Spatial Pyramid Pooling它在 feature map 上并行跑四个分支一个 1x1 卷积、三个 3x3 空洞卷积dilation rate 分别为 6、12、18、一个全局平均池化。这样做的目的是在同一层捕获多尺度上下文——小 dilation 关注局部细节大 dilation 覆盖远处物体。为什么是 6、12、18这是原论文在 Cityscapes 上通过网格搜索得到的经验值对应输出步长 16 时的感受野大约覆盖 30、60、90 像素。如果你把 output stride 改成 8dilation rate 通常要翻倍到 12、24、36否则感受野不够。ASPP 里每个分支后面都接 BatchNorm 和 ReLU最后把五个分支 concat 起来再过一个 1x1 卷积降维到 num_classes。import torch import torch.nn as nn import torch.nn.functional as F class ASPP(nn.Module): def __init__(self, in_channels, out_channels256, rates[6, 12, 18]): super().__init__() self.branches nn.ModuleList() # 1x1 卷积分支 self.branches.append(nn.Sequential( nn.Conv2d(in_channels, out_channels, 1, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue) )) # 三个空洞卷积分支 for r in rates: self.branches.append(nn.Sequential( nn.Conv2d(in_channels, out_channels, 3, paddingr, dilationr, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue) )) # 全局池化分支 self.global_pool nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Conv2d(in_channels, out_channels, 1, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue) ) self.project nn.Sequential( nn.Conv2d(out_channels * 5, out_channels, 1, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue), nn.Dropout(0.5) ) def forward(self, x): size x.shape[-2:] feats [branch(x) for branch in self.branches] gp self.global_pool(x) gp F.interpolate(gp, sizesize, modebilinear, align_cornersFalse) feats.append(gp) out torch.cat(feats, dim1) return self.project(out)paddingr和dilationr必须相等这样卷积核在空洞展开后输出尺寸不变。全局池化分支最后要插值回原 feature map 大小再 concatalign_cornersFalse是 PyTorch 分割任务里的推荐设置。Dropout 0.5 只在训练时生效推理时自动关闭。如果你显存吃紧可以把out_channels从 256 降到 128但 mIoU 通常会掉 1 到 2 个点。3.2 backbone 用 ResNet-101 还是 MobileNetV2选 backbone 本质是在精度和速度之间做权衡。ResNet-101 是 DeepLabV3 原论文的标配在 Cityscapes val 上能到 78% 左右的 mIoU但参数量大训练时 batch size 只能设 4 到 8多卡。MobileNetV2 轻量单卡 11GB 能跑 batch size 16mIoU 大概 70% 到 72%适合快速验证想法或部署到边缘设备。我一般会先用 MobileNetV2 跑通全流程确认数据管线和 loss 没问题再换 ResNet-101 冲指标。换 backbone 时注意两点一是输出步长ResNet 的layer4输出 stride 是 16MobileNetV2 也是 16但后者最后几个 block 的通道数不同ASPP 的in_channels要对应改二是预训练权重torchvision.models里都有加载时用strictFalse跳过分类头。import torchvision.models as models def build_deeplabv3(backboneresnet101, num_classes19): if backbone resnet101: net models.resnet101(pretrainedTrue) in_channels 2048 # 去掉最后两个下采样保持 output stride16 net.layer4[0].conv2.stride (1, 1) net.layer4[0].downsample[0].stride (1, 1) elif backbone mobilenetv2: net models.mobilenet_v2(pretrainedTrue).features in_channels 320 aspp ASPP(in_channels, 256) classifier nn.Conv2d(256, num_classes, 1) return net, aspp, classifierResNet 改 stride 那两行是关键不改的话 output stride 是 32小物体分割会明显变差。MobileNetV2 直接取features部分输出通道 320。实际训练时我会把 backbone、ASPP、classifier 包在一个nn.Module里forward 时先过 backbone 拿 feature再过 ASPP最后分类头输出 logits再双线性插值回原图大小算 loss。4. 训练配置与调参学习率、loss 和 batch size 怎么定4.1 学习率策略poly 衰减比 step 更稳DeepLabV3 原论文用的是 poly 衰减lr base_lr * (1 - iter / max_iter) ** 0.9。相比每 30 个 epoch 降一次的 step 策略poly 衰减在分割任务里更平滑不容易在后期震荡。base_lr 一般设 0.01batch size 16 时如果 batch size 减半lr 也减半。优化器用 SGD momentum 0.9weight decay 1e-4。注意 backbone 和 ASPP 可以用不同的学习率backbone 设小一点比如 base_lr 的 0.1 倍ASPP 和分类头用完整 lr这样预训练权重不会被破坏太快。def poly_lr(base_lr, iter, max_iter, power0.9): return base_lr * (1 - iter / max_iter) ** power optimizer torch.optim.SGD([ {params: backbone.parameters(), lr: 0.001}, {params: aspp.parameters(), lr: 0.01}, {params: classifier.parameters(), lr: 0.01} ], momentum0.9, weight_decay1e-4)训练循环里每个 iteration 手动调param_group[lr]不要用lr_scheduler的 step 模式因为 poly 是按 iteration 算的。max_iter 通常设 30000 到 60000Cityscapes 精细标注 2975 张训练图batch size 8 时大约 370 个 iteration 一个 epoch30000 iter 约 80 个 epoch。4.2 loss 用 CrossEntropy 还是加辅助 loss标准做法是 CrossEntropyLoss 加ignore_index255再配合class_weight处理类别不平衡。Cityscapes 里 road、building、sky 像素多pole、traffic light 像素少不加权重的话小类别 mIoU 会很低。常见权重是按类别频率的倒数开根号或者直接用torch.nn.CrossEntropyLoss(weighttorch.tensor([...]))。另一个技巧是加辅助 loss在 backbone 中间层接一个小的分类头和主 loss 加权求和权重 0.4。这能缓解梯度消失尤其当 backbone 很深时。但辅助 loss 不是必须的我跑 MobileNetV2 时通常不加ResNet-101 会加。criterion nn.CrossEntropyLoss(ignore_index255, weighttorch.tensor(class_weights)) # 辅助 loss 示例 aux_logits aux_classifier(aux_feat) aux_logits F.interpolate(aux_logits, sizelabel.shape[-2:], modebilinear, align_cornersFalse) loss criterion(logits, label) 0.4 * criterion(aux_logits, label)class_weights需要你统计训练集里每个类别的像素占比然后取倒数归一化。注意权重不要设得太极端否则模型会偏向小类别而牺牲整体 mIoU。验证时用argmax拿预测类别再算混淆矩阵最后按类别算 IoU 取平均。5. 避坑与排查训练不收敛、mIoU 异常、显存爆炸的五个血泪经验5.1 现象loss 从第一轮就 NaN原因通常是学习率太大或者数据里有非法值。先检查输入图像归一化后有没有 NaN再确认 label 里有没有超出num_classes的 ID。如果用了labelIds而不是labelTrainIds里面会有 34 类的 ID而你的分类头只有 19 类CrossEntropy 会直接报错或产生 NaN。解决方法是统一用labelTrainIds或者在 Dataset 里做一次映射。5.2 现象训练 loss 下降但验证 mIoU 一直卡在 10% 左右这大概率是标签和图像没对齐。随机裁剪时如果只裁了图像没裁标签或者用了不同的随机种子标签就会错位。另一个可能是归一化用了错误的均值方差导致模型学不到东西。排查方法取一个 batch把图像和标签可视化出来肉眼确认道路区域对应的标签值是不是 0。如果标签图看起来是乱的回去检查RandomCrop.get_params是否同时作用于 img 和 label。5.3 现象显存爆炸batch size 只能设 1Cityscapes 原图太大必须裁剪。如果你已经裁剪了 512x1024 还是爆检查 ASPP 的out_channels是不是设了 512 或更大以及有没有在 forward 里保留中间 feature 不释放。另一个常见原因是num_workers设太大导致内存不足而不是显存。用torch.cuda.empty_cache()清理缓存或者把crop_size降到 256x512 先跑通。5.4 现象mIoU 比论文低 5 个点以上先确认评估时有没有把 ignore 区域算进去。正确做法是只在非 ignore 像素上算混淆矩阵。其次检查 output stride 和 dilation rate 是否匹配改了 stride 没改 dilation 会严重掉点。最后看学习率策略poly 衰减的max_iter如果设得太小模型还没收敛就结束了。我一般会跑 60000 iter 再评估30000 iter 时 mIoU 通常还在上升。5.5 现象多卡训练时 mIoU 反而比单卡低这是 BatchNorm 的锅。多卡时如果用了DataParallel每张卡上的 batch size 变小BN 统计量不准。换成DistributedDataParallel并同步 BN或者直接把 BN 换成 GroupNorm。另一个可能是学习率没随卡数线性放大8 卡时 base_lr 应该设 0.08 而不是 0.01。6. 进阶技巧用 ONNX 导出和 TensorRT 加速推理训练完拿到 78% mIoU 的模型只是第一步实际部署时推理速度往往比精度更关键。PyTorch 模型直接跑一张 1024x2048 的图在 V100 上大概 200ms用 TensorRT 能压到 50ms 以内。中间桥梁是 ONNX。导出时注意三个点一是固定输入尺寸动态 shape 会让 TensorRT 优化受限二是把interpolate的align_corners显式写死ONNX 对默认值处理不一致三是导出后先用onnxruntime验证输出和 PyTorch 一致再转 TensorRT。import torch.onnx model.eval() dummy torch.randn(1, 3, 512, 1024).cuda() torch.onnx.export( model, dummy, deeplabv3_cityscapes.onnx, input_names[input], output_names[output], opset_version11, dynamic_axesNone # 固定尺寸 )导出后检查 ONNX 模型里有没有Resize节点如果align_corners没设对推理结果会和 PyTorch 有细微偏差。TensorRT 转换用trtexec命令行工具--fp16开启半精度--workspace2048给足显存。转换完用polygraphy对比 TensorRT 和 ONNX 的输出确保 max diff 小于 1e-3。如果部署在边缘设备上还可以把 backbone 换成 MobileNetV2 再走一遍这个流程速度能再快一倍。我自己踩过的坑是导出时忘了model.eval()BN 层还在用 batch 统计量导致推理结果完全不对。这个后悔药就是导出前一定加eval()和torch.no_grad()。希望帮到你。本文还有配套的精品资源点击获取