PyTorch从零搭建DeepLabV3+:自定义数据集训练与评估全指南
开始之前认清DeepLabV3到底在做什么很多刚接触语义分割的朋友上来就翻开源码看网络结构看着看着就懵了。我当初也是这么过来的。其实在学习一个模型之前最好先搞清楚一个问题这个模型到底解决的是什么问题语义分割的任务简单说就是给图像里的每一个像素打上类别标签。比如我手里有一张街景照片模型要告诉我哪些像素是道路、哪些是行人、哪些是建筑物、哪些是天空。它不像目标检测那样只给个矩形框而是真正把物体的边界轮廓都圈出来。DeepLabV3就是目前做这件事最主流的模型之一平衡了精度和速度工程落地场景非常广。这个项目的核心目标是用PyTorch从零搭建DeepLabV3并且在自己的数据集上完成训练和评估。这里的关键词是从零和自己的数据集。从零意味着你真要自己动手把网络结构搭出来而不是直接pip装一个现成的库完事自己的数据集意味着你要处理标注、写数据加载器、适配类别、调评估指标这些都是实际项目中躲不开的硬功夫。适合谁来学我觉得三类人最合适一是刚入门语义分割、想把论文里的结构真正变成代码的研究生二是要在业务中做落地验证、需要自定义类别的算法工程师三是对PyTorch有基础了解、想通过一个完整项目把数据流、训练流程串起来的学习者。如果你连PyTorch的Tensor都不会操作建议先花一周把基础补上再来不然容易被劝退。1. 准备工作先把环境和思路理清1.1 技术选型为什么用PyTorch不用其他框架标题里明确写了torch这个选择本身就是合理的。PyTorch属于动态图框架调试的时候可以随时打印中间变量的shape断点打断点也不会有任何心理负担。这对于我们这种需要手工搭建、边写边验的场景非常友好。相比之下静态图框架在调试上会繁琐一些写着写着就陷入编译报错—查文档—再编译的循环。版本搭配上我自己的推荐是组件推荐版本备注Python3.8 - 3.103.10以下最稳妥有些老代码对3.11兼容不好PyTorch1.13 - 2.x2.x可以用torch.compile加速但注意算子兼容CUDA11.7 或 12.1和PyTorch版本对应即可torchvision与PyTorch同版本配套里面带了ResNet和ASPP相关定义标注工具Labelme 或 其他可视化标注工具看个人习惯我测试过PyTorch 2.0以后在RTX 30系、40系显卡上跑DeepLabV3都挺稳但如果你是在公司服务器上部署建议先确认驱动和CUDA版本再决定装哪个版本不然装到一半爆出个version mismatch真的很影响心态。安装没什么好说的去PyTorch官网选对应配置生成命令即可。但有一点我要特别提醒虚拟环境一定要建最好用conda或者python -m venv隔离开。我吃过亏直接把各种包装进base环境后来某个项目要装旧版numpy结果把其他项目的依赖搞炸了重装了半天。1.2 项目结构规划条理清晰的目录能救你一命不要小看项目结构这件事。像DeepLabV3这种项目涉及数据加载、模型定义、训练脚本、评估脚本、可视化工具如果不分模块全塞在几个大文件里后面改一处就要翻半天。我建议按下面的方式组织deeplabv3plus-project/ ├── config.py # 所有超参数集中管理 ├── datasets/ │ ├── __init__.py │ ├── augmentation.py # 数据增强 │ └── voc_like_dataset.py # 自定义数据集加载器 ├── models/ │ ├── __init__.py │ ├── backbone.py # ResNet骨干网络 │ ├── aspp.py # ASPP模块 │ ├── decoder.py # 解码器 │ └── deeplab_v3plus.py # 主模型组装 ├── tools/ │ ├── train.py # 训练入口 │ ├── evaluate.py # 评估入口 │ └── predict.py # 单图推理入口 ├── checkpoints/ # 模型权重保存 ├── runs/ # TensorBoard日志 └── data/ ├── images/ # 原始图片 ├── masks/ # 标注mask └── index/ # 训练/验证划分文件这种组织的核心思路是参数、数据、模型、流程四者解耦。以后你想换数据集只改datasets和config想换骨干只动backbone想调训练策略只改train.py。这事前期的几分钟规划能给你后期省下几个小时。2. 数据处理从一堆图片到能训练的干净样本2.1 标注工具与格式转换要点训练语义分割模型最绕不开的就是标注。我自己最常用的方案是Labelme因为它操作起来直观打开图片沿着物体边缘打点画多边形保存后生成一个同名JSON文件。但这个JSON格式不能直接喂给模型需要转换成两个东西一张和原图等尺寸的灰度标签图每个像素的值代表类别索引比如0代表背景、1代表目标、2代表另一类如果是Pascal VOC风格还需要一个SegmentationClass目录存放转换后的PNG以及一个ImageSets/Segmentation/train.txt和val.txt记录参与训练和验证的图片名不带扩展名。说到转换我踩过一个大坑Labelme默认生成的标签图背景值是0但如果你有类别恰好是0号就会和背景冲突。所以转换前一定要规划好类别编号。我的习惯是背景永远是0从1开始给真实类别编号。如果你的业务场景里背景不重要还有个推荐做法是把不关心的区域统一标为255然后在损失函数里设置ignore_index255这样模型在计算损失时会自动忽略这些像素。这个技巧在做医学图像时特别常用因为很多切片上大面积的无关区域对分割结果没有意义。2.2 图片尺寸与裁切策略的考量DeepLabV3理论上支持任意输入尺寸因为全卷积网络对输入分辨率没有硬性限制。但实际操作中图像尺寸直接决定显存占用和训练速度。一张512×512的图batch size设8在12G显存上训练ResNet-101骨干的DeepLabV3勉强能跑。如果图片是1024×1024同样的batch size显存直接翻倍以上。所以我强烈建议把训练尺寸先统一到512×512或640×640。原图长宽比不一致怎么办两个思路直接resize实现简单但会改变物体的形变比例对边界精度略有影响先resize到长边固定再随机crop固定尺寸这样保留更多原始比例信息voc_like数据集加载时配合随机裁剪效果更好。我做项目时一般用后者。训练阶段长边resize到512然后随机crop到512×512。测试推理阶段则直接整图resize到512×512或者等比例resize到长边为512再padding到512×512。关于padding的影响后面讲推理时会细说。另外数据里不能有纯黑的图和全透明的标注清洗时要过滤掉。判断方法很简单读进来看看像素的类别数量如果少于2类这张图对训练基本没有贡献建议剔除。3. 自定义数据加载器让PyTorch读懂你的数据3.1 重写Dataset类核心逻辑拆解PyTorch的torch.utils.data.Dataset是所有数据来源的基类。我们需要继承它并实现两个方法__len__返回样本总数__getitem__返回第i个样本的(图片, 标签)对。这一段是纯手动搭建项目的第一道关卡很多人卡在这里。直接上一份可以抄作业的代码我用的是VOC风格的目录但稍作适配就能接自己的数据import torch import cv2 import numpy as np import os from torch.utils.data import Dataset from PIL import Image class VocLikeDataset(Dataset): def __init__(self, images_dir, masks_dir, file_list, base_size512, crop_size512, transformNone, label_transformNone): with open(file_list, r) as f: self.ids [line.strip() for line in f.readlines()] self.images_dir images_dir self.masks_dir masks_dir self.base_size base_size self.crop_size crop_size self.transform transform self.label_transform label_transform def __len__(self): return len(self.ids) def __getitem__(self, index): img_id self.ids[index] img_path os.path.join(self.images_dir, img_id .jpg) mask_path os.path.join(self.masks_dir, img_id .png) # 注意用PIL读图能保证通道顺序为RGB image Image.open(img_path).convert(RGB) mask Image.open(mask_path) # 这里做你定义好的数据增强 image, mask self._sync_transform(image, mask) # 转成Tensor并归一化 image np.asarray(image, dtypenp.float32) / 255.0 image image.transpose(2, 0, 1) image torch.from_numpy(image) mask np.asarray(mask, dtypenp.int64) mask torch.from_numpy(mask) return image, mask里面的_sync_transform就是你自定义的数据增强流程关键是图片和mask必须进行完全一致的几何变换否则图像旋转了标签没跟着转模型学出来的就是一堆噪点。3.2 数据增强花小力气解决大问题分割任务的数据增强和图像分类有不少区别。分类任务可以随便裁剪、翻转、调色因为语义不依赖空间位置但分割任务里每一个像素都要有对应的标签所以几何变换必须同步作用在mask上。我常用的组合有这么几个按性价比排序随机水平翻转实现成本几乎为零能让有效样本翻倍随机小角度旋转几度到十几度之间增强边界的多视角性随机裁剪在放大后的图上裁剪固定尺寸相当于做了尺度增强颜色抖动调整亮度、对比度、饱和度增强对光照变化的鲁棒性。这里重点说说随机小角度旋转的一个细节旋转后图像四个角会出现黑色的填充区域这些区域对应到原图是无效的。如果你不处理模型就会把黑色默认成背景造成误判。我的方案是在填充区域用与标签语义一致的类填填充。实现上可以用OpenCV的仿射变换设置borderModecv2.BORDER_CONSTANT图片侧填充一个全局均值比如127mask侧填充255忽略类。这样那些填充区域不会参与损失计算也不会引入虚假的背景信息。还有一类增强叫class balancing crop专门解决类别不平衡。如果你的类别面积差异巨大比如背景占95%、主体占5%每次随机裁剪大概率裁到背景。为了提高主体被采样的概率可以设置一种采样策略按类别分布地图生成采样权重前景区域权重高。这属于进阶玩法前期先别折腾等模型训练后发现问题再优化也不迟。4. 模型搭建把DeepLabV3掰开揉碎4.1 骨干网络残差与空洞卷积的配合DeepLabV3的骨干通常是ResNet但和普通图像分类用的ResNet有一个关键区别它把最后两个下采样阶段的步长改成了1用空洞卷积保持感受野。为什么要这样因为如果完全按ResNet的下采样节奏走到输出层时特征图尺寸只有输入的1/32对于分割任务来说这个分辨率太低了边界细节几乎全丢。保持到1/8或者1/16细节保留好但同时要保证感受野不缩小于是空洞卷积就登场了。空洞卷积可以理解为在卷积核的像素之间插入空格扩大感受野而不增加参数量。比如一个3×3的卷积dilation2时实际能看到的范围相当于5×5但参与计算的参数仍然是9个。这就让网络在不降低特征分辨率的条件下看得更远。在构建时我习惯先定义_resnet_backbone函数用torchvision自带的resnet50或resnet101加载预训练权重然后分别取出各层输出。torchvision把ResNet拆成了layer1、layer2、layer3、layer4操作起来很方便。需要注意的坑是torchvision自带的ResNet会根据传入的replace_stride_with_dilation参数自动调整最后两层的空洞率但我们后面还会单独换掉layer4的卷积层因为DeepLabV3原作中layer4的输出空洞率是16而如果我们把layer4的下采样步长改为1理论空洞率要改为2才能保持感受野不变。这个细节容易算错建议每一步打印feature map的形状验证。4.2 ASPP模块多尺度信息的汇聚ASPPAtrous Spatial Pyramid Pooling是DeepLab系列的标志性模块它的设计哲学很简单同一张特征图用不同膨胀率的空洞卷积并行提取然后把结果拼起来。膨胀率分别为6、12、18再加上一个全局平均池化分支最后再接一个1×1卷积把通道数统一到256。这样做的好处是同一层的特征被拆成了多个不同感受野的视角大率分支负责捕捉大范围上下文小率分支保留局部细节最终汇总后信息更完整。就有点像你看一张全景照片先看整体格局再拉近看细节两者结合才能判断出某个区域是道路还是河滩。实现代码主体部分长这样class ASPP(nn.Module): def __init__(self, in_channels, out_channels256): super().__init__() self.branch1 nn.Sequential( nn.Conv2d(in_channels, out_channels, 1, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue) ) self.branch2 self._make_atrous(in_channels, out_channels, 6) self.branch3 self._make_atrous(in_channels, out_channels, 12) self.branch4 self._make_atrous(in_channels, out_channels, 18) self.pool_branch 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 _make_atrous(self, in_ch, out_ch, dilation): return nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, paddingdilation, dilationdilation, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue) ) def forward(self, x): size x.shape[-2:] b1 self.branch1(x) b2 self.branch2(x) b3 self.branch3(x) b4 self.branch4(x) pool self.pool_branch(x) pool F.interpolate(pool, sizesize, modebilinear, align_cornersFalse) out torch.cat([b1, b2, b3, b4, pool], dim1) return self.project(out)全局平均池化分支有一个小坑要注意池化后是1×1必须用F.interpolate把它还原回原尺寸再拼接。否则每个分支的尺寸对不上torch.cat直接报错。4.3 解码器与上采样策略DeepLabV3的指的就是在原来DeepLabV3的基础上增加了一个简单而高效的解码器模块。它把主干网络中低层的高分辨率特征通常是layer2的输出通道数是中间层的比如512用1×1卷积降维至48通道然后与ASPP的输出上采样后拼接起来再经过两个3×3卷积和最终的上采样预测。为什么拼接低层特征因为ASPP的输出分辨率是原图的1/16虽然语义信息丰富但边界轮廓比较粗糙。低层特征虽然语义层次低但边缘信息保留得更好。两者拼接相当于给最终分割结果补了一层素描线条。上采样我推荐F.interpolate配合bilinear模式节省显存。不过需要注意align_corners这个参数。我一直推荐设置align_cornersFalse它在大部分情况下的对齐效果更好而True是旧版PyTorch的默认值容易让边缘有轻微锯齿。5. 训练主流程参数、损失与迭代细节5.1 损失函数交叉熵加ignore_index语义分割的损失函数最经典的就是逐像素交叉熵。对每个像素模型输出所有类别的得分经过softmax得到概率分布再和真实类别计算交叉熵。如果类别不平衡可以给不同类别分配不同的权重权重向量里背景小、前景大这样模型为了降低损失会更努力地学习小类别。PyTorch里用现成的criterion nn.CrossEntropyLoss(ignore_index255)这个ignore_index255前面已经提到过是处理标注不确定区域的神器。如果某个像素在mask里的值是255训练时这个像素就被跳过。这在弱监督或半自动标注场景下特别实用。我需要额外提醒的是不要一开始就让ignore_index和你的背景类混用。如果你的背景就是0号就把背景参与训练。只有在你确实有不care的区域时才把它们标成255并忽略。5.2 优化器与学习率策略怎么调才不玄学优化器大家默认都是Adam或者SGD分割任务里两种都可以但我更推荐先用SGD动量方法来训练语义分割模型因为它的收敛曲线更平滑泛化也好一些。具体配置optimizer torch.optim.SGD(model.parameters(), lr0.01, momentum0.9, weight_decay1e-4)学习率用多项式衰减策略。论文里的做法是lr base_lr * (1 - iter / total_iter) ^ 0.9这样的好处是前期步子大、后期步子小。在代码里可以每次迭代后手动更新也可以借助torch.optim.lr_scheduler.LambdaLR来化简def poly_lr(epoch): return (1 - epoch / total_epochs) ** 0.9 scheduler torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambdapoly_lr)之前有朋友问我为什么不用CosineAnnealing其实两种都可以但poly更适合分割任务这种动辄100轮以上的长训练因为它在训练末尾仍然有相对可用的学习率不会像cosine那样跌到几乎为0导致后期基本没有微调能力。当然这也是经验之谈具体效果还是要以你的实验为准。5.3 训练循环验证、保存与恢复训练循环的骨架是标准的但有几个容易被忽视的细节。首先每训练一个epoch必须在验证集上跑一次mIoU不能只盯着训练loss。我见到很多初学者只看train loss结果最后发现模型过拟合严重训练损失很漂亮验证集一测出来就崩。其次模型保存要用best checkpoint机制跟踪验证集mIoU超过历史最佳就覆盖保存同时定期保存最后一个epoch的权重。恢复训练时加载模型参数、优化器状态、scheduler状态以及epoch和best_miou的数值。这个习惯在你训练到一半断电或者显存爆掉重跑时能省下好几个小时。还有一个关键点是梯度裁剪。如果batch size较大或学习率较高偶尔会遇到loss为NaN的情况。这时可以在反向传播后加一句torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm10)虽然不能根治问题但能防止训练直接崩溃。6. 评估标准mIoU的计算你搞懂了吗6.1 混淆矩阵与IoU的数学直觉训练完后怎么评价模型好不好分类任务看准确率就行但分割任务里准确率会骗人。举个例子假设一张图里道路占90%、行人和车辆加起来占10%模型把所有像素都预测成道路准确率是90%看着很高但实际上一个物体都没分割出来。所以分割任务用mIoU平均交并比作为核心指标。对每个类别计算模型预测结果和真实标注的交集面积除以并集面积得到该类别的IoU然后对所有类别取平均就是mIoU。IoU的直觉理解就是预测区域和真实区域的重叠程度完全重叠就是1完全错开就是0。它的数学实现依赖于混淆矩阵一个K×K的矩阵K为类别数第i行第j列表示真实类别i被预测为类别j的像素数。类别i的IoU就是[ IoU_i \frac{conf[i][i]}{\sum_{j} conf[i][j] \sum_{j} conf[j][i] - conf[i][i]} ]6.2 谐波trick用one-hot矩阵高效计算在代码里直接遍历所有测试图片收集像素级别的prediction和target然后用sklearn.metrics或者手动矩阵运算都行。我建议手动用one-hot展开的方式写可控性更高避免引入不必要的依赖def compute_miou(pred_mask, target_mask, num_classes, ignore_index255): pred_mask pred_mask.flatten() target_mask target_mask.flatten() valid target_mask ! ignore_index pred_mask pred_mask[valid] target_mask target_mask[valid] pred_onehot np.eye(num_classes, dtypenp.float32)[pred_mask] target_onehot np.eye(num_classes, dtypenp.float32)[target_mask] intersection pred_onehot * target_onehot union pred_onehot target_onehot - intersection iou_per_class intersection.sum(axis0) / (union.sum(axis0) 1e-10) return np.nanmean(iou_per_class)这个技巧是通过one-hot向量做比特级的运算两个one-hot向量的逐元素乘积只有在类别完全一致时才有非零值从而天然算出交集逐元素加法然后减法则能算出并集。1e-10这个小epsilon防止除零尤其是某个类别在验证集中从未出现时很有用。评估时注意最终mIoU应该是对整个验证集所有图片的混淆矩阵一次性统计而不是单张图算完再平均。这两种算法在类别分布不均匀时差异明显前者更接近真实水平。7. 推理与部署让模型用起来7.1 单图推理的预处理细节模型训练完最终是要拿来用的。推理时的预处理和训练时有所不同。由于我们训练时用固定尺寸crop推理时整图resize到相同尺寸再输入就行。但有个问题如果原图尺寸和训练尺寸长宽比不一样直接resize会变形。更稳的做法是等比例缩放使长边接近目标尺寸然后四周padding到目标尺寸padding区域的值填0或者127记录padding的位置最后把预测结果裁剪回去。def preprocess_for_inference(image_path, target_size512): img cv2.imread(image_path) h, w img.shape[:2] scale target_size / max(h, w) new_h, new_w int(h * scale), int(w * scale) resized cv2.resize(img, (new_w, new_h)) pad_h target_size - new_h pad_w target_size - new_w padded cv2.copyMakeBorder(resized, 0, pad_h, 0, pad_w, cv2.BORDER_CONSTANT, value(127,127,127)) return padded, (new_h, new_w, scale)推理时的预测可以采用TTATest Time Augmentation把图片水平翻转后再预测一次把两次预测的概率取平均后再取argmax。这个操作对分割结果的提升幅度通常有0.5%-1%的mIoU代价只是推理时间翻倍。如果你对部署效率敏感可以只在离线评估时用在线部署时不启用。7.2 可视化与导出从Tensor到业务模块把预测的tensor还原成可视化的彩色mask我通常会准备一份调色板字典给每个类别一个RGB颜色。背景通常用黑色0,0,0其他类别用醒目的红绿色系。Python配合OpenCV把预测值映射成彩色图然后与原图叠加透明度输出一张叠加图方便肉眼检查效果。导出模型时记住不要把整个训练状态导出只导出model.state_dict()就够了。如果还需要进一步加速推理可以把模型转成TorchScript或者ONNX后接TensorRT。需要注意在转换时要把BatchNorm层设置为eval模式不然BatchNorm的运行均值方差还在更新导出的模型推理结果会不稳定。我遇到过坑导出后推理效果和平时的验证集评分差了一大截排查半天才发现是忘了model.eval()。8. 常见问题与排查技巧实录现象可能原因解决方案训练时loss变成NaN学习率过大、输入数据有NaN、梯度爆炸降低学习率、清洗数据、加梯度裁剪mIoU一直很低且不涨标注有错、类别严重不平衡、数据增强过度检查mask可视化、加类别权重、简化增强显存不足OOMbatch size过大、图片过大调小batch size、减小输入尺寸、开启混合精度预测结果边缘粗糙低层特征利用不足、解码器维度不匹配检查低层特征是否拼接对、加大深层特征分辨率训练和验证差距巨大过拟合加dropout、数据增强、权重衰减这里我特别想展开说说我在实际踩过的两个问题这两个问题新手几乎都会遇到但网上很少有人明确点破。第一个是Mask的读取格式问题。很多人像读普通图片一样用cv2.imread读标注mask但OpenCV读图默认是三通道BGR而标注mask是单通道的索引图。如果标注工具保存的是彩色PNG但你用三通道读进来后面的数值索引就会乱七八糟。正确做法是读的时候用cv2.IMREAD_GRAYSCALE或者直接用PIL打开并且要确认标注图的像素值确实是整数类别索引而不是空隙类别被压缩成了0到255之间的浮点颜色值。第二个是预训练权重加载后BatchNorm的有效性。直接用torchvision的resnet50(weightsResNet50_Weights.IMAGENET1K_V1)加载预训练权重encoder前几层很快能收敛。但当你修改了layer4的下采样策略后torchvision的权重仍然能用因为layer4内部的卷积参数没有变化只改了下采样的步长和空洞率。这点很多教程含糊带过实际上直接加载官方预训练权重即可不要自己瞎训练骨干。我自己试过不加载预训练从头训练同样的训练轮数mIoU大概会低8到10个点而且收敛极慢。9. 后续还能怎么玩从基础版到进阶版如果你顺利跑通了上面的全流程你已经算入门了。接下来有几条进阶路线根据自己的需求选第一条是把骨干换成MobileNetV3或者轻量级网络压缩参数量让模型在嵌入式设备上跑得动。DeepLabV3的解码器部分是设了低层特征通道数的MobileNetV3的低层通道数少拼接时需要调整1×1卷积的输出通道数。我自己做过一次这个改动模型体积从200MB降到不到30MB推理帧率提升5倍以上精度降幅在3个点以内对很多移动端场景完全够用。第二条是引入类别引导的损失优化比如Dice Loss或者Focal Loss。如果你的数据里某几个类别特别小交叉熵很难把那些像素的前景响应拉起来。Dice Loss直接把IoU作为优化目标对小物体更友好但具体效果要看实际的数据分布。我常用的做法是让CrossEntropyLoss和DiceLoss按0.7:0.3的比例相加两种损失互补。第三条是尝试边界感知机制在解码器后额外接一个边界预测头或者设计专用的边界损失项。这个方向在医学图像分割和遥感图像分割中效果显著感兴趣的朋友可以搜一下边界损失相关的英文关键词去查论文。我在实际项目中最常碰到的场景其实是第一种——精度够用、越快越好。毕竟算法要落到业务里不能只在实验室里跑速度、内存、稳定性都是要综合考虑的。最后再分享一个小技巧在训练之前先拿一两张图跑一遍完整的前向和反向传播确保没有shape错误和维度不匹配再启动全量远训练这一点能帮你省下大量试错的时间。