ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

DeepLabv3+图像分割实战:Pytorch训练与Cityscapes数据集应用

DeepLabv3+图像分割实战:Pytorch训练与Cityscapes数据集应用 简介面向图像分割入门与进阶开发者这份实战资源基于Pytorch在VOC和Cityscapes数据集上完整实现DeepLabv3算法。包含从数据加载、网络搭建、模型训练到预测评估的全流程Python脚本以及分割结果可视化样例和详细流程教程适合希望系统掌握语义分割项目实践的读者。zip压缩包共55个文件以23个py源码、17个png结果图为主另有txt数据列表、zbak备份文件及README说明整体仅2.25MB结构清晰便于按模块学习。已有139人学习资源内附网络骨干代码含ResNet、Xception、MobileNetV2等、损失函数、调度器、可视化工具等并配有VOC与Cityscapes的标注文件与数据划分可直接用于复现和二次开发。压缩包还包含附赠内容便于扩展练习。1. DeepLabv3 结合 Pytorch 做图像分割训练我没有先读论文而是直接看 torchvision 源码很多人拿到 DeepLabv3 这种经典的图像分割模型第一反应是去重读论文、找仓库、跑 demo最后卡在数据格式上。我在实际项目里反而不这样起步torchvision 里已经封装了 DeepLabv3 的骨干结构Pytorch 生态对 VOC 和 Cityscapes 这两个数据集的加载也有现成接口真正需要自己动手的部分其实只有三块——把 DeepLabv3 的 Decoder 补全把 Cityscapes 的 34 类原始标注映射到 19 类训练标签以及写一个带 poly 学习率衰减和 ignore_index 掩码的训练循环。这篇文章就按这条链路把整条训练流程摊开讲。适合有 Pytorch 基础、想快速在分割任务上出结果而不是纠结框架细节的工程向读者也适合刚看完语义分割论文、准备写第一个训练脚本的研究生。2. DeepLabv3 的结构拆解与 VOC、Cityscapes 数据集的 Pytorch 加载2.1 DeepLabv3 比 DeepLabv3 多了什么Decoder 要自己补torchvision 的torchvision.models.segmentation.deeplabv3_resnet101实现的是 DeepLabv3 结构ResNet101 作为骨干在最后一层输出上接 ASPP空洞空间金字塔池化然后直接 1x1 卷积分类再双线性上采样 8 倍恢复原尺寸。这套设计的短板很明显——ASPP 输出的特征图是 1/8 分辨率output_stride8 时上采样回来会丢失很多边缘细节。DeepLabv3 论文的 Decoder 正是用来补这个短板的取骨干网络 layer1也就是 1/4 下采样的低层特征通过 1x1 卷积降到 48 通道再与双线性上采样 4 倍之后的 ASPP 输出拼接过两层 3x3 卷积后恢复类别数。所以我一般会在 torchvision 的 DeepLabv3 基础上手动补一个 Decoder这也是整个训练流程里最像“源码”的部分。核心实现如下import torch import torch.nn as nn import torch.nn.functional as F class DeepLabv3PlusDecoder(nn.Module): def __init__(self, num_classes, low_level_channels256, aspp_channels256): super().__init__() # low_level_channels256 对应 ResNet101 的 layer1 输出通道数 self.low_level_conv nn.Sequential( nn.Conv2d(low_level_channels, 48, 1, biasFalse), nn.BatchNorm2d(48), nn.ReLU(inplaceTrue) ) self.fuse nn.Sequential( nn.Conv2d(aspp_channels 48, 256, 3, padding1, biasFalse), nn.BatchNorm2d(256), nn.ReLU(inplaceTrue), nn.Dropout(0.5), nn.Conv2d(256, 256, 3, padding1, biasFalse), nn.BatchNorm2d(256), nn.ReLU(inplaceTrue), nn.Dropout(0.1), ) self.classifier nn.Conv2d(256, num_classes, 1) def forward(self, aspp_out, low_level_feat): low_level_feat self.low_level_conv(low_level_feat) h, w low_level_feat.shape[-2:] # ASPP 输出分辨率是低层特征的 1/2先上采样对齐再拼接 aspp_out F.interpolate( aspp_out, size(h, w), modebilinear, align_cornersTrue ) return self.classifier(self.fuse(torch.cat([aspp_out, low_level_feat], dim1)))这个 Decoder 的输入aspp_out来自 torchvision 的 DeepLabv3 模型去掉最后的 1x1 分类层原模型的model.classifier[4]之前的部分low_level_feat则来自model.backbone.layer1的输出。训练时把这两个输出接进 Decoder 再算损失就完成了一个严格的 DeepLabv3 结构。ResNet101 的 layer1 输出通道是 256这是low_level_channels256这个参数的直接来源如果换用 MobileNetV3 之类的轻量骨干这个数值要跟着改。Decode 层的两个 Dropout 也很关键。第一个 0.5 的 Dropout 作用于多尺度特征融合后的高层语义第二个 0.1 的 Dropout 在分类前做最后一次正则这个 0.1 的设置是从 Xception 版 DeepLabv3 的工程实现里沿袭过来的换数据集时不必调。2.2 VOCSegmentation 与 Cityscapes 的加载差异目录结构和类别映射VOC 和 Cityscapes 的加载方式差异很大根源在于标签格式不同。VOC 的 ground truth 是带调色板的 PNG 图片PIL 读出来是modeP像素值直接就是类别索引而 Cityscapes 的 gtFine 是单通道 L 模式的 PNG像素值对应的是 34 类原始标注 id需要映射到 19 个训练类剩下的类别统一置为 255 让损失函数忽略。VOC 用 torchvision 封装即可train 集 1464 张、val 集 1449 张适合快速验证训练流程能否跑通。Cityscapes 的 torchvision 接口要求数据目录结构必须是gtFine/train与leftImg8bit/train同级摆放否则自检逻辑会报找不到文件的错。Cityscapes 数据量比 VOC 大一个量级train 集 2975 张val 集 500 张单卡 resnet101 输入 769x769 大概要跑接近一天才能看到像样的 mIoU建议先用 VOC 调通代码再切 Cityscapes。Cityscapes 类映射是这段流程里最大的坑我把它单独拎出来import numpy as np import pandas as pd # Cityscapes 官方 id 到训练类 id 的映射-1 表示 void 或 ignore # 数据来自 cityscapesScripts 仓库的 labels.py这里只列常用类别 class_id_map { 0: 0, # road 1: 1, # sidewalk 2: 2, # building 3: 3, # wall 4: 4, # fence 5: 5, # pole 6: 6, # traffic light 7: 7, # traffic sign 8: 8, # vegetation 9: 9, # terrain 10: 10, # sky 11: 11, # person 12: 12, # rider 13: 13, # car 14: 14, # truck 15: 15, # bus 16: 16, # train 17: 17, # motorcycle 18: 18, # bicycle 255: 255, # unlabeled } def cityscapes_remap(label): # label 为 PIL Imagenp.array 后像素值是原始 id label np.array(label, dtypenp.int64) mapped np.full(label.shape, 255, dtypenp.int64) for orig_id, train_id in class_id_map.items(): # 只映射存在于表的类别其余原始 id 保持 255 mapped[label orig_id] train_id return mapped这里有个容易忽略的点Cityscapes 原始 34 类中像建筑外墙、护栏后的地面、人行道旁的植被这类ignoreInEval类别官方标注里用的是 1、4、9 这类单独 id经映射后应该变成 255 而非某个训练类。如果直接用np.where把未在映射表里的类别全归 255代码没错但不透明上面这种显式字典映射的方式跑完一版之后还能拿来做类别分布统计。VOC 则不需要这个步骤加载后像素值直接就是类别 id。增强策略上两个数据集通用的是随机左右翻转和随机缩放后裁剪。VOC 我习惯用 513x513 裁剪Cityscapes 用 769x769两者都配合 0.5 到 2.0 的随机缩放。Pytorch 里做这类组合增强torchvision.transforms.Compose配合RandomResizedCrop可读性更好但注意RandomResizedCrop的 scale 参数要限制在 0.5 到 2.0 之间并关闭随机长宽比变化否则分割标签会被拉伸变形。3. Pytorch 训练循环的搭建poly 学习率、ignore_index 与 AMP 精度混合3.1 损失函数的选择CrossEntropyLoss 配 ignore_index255 是分割训练的标准组合DeepLabv3 在 VOC 和 Cityscapes 上的损失函数我始终用nn.CrossEntropyLoss(ignore_index255)原因有三一是这个损失函数与 mIoU 评估指标的相关性在分割任务上被验证得最充分二是 Pytorch 的 CrossEntropyLoss 原生支持 mask 机制255 位置的梯度整片置零不污染梯度统计三是它不需要像 Dice Loss 那样关心类别的空间重叠模式在长尾小目标上的收敛行为更稳定。VOC 的 21 类里背景占了绝大部分像素Cityscapes 里 road、building 这类大物体占比极高单独用torch.optim.AdamW配合 CrossEntropyLoss 即可不需要额外加权。如果类别不均衡到影响评估结果我会在损失里对每个类别乘一个权重torch.tensor([...])权重用训练集像素占比的倒数近似。VOC 的 person 和 Cityscapes 的 traffic light 这类目标建议权重值控制在 1.0 到 5.0 之间超过 5.0 会导致梯度爆炸域。但这个技巧默认不启用先跑通基线再说。3.2 训练循环代码五段式结构直接抄训练循环的骨架分成五步前向、算损失、反传、优化器 step、EMA 或者学习率调度更新。poly 学习率是这个流程里和别的分类任务差异最大的一处——它不像 StepLR 那样到某个 epoch 骤降而是每个 step 按照(1 - iter/total_iters) ** 0.9平滑衰减这是 DeepLabv3 论文的原始设计。写成代码import math import torch from torch.cuda.amp import autocast, GradScaler def train_one_epoch(model, loader, optimizer, criterion, epoch, total_epochs, total_iters, scaler): model.train() running_loss 0.0 for batch_idx, (images, labels) in enumerate(loader): images images.cuda() labels labels.cuda() # labels 里 255 的像素会被损失函数自动忽略 # poly 学习率base_lr * (1 - iter/total_iters)^0.9 current_iter epoch * len(loader) batch_idx lr optimizer.param_groups[0][lr] if current_iter 0 else 0 for param_group in optimizer.param_groups: param_group[lr] param_group[initial_lr] * (1 - current_iter / total_iters) ** 0.9 optimizer.zero_grad() with autocast(): outputs model(images) # outputs 如果是 dicttorchvision 风格取 [out] logits outputs[out] if isinstance(outputs, dict) else outputs loss criterion(logits, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() running_loss loss.item() if batch_idx % 50 0: print(fEpoch {epoch} | Iter {batch_idx} | Loss {loss.item():.4f}) return running_loss / len(loader)几个要重点说明的参数total_iters是总迭代次数而不是总 epoch 数它的计算方式是epoch * len(train_loader)在训练开始前就要算好否则 poly 衰减的曲线会错位param_group[initial_lr]需要在初始化优化器之前把初始学习率存进去不能等到循环里再读因为第一个 step 完成之后param_group[lr]已经被改掉。AMP 的GradScaler在scaler.scale(loss).backward()和scaler.step(optimizer)的顺序上不能写反先 scale 再反传才能保证 fp16 的梯度不溢出。model(images)返回 dict 是 torchvision 分割模型的固有行为aux_classifier在训练阶段的输出可以直接忽略或者拿来算辅助损失我一般直接不用。3.3 验证循环里的 mIoU 计算要点混淆矩阵在 CPU 上累加验证循环的正确性直接影响能否判断模型收敛。分割任务里只观察 loss 完全没有意义loss 降得再低 mIoU 都可能原地踏步因为 255 的像素被忽略后 loss 的绝对数值不再可比。我在验证阶段用滑窗评估单尺度输入分辨率与训练一致但把每一步的预测 argmax 从 GPU 拷回 CPU在 CPU 上维护一个confusion_matrix np.zeros((num_classes, num_classes))来累计每个像素的预测对和实际对。代码片段长这样def validate_epoch(model, loader, criterion, num_classes): model.eval() conf_matrix torch.zeros(num_classes, num_classes, dtypetorch.int64) total_loss 0.0 with torch.no_grad(): for images, labels in loader: images images.cuda() labels labels.cuda() logits model(images)[out] loss criterion(logits, labels) total_loss loss.item() pred logits.argmax(dim1) # (B, H, W) target labels mask target ! 255 # 只统计非 ignore 像素 pred_masked pred[mask] target_masked target[mask] # flat 索引用于矩阵累加bincount 内存占用更小 idx target_masked * num_classes pred_masked conf_matrix torch.bincount(idx, minlengthnum_classes * num_classes).reshape(num_classes, num_classes) ious [] for cls in range(num_classes): tp conf_matrix[cls, cls].item() gt_sum conf_matrix[cls, :].sum().item() pred_sum conf_matrix[:, cls].sum().item() union gt_sum pred_sum - tp ious.append(tp / union if union 0 else 0.0) miou sum(ious) / len(ious) return miou, total_loss / len(loader)用torch.bincount替代传统的 for 循环逐类累加在大类别数Cityscapes 19 类上能省下不少验证耗时。union gt_sum pred_sum - tp是 IoU 的原始定义比直接调用sklearn.metrics.jaccard_score要快也更好控边界条件。验证时务必用model.eval()和torch.no_grad()这个不用说但容易漏的是 Dropout 和解码器里那几个 Dropout 在 eval 模式下被正确关闭——DeepLabv3 结构里的 Dropout 分布在 backbone 之后Pytorch 的eval()会统一处理不需要单独手动管。4. 训练策略与参数配置从 VOC 到 Cityscapes 的迁移和显存平衡4.1 骨干网络预训练权重与冻结策略的选择ResNet101 的 backbone 强烈建议加载 ImageNet 预训练权重torchvision.models.segmentation.deeplabv3_resnet101(weightsDeepLabV3_ResNet101_Weights.COCO_WITH_VOC_LABELS_V1).backbone甚至可以直接拿到在 COCO 上训过分割任务的骨干这个权重的收敛速度比纯 ImageNet 权重快很多。如果用的是纯 torchvision 的 DeepLabv3 权重它的 voc 分支已经对齐了 VOC 的 21 类跑 VOC 任务时可以直接复用跑 Cityscapes 时 backbone 可以复用但最后的分类层必须重新初始化。小数据集上我习惯冻结 backbone 的前三个 stage。冻结的方式不是设requires_grad False就完事BatchNorm 的 running_mean 和 running_var 仍然会被更新这会导致前向统计量漂移。正确做法是给被冻结模块里的所有 BatchNorm 显式设track_running_stats False或者干脆把整个冻结模块切到 eval 模式。VOC 数据量小冻结 conv1 到 layer1 是个保险策略Cityscapes 2975 张图足够微调我一般只加载预训练权重做全量训练不冻结任何层衰减设 1e-4 就够。4.2 输入尺寸、Batch Size 与学习率的三角关系DeepLabv3 的输入尺寸直接决定显存占用和感受野覆盖VOC 和 Cityscapes 的最佳配置完全不同。下面这个参数表是我在两个数据集上验证过的起点配置数据集输入尺寸Batch Size初始学习率权重衰减骨干输出 strideVOC513x513162 卡 x80.0071e-48 或 16Cityscapes769x76982 卡 x40.011e-48VOC单卡513x51340.0031e-416单卡 11G 显存上 resnet101 769x769 最多塞下 batch size 2此时如果还照搬论文的 lr0.01 基本第一轮就爆。原因是 BatchNorm 在 batch size 2 时统计量几乎不可用两个方向的建议做法是要么把输入降到 716x716 配合 batch size 4要么使用同步 BatchNorm 撑到两卡以上。这个坑极其常见表现为训练 loss 一会儿大一会儿小最后 mIoU 卡在 50 上下不去。output_stride参数控制的是骨干网络最后两个残差块的 stride设为 8 时 ASPP 的特征图分辨率高一倍精度能涨 1 到 2 个点但显存开销对应上升。VOC 的训练用 output_stride16 是对的512x512 输入下感受野足够Cityscapes 建议直接用 output_stride8因为小目标比如 pole、traffic light 在 1/16 分辨率下几乎消失。4.3 多卡训练与同步 BatchNorm 的 Pytorch 写法要真正在 Cityscapes 上跑到像样的结果单卡很难在合理时间内迭代多轮。Pytorch 的标准做法是用DistributedDataParallel而不是DataParallel——后者每张卡的梯度通过主卡汇总瓶颈在通信前者是真正的进程级并行CPU 利用率也更稳定。配合同步 BatchNormtorch.nn.SyncBatchNorm.convert_sync_batchnorm(model)解决多卡之间 BN 统计量独立的问题这里有个权衡SyncBN 会拖慢约 15% 的 forward 速度但如果每张卡 batch size 只有 2收益远大于损耗。启动多卡训练的样板命令是torchrun --nproc_per_node2 train.py进程内通过torch.distributed.init_process_group(backendnccl)初始化。标注 rank 和 local_rank 的环境变量会在每次迭代取数据时决定当前进程拿哪一部分。DDP 下学习率通常不用像单卡那样调小因为梯度同步后等效于大的全局 batch sizepoly 学习率尤其吃这一套。4.4 几个常见的训练陷阱VOC 和 Cityscapes 的 label 类型必须转成torch.long再进CrossEntropyLoss否则会报数据类型错误或者静默精度下降。Cityscapes 的标签在Dataset.__getitem__里经过了 remap 转成 numpy 数组numpy 默认是 int64张量化后直接是 longVOC 的 PIL 标签转 tensor 时transforms.ToTensor()会把 0-20 归一化到小数所以 VOC 的标签不能在数据加载时用 ToTensor要手动torch.as_tensor(np.array(img), dtypetorch.long)。数据加载器的num_workers在 Windows 上必须设为 0因为多进程 DataLoader 在 Windows spawn 模式下每次都会重开解释器会拖慢而且调试麻烦Linux 上 8 到 16 都可以超过物理核心数反而会增加调度开销。pin_memoryTrue建议恒开对 GPU 和 CPU 内存之间的拷贝有明显提速代价是多占用一部分主机内存VOC 和 Cityscapes 的整个训练集远小于主机内存没有压力。5. 评估与推理mIoU 复核、调色板可视化与模型导出训练完一版之后我建议先做三件事再决定要不要推倒重来用原始验证集重新计算每类 IoU把预测结果按数据集原始的 RGB 调色板可视化最后把 Pytorch 模型导出成 TorchScript 或 ONNX 供部署验证。这三个动作能最快暴露过拟合、类别映射错误和导出兼容问题。mIoU 的结果要看逐类的分布而不只看均值。Cityscapes 上常见的情况是 road 达到 97 而 pole 只有 45这两类合理差距摆在那里如果 road 掉到 90 以下大概率是 255 的 ignore 掩码写错导致背景噪声被算进公路类里。可视化时 Cityscapes 要手工套官方调色板torchvision 的draw_segmentation_masks默认是 torchvision 的随机色不利于人眼判断。我把 Cityscapes 19 类的 RGB 值直接写进推理脚本用PIL.Image.putpalette生成带调色板的 PNG比 matplotlib 逐像素映射要快一个数量级。导出模型时标题说的项目源码里一般会给一套eval.py或者predict.py我自己往往在评估完所有训练指标后补一步 ONNX 导出。TorchScript trace 对动态输入尺寸不友好固定 769x769 输入用torch.jit.trace没问题要支持任意尺寸则改用torch.jit.script把 Decoder 的整体前向包进torch.jit.script糖衣里。推理加速的最后一招是半精度直接model.half()在验证集上跑一遍Pytorch 默认的 fp32 在 GPU 上多占一倍的显存转 fp16 后显存占用减半速度大概能提 40%唯一要留意的是输入图像也得到.half()否则类型不匹配会报错。本文还有配套的精品资源点击获取
返回列表