简介面向图像分割初学者与进阶开发者的Pytorch实战资源以DeepLabv3为核心算法覆盖VOC与Cityscapes两个公开数据集的完整训练流程解决从模型搭建到训练评估、推理预测的全程落地问题。压缩包共55个文件、大小约2.25MB主要包含23个Python源码、17张训练样本与预测结果对比图、9个备份文件以及TXT/README说明文档源码按datasets、network、utils等模块组织便于阅读与二次开发。核心代码涵盖DeepLabv3的建模、ResNet/Xception等骨干网络、数据加载与增强、损失函数、学习率调度、可视化以及预测脚本可直接用于实验复现或改造为自有数据集。目前已有139人下载学习适合需要一份简洁可跑的图像分割工程模板的读者。1. 图像分割实战为什么我建议从这份 DeepLabv3 源码入手图像分割项目的难点从来不在模型结构而在环境、数据格式、训练参数和评估这一整条链路上。见过太多人拿 DeepLabv3 跑 VOC 没问题一换到 Cityscapes 就翻车原因大多不是网络写得不好而是数据集的标注体系没对齐、类别索引错位、mIoU 一路为 0。这份基于 Pytorch 的 DeepLabv3 项目源码把 VOC 和 Cityscapes 两条数据流都打通了训练、预测、可视化、指标评估一应俱全适合有 Pytorch 基础、想真正落地图像分割训练流程的读者。我从这份源码里拆出骨干网络替换、数据索引对齐、损失函数配置、断点续训等几个关键环节按自己实际跑通的经验把步骤和坑写出来。2. 先读懂 DeepLabv3骨干网络、ASPP 与解码器的设计逻辑2.1 Encoder-Decoder 结构空洞卷积和 ASPP 到底在干嘛DeepLabv3 相比前代版本最大的变化是把模型改成了 Encoder-Decoder 结构。Encoder 部分沿用 DeepLabv3 的带空洞卷积的骨干网络加 ASPP 模块Decoder 部分则负责把低层特征和高层语义特征融合起来恢复分割结果的边缘细节。ASPP 是 Atrous Spatial Pyramid Pooling 的缩写核心思想是在同一个特征图上用不同膨胀率的空洞卷积并行采样捕获多尺度上下文。项目源码里可以在network/_deeplab.py中看到 ASPP 的实现它一般包含 1x1 卷积、三个不同 rate 的 3x3 空洞卷积以及一个全局平均池化分支最后把所有分支的输出拼接起来再过一层卷积。简单理解就是用四个不同视野范围的“眼睛”同时看同一张图然后把看到的结果拼在一起这样模型对大目标和小目标都有感知能力。空洞卷积改变 rate 之后感受野会变大但特征图的尺寸不变。这一点对分割任务很关键因为分割需要保留空间位置信息不能像分类那样一路下采样到底。代码里通常会通过aspp_padding、aspp_rates这样的参数来控制膨胀率组合常见配置是 6、12、18。Decoder 部分则把骨干网络浅层特征和 ASPP 输出做拼接再经过卷积和上采样恢复到原图尺寸。源码里 Decoder 一般有一段用low_level_features的逻辑把骨干网络第一层输出的浅层特征拿过来与上采样后的高级特征拼接目的是补回被多次下采样损失掉的边界细节。2.2 骨干网络选型ResNet、Xception、MobileNetV2 怎么挑项目目录network/backbone/下放了resnet.py、xception.py、mobilenetv2.py和hrnetv2.py这就是这个源码库比一般教程值钱的地方——同一个训练流程可以切换四个骨干。选哪个骨干不是越深越好而是看你要在什么设备上跑、对精度和速度的要求是多少。最省事的是 ResNet。用resnet.py里预训练好的 ResNet-101 作为骨干精度不错显存消耗属于中等水平大部分单卡 11GB 或 16GB 的机器可以扛住。在 VOC 上做 PASCAL VOC 2012 增强集的训练大概 3 万到 5 万步能收敛到不错的 mIoU。Xception 是 DeepLabv3 原论文用的骨干效果理论上最好但计算量很大训练速度明显慢一截显存占用也高。MobileNetV2 是为轻量级场景准备的适合在边缘设备或者实时性要求高的场景下推理。实测用 MobileNetV2 作骨干mIoU 会比 ResNet-101 低四到五个点但单张推理速度能压到几十毫秒。HRNet 则是另一条路线通过并行高分辨率分支保持空间细节在小目标分割场景表现不错但显存占用也不低。我的习惯是先在 ResNet-101 上把整套流程跑通确认数据、训练、评估都正常了再剪枝或替换成 MobileNetV2 做加速。骨干的切换入口通常在训练参数里指定--backbone代码内部通过build_backbone这类函数把骨干网络和 ASPP 接起来。如果你要换骨干注意两个地方一是预训练权重的下载路径是否正确二是骨干输出通道数是否和 ASPP 的输入通道数匹配。2.3 代码库目录拆解main.py、network 与 datasets 各自的边界拿到这份源码之后先别急着运行花十分钟把目录结构理清楚。这个项目的核心模块划分是这样的目录 / 文件职责main.py训练和评估的统一入口负责解析参数、加载数据、初始化模型、执行训练循环和验证循环predict.py推理入口加载训练好的权重对单张图片或一组图片做分割预测并保存可视化结果network/模型定义包含_deeplab.py网络主体、backbone/骨干网络实现、modeling.py组装逻辑datasets/数据集加载voc.py、cityscapes.py、utils.py分别对应 VOC、Cityscapes 和通用数据工具utils/工具层包括loss.py损失函数、scheduler.py学习率调度、ext_transforms.py数据增强、visualizer.py结果可视化metrics/评估指标stream_metrics.py实现了流式的 mIoU、mPA 等指标累计计算注意utils/目录里有不少.zbak后缀的备份文件比如main.py.zbak、__init__.py.zbak这是作者调试时留下的不用理会也不影响运行。main.py是训练入口它会解析命令行参数完成模型初始化和数据加载然后进入训练循环。network/目录只负责模型前向计算不管数据和损失。utils/loss.py定义损失函数metrics/stream_metrics.py定义评估指标。这种解耦的好处是你可以单独替换损失策略或者评估方式不用动整体流程。另外samples/目录里已经有训练好的效果图比如1_overlay.png、23_pred.png、114_target.png、city_6_overlay.png这些分别是原图叠加分割结果、预测结果和标注标签的样例。跑完predict.py之后输出格式和这些示例是一致的可以用这些图对照检查自己的预测效果是否正常。3. 环境搭建与双数据集准备从 Pytorch 安装到 label 索引对齐3.1 Pytorch 环境与依赖安装这份源码是基于 Pytorch 的动手之前先把环境装好。安装 Pytorch 之前先确定自己的 CUDA 版本在终端里执行nvidia-smi看右上角的 CUDA Version这个值是你驱动支持的最高版本。然后根据这个版本去官网上选对应的安装命令。比如驱动支持 CUDA 11.8就可以安装对应版本的 Pytorch。# 创建虚拟环境Python 版本建议 3.8 或 3.9 conda create -n deeplab python3.9 conda activate deeplab # 以 CUDA 11.8 为例安装 Pytorch pip install torch2.0.1 torchvision0.15.2 --index-url https://download.pytorch.org/whl/cu118 # 安装项目依赖 pip install numpy matplotlib pillow tqdm visdom tensorboard依赖列表里有几个重点说明。visdom是用于可视化训练曲线的main.py里通过visualizer.py连接 visdom 服务训练过程中能看到 loss 曲线和 mIoU 变化。visdom需要先启动服务端再运行训练脚本。tensorboard是可选的如果你更习惯 TensorBoard 就用它两个可以同时开不影响训练主流程。装完Pytorch之后验证一下 GPU 是否可用import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0) if torch.cuda.is_available() else CPU)打印出来的torch.cuda.is_available()必须为True否则后续训练会掉到 CPU 上速度慢到无法接受。这一步经常出问题的地方是 Pytorch 版本和 CUDA 版本不匹配常见的报错是CUDA initialization: Unexpected error。遇到这种情况先检查驱动是否安装成功再确认 Pytorch 对应版本的 CUDA 是否和驱动兼容。3.2 VOC 数据集目录结构与标注格式VOC 数据集的标注是 PNG 格式的调色板图也就是说像素值不是类别 id而是调色板索引。训练之前一般要下载 PASCAL VOC 2012 的原始数据集加上 SBD 增强数据合并成训练集。项目里datasets/voc.py就是负责加载 VOC 数据的模块train_aug.txt文件里记录的是增强训练集的图片索引列表。VOC 原始目录结构比较简单但有几个容易忽略的地方VOCdevkit/ ├── VOC2007/ ├── VOC2012/ ├── ImageSets/ │ └── Segmentation/ │ ├── train.txt │ ├── val.txt │ └── train_aug.txt ├── JPEGImages/ └── SegmentationClass/JPEGImages存放原始 RGB 图片SegmentationClass存放分割标签 PNG。标签图片每个像素的颜色对应一个类别读取时需要把 RGB 颜色转换回类别索引。datasets/voc.py中一般会实现一个class VOCSegmentation类核心逻辑是读图、读标签、做数据增强、返回图像张量。文件里通常会定义color_map或直接调用 PASCAL VOC 的标准调色板来把标签图像转为类别索引。转换的逻辑大致是# 读取标签图像转为调色板索引 target Image.open(os.path.join(root, SegmentationClass, name .png)) target_np np.array(target, dtypenp.uint8)这里要特别留意VOC 的标签 PNG 用PIL的Image.open打开后是P模式直接转np.array得到的是调色板索引数组数值范围是 0 到 20其中 255 表示忽略区域。如果你用cv2.imread去读读出来的是 BGR 三通道的 RGB 值需要先通过调色板映射回索引这一步做错了后面损失函数计算就会出问题。train_aug.txt在项目里被直接放在datasets/目录下作者应该是已经把增强集文件准备好了。如果你自己下载 SBD 增强数据需要保证train_aug.txt中列出的图片名称在JPEGImages和SegmentationClass中都存在否则数据加载会报找不到文件。3.3 Cityscapes 的 trainId 与 id最容易踩的数据坑Cityscapes 是自动驾驶场景的经典数据集和 VOC 最大的区别是它的标注体系有 id、trainId、category 等多层映射。原始像素值并不是直接可用的类别编号必须把id转换成trainId才是训练目标的类别。Cityscapes 数据集下载完成后目录结构是leftImg8bit/ ├── train/ ├── val/ gtFine/ ├── train/ ├── val/其中gtFine下的*_gtFine_labelIds.png存的是原始 id 的标签图而*_gtFine_labelTrainIds.png是已经转换好的 trainId 标签图。你训练时应该使用labelTrainIds.png不是labelIds.png。如果错误地用了原始 id类别数会变成 34 个而不是 19 个因为原始 id 包含了未标注区域、不同 group 的 id 等。datasets/cityscapes.py里的加载逻辑会做类似下面的处理target Image.open(gt_path) target_np np.array(target) # 只保留 trainId 在 0~18 范围内的像素其余全部置为 ignore_index target_trainId target_np.copy() for id, trainId in id_to_trainId_map.items(): target_trainId[target_np id] trainId注意项目里的ignore_index一般设为 255Cityscapes 中 trainId 为 255 的像素表示忽略区域包括未标注的物体、边界和难以分类的区域。这些像素不参与损失计算也不会被计入 mIoU 的评估。Cityscapes 的图片尺寸是 2048x1024直接整图训练显存压力很大。项目里通常会做 crop 或者 resize常见做法是随机裁剪成 512x1024 或 768x768 的 patch。如果你的显存是 11GB建议 crop 尺寸控制在 512x1024 以内Batch Size 取 2 或 4。另外 Cityscapes 的类别分布极度不均衡像sidewalk、terrain这些类别的像素占比远低于road训练的收敛速度会比 VOC 慢不少这是正常现象。4. 训练到预测的完整闭环main.py 参数、损失函数与 mIoU 评估4.1 训练入口参数backbone、lr、epochs 与 Batch Size 的配合main.py是整套训练的入口它用 argparse 解析参数。命令行参数比较多下面是常用的一组训练命令python main.py \ --model deeplabv3plus \ --backbone resnet \ --dataset voc \ --data-root /path/to/VOCdevkit \ --train-list datasets/train_aug.txt \ --val-list datasets/val.txt \ --batch-size 16 \ --lr 0.01 \ --epochs 50 \ --crop-size 513 \ --gpu-id 0参数说明--model指定模型结构--backbone指定骨干网络--dataset用来切换 VOC 和 Cityscapes 的数据加载逻辑--data-root是你存放数据集的根路径--train-list和--val-list指向训练集和验证集的图片列表文件。--batch-size和--lr是训练的两个核心参数它们需要配合调整。Batch Size 对学习率的设置有直接影响。DeepLabv3 在 Pytorch 实现里常用 poly 学习率策略初始学习率一般设置在 0.007 到 0.01 之间如果 Batch Size 减半学习率也应该相应降低。比如 Batch Size 从 16 降到 8学习率最好从 0.01 降到 0.007 左右否则梯度更新步长偏大前期 loss 容易震荡。--crop-size是随机裁剪尺寸VOC 通常用 513Cityscapes 建议用 512 或 768。这里注意裁剪尺寸一定要能被 8 整除因为 DeepLabv3 的下采样倍率是 8输入输出尺寸不成比例会导致 decoder 上采样后尺寸不匹配报错。训练过程中main.py会每过若干个 epoch 在验证集上跑一次 mIoU 评估并打印当前学习率、loss、mIoU 等信息。如果你的显存只够跑小 Batch Size收敛速度会明显变慢这时可以把--eval-interval调大减少验证频率把更多时间花在训练上。4.2 损失函数与学习率调度交叉熵、OHEM 和 poly 策略utils/loss.py里封装了损失函数。默认使用的是带 ignore_index 的交叉熵损失。交叉熵处理多分类分割任务没问题但遇到类别严重不均衡的数据集时效果就一般。项目里一般会预留 OHEM 的开关OHEM 的意思是 Online Hard Example Mining即在线难例挖掘。开启 OHEM 之后损失函数会挑选 loss 值最大的前 N 个像素参与反向传播其他简单像素不贡献梯度。这能有效解决类别不平衡问题因为大多数像素都是背景或大块类别只算这些像素会让模型忽视小目标。OHEM 在 Cityscapes 这种类别分布差异巨大的数据集上效果提升明显但对显存有额外要求因为需要先完整算出所有像素的 loss再做 topk 选取。# utils/loss.py 中 OHEM 的简化逻辑 def ohem_loss(logits, target, ignore_index255, thresh0.7, n_min10000): # 先按交叉熵算每个像素的损失 per_pixel_loss F.cross_entropy(logits, target, ignore_indexignore_index, reductionnone) # 选出难例像素参与反向传播 hard_pixels per_pixel_loss thresh if hard_pixels.sum() n_min: hard_pixels per_pixel_loss.topk(n_min).indices return per_pixel_loss[hard_pixels].mean()这里的核心参数是thresh和n_min。thresh控制哪些像素被认定为“难例”loss 大于这个阈值的像素才会被选中。n_min是最少保留的像素数量防止全图都是简单样本时选不出足够的难例。utils/scheduler.py实现了 poly 学习率策略poly 策略的公式是lr base_lr * (1 - iter/total_iter)^powerpower 默认取 0.9。和 step 策略相比poly 的曲线变化更平滑前期的学习率下降较慢后期加速衰减比较适合分割任务的长训练过程。理论上如果你用 step 策略需要在特定 epoch 手动降低学习率两次具体时机要靠经验判断而 poly 策略基本不用管训练结束后学习率自动降到接近零。项目里两种调度方式都有实现选择哪个可以在main.py的参数里控制。4.3 边训练边看效果visualizer 与 stream_metrics 的配合训练过程中除了看 loss 数字更直观的判断方式是可视化预测结果。项目里utils/visualizer.py封装了 visdom 可视化逻辑。使用 visdom 之前需要先启动服务端# 先启动 visdom 服务端 python -m visdom.server -port 8097然后在另一个终端运行main.py训练脚本会自动连接 visdom 服务把 train_loss、mIoU、learning_rate 等曲线实时推送到浏览器。metrics/stream_metrics.py实现的评估逻辑是值得借鉴的它没有用简单粗暴的逐图计算再取平均而是维护一个全局的混淆矩阵把所有验证图片的预测结果累积到同一个矩阵里最后统一计算各类别的 IoU 和 mIoU。这么做和逐图平均的结果差别很大——逐图平均会让大图类别占比影响结果而混淆矩阵累积的方式才是严格意义上的 mIoU。# metrics/stream_metrics.py 的核心逻辑 class StreamSegMetrics: def __init__(self, num_classes): self.confusion_matrix np.zeros((num_classes, num_classes), dtypenp.int64) def update(self, pred, target): # 将预测和真实标签展平忽略 ignore_index pred pred.flatten() target target.flatten() mask target ! 255 pred pred[mask] target target[mask] # 累加混淆矩阵 np.add.at(self.confusion_matrix, (target, pred), 1) def get_results(self): # 从混淆矩阵计算每类 IoU、mIoU、mPA 等指标 intersection np.diag(self.confusion_matrix) union self.confusion_matrix.sum(1) self.confusion_matrix.sum(0) - intersection iou intersection / union miou np.nanmean(iou)验证集评估的逻辑一般写在main.py的validation函数里模型预测输出经过 argmax 得到类别索引然后传给StreamSegMetrics.update()。注意这里预测输出的 softmax 概率要转换成类别索引也就是torch.argmax(output, dim1)。VOC 数据集的 mIoU 计算通常只统计 21 个类别中的 20 个背景 其他类背景类也会参与评分。Cityscapes 则只评估 19 个类别未标注的区域不参与。所以你训练完成后报告的 mIoU 数值在这两个数据集上不是一个可比的口径对比不同论文结果时要注意这一点。4.4 用 predict.py 跑通单张与批量推理训练完成之后权重会保存为.pth文件用predict.py可以加载权重对图片做推理。预测的命令示例如下python predict.py \ --checkpoint checkpoints/best_deeplabv3plus_resnet_voc.pth \ --input samples/1_image.png \ --output samples/1_pred.png \ --dataset voc \ --backbone resnet \ --crop-size 513--checkpoint指向保存的权重文件--input是输入图片路径--output是预测结果保存路径。--dataset和--backbone要和训练时保持一致否则权重加载会出现 shape 不匹配的错误。推理代码的内部逻辑是先加载模型结构和权重然后对输入图片做预处理resize 到固定尺寸、减去均值、除以标准差、转成 tensor 并增加 batch 维度。前向计算之后用 softmax 拿到每个像素的概率分布再 argmax 得到类别索引。最后根据数据集的调色板把类别索引映射成可视化颜色生成分割效果图。项目里samples/目录下的23_pred.png、1_pred.png就是这种推理输出的例子。另外还有overlay文件比如1_overlay.png是把原始图片和半透明分割结果叠加在一起的混合图。predict.py里也实现了这个功能输出时可以同时生成纯分割图和叠加图方便直接观察分割边界和实际物体的贴合程度。推理时还有一个重要参数是输入尺寸。推理时不用必须和训练时的 crop size 完全一致但建议保持一致因为模型在固定输入尺寸上训练的突然换一个更大的输入尺寸可能会降低精度。如果你希望提升小目标的检测效果可以尝试推理时把输入尺寸调大比如从 513 放大到 769mIoU 通常会有一两个点的提升代价是推理时间变长。5. 避坑手册与进阶技巧从权重恢复训练到 ONNX 导出5.1 三条血泪踩坑记录坑一Cityscapes 的 mIoU 一开始就是 0所有类别都不对现象训练和验证都能跑但验证集 mIoU 始终是 0 或者一直在个位数徘徊分割输出都是黑色或噪声。原因Cityscapes 的原始labelIds.png和labelTrainIds.png混用了。用原始 id 训练时类别数和模型的输出通道数不匹配模型实际上在预测 34 类而评估时只用 19 类导致索引全部错位。解决检查数据集加载时用的是否是*_gtFine_labelTrainIds.png。如果不是改成 trainId 版本的标签图。改完之后在加载函数里打印一下np.unique(target_np)确认数值范围在 0 到 18 之间如果有相应的 id 映射逻辑重新映射一遍。坑二visdom 启动后训练卡住或者训练正常但浏览器里看不到曲线现象python -m visdom.server启动服务后训练脚本卡在初始化阶段或者训练能跑但浏览器打开 8097 端口是空白页面。原因visdom 服务端和训练脚本之间的连接初始化需要握手通信如果 visdom 版本和 Pytorch 环境不兼容或者服务端没有正常启动训练脚本会在等待连接时阻塞。解决先确认 visdom 服务是否启动成功浏览器访问http://localhost:8097能看到页面再跑训练脚本。如果初始化阶段一直阻塞可以在代码里把 visdom 初始化放在训练循环之后或者用环境变量VISDOM_USE_ENV_OPTS调整。更稳妥的方案是把visualizer.py里的 visdom 逻辑换成 TensorBoardSummaryWriter的兼容性更好几乎没有连接问题。坑三损失函数一直在下降但 mIoU 不涨甚至越训越差现象训练 loss 稳步下降验证集的 mIoU 却停滞不前过了一定 epoch 之后开始缓慢下降明显过拟合了。原因数据增强不够或者训练集太小。VOC 原始训练集只有 1464 张图如果只用原始 train 集不增强模型很容易在训练集上过拟合。解决项目里utils/ext_transforms.py实现了多种分割专用的数据增强方法比如随机翻转、随机缩放、随机裁剪、颜色抖动等。训练时打开随机缩放和颜色抖动把训练集扩增到原始数据的十几倍。同时注意增强操作要保持图片和标签的同步ext_transforms.py里应该有Compose之类的类来同时处理 image 和 target。5.2 常用技巧恢复训练、学习率热启动与骨干替换训练到一半断掉是常态代码里一般预留了--resume参数可以加载上次保存的权重继续训练python main.py \ --resume checkpoints/checkpoint_epoch_30.pth \ --epochs 100 \ --start-epoch 31--start-epoch需要手动指定否则模型会从 epoch 0 开始重新计数学习率和数据增强的随机状态会错乱。恢复训练时优化器状态要不要加载取决于你是不是对模型结构做过改动。如果只是中断后续训建议把优化器状态也保存下来否则学习率曲线会有一个较大的跳变。学习率热启动是一个实用技巧做法是所有 epoch 的前几个迭代用很小的学习率做 warmup然后恢复到正常初始学习率这样能避免在预训练权重上过早震荡。部分实现里会在main.py的adjust_learning_rate里加一个warmup_epochs参数前 3 个 epoch 把学习率从 1e-5 线性升到基础学习率。5.3 导出 ONNX 做部署验证训练出满意的权重之后经常需要把模型导出成 ONNX 格式做部署测试。这是 Pytorch 自带的功能导出命令如下import torch from network.modeling import deeplabv3plus # 初始化模型并加载权重 model deeplabv3plus( n_classes21, backboneresnet, backbone_nameresnet101 ) checkpoint torch.load(checkpoints/best.pth, map_locationcpu) model.load_state_dict(checkpoint[model_state]) model.eval() # 构造示例输入导出 ONNX dummy_input torch.randn(1, 3, 513, 513) torch.onnx.export( model, dummy_input, deeplabv3plus_resnet101.onnx, opset_version11, input_names[input], output_names[output] )导出后建议先用 ONNX Runtime 验证一下结果和 Pytorch 是否一致特别是分割模型的输出是 4D 张量要注意 reshape 和多 batch 的情况。5.4 借助预计算 overlay 快速判断训练方向项目samples/目录下的 overlay 图片不是随便放的它们是训练流程是否正常的可视化凭证。训练过程中每隔固定 epoch 保存*_pred.png和*_overlay.png这个方法很实用。如果 overlay 图上分割边界和物体轮廓贴合度较高说明模型学到的语义信息是可靠的如果分割结果呈现大块色斑或者边缘锯齿严重说明解码器部分还没有收敛。我自己的习惯是在训练到 20%、50%、80% 时都保存一组预测图和 overlay 图对比它们的变化趋势。如果 20% 和 50% 的效果几乎没有差别大概率是学习率设置偏小或者骨干网络没有加载预训练权重。如果 80% 的效果反而比 50% 差要警惕过拟合。从那以后我每次跑分割训练无论什么数据集都强制把评估图和 overlay 图一起存档用直观效果配合 mIoU 做双通道判断不再只看 loss 曲线一个人瞎猜。希望这份源码里的训练和可视化闭环也能帮到你。本文还有配套的精品资源点击获取