ARTICLE DETAIL

资讯详情

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

DeepLabv3+实战指南:基于Pytorch的语义分割训练与部署全流程

DeepLabv3+实战指南:基于Pytorch的语义分割训练与部署全流程 简介一套基于Pytorch实现DeepLabv3图像分割算法的实战项目资源面向视觉分割方向的学习者、科研人员及开发者可在VOC与Cityscapes两个主流数据集上完成模型训练、验证与预测推理尤其适合毕设项目、课程实验或算法对比研究。资源包共55个文件压缩后约2.25MB其中23个Python文件覆盖数据集预处理、骨干网络构建、训练调度、损失函数、指标评估与可视化等关键模块17张PNG图片直观展示输入图像、真实标签、预测结果及叠加效果便于核验算法表现另有多个zbak备份文件、txt说明与依赖配置兼顾代码可读性与容错需求。目前已有139人学习浏览。配套附赠内容与README还提供了完整目录结构和基础使用指引可帮助快速上手DeepLabv3的工程实现并迁移到自有数据集中进行二次开发。1. 图像分割实战入口VOC、Cityscapes 上把 DeepLabv3 跑通意味着什么说得直白点图像分割semantic segmentation就是把画面里的每个像素分类建筑物、道路、人、车各归其位自动驾驶感知、遥感分析和医学影像系统基本都靠它。DeepLabv3 至今仍是工程落地里最稳的图像分割算法之一结构成熟、预训练好找、显存门槛不高、导出部署方便在 PASCAL VOC 2012 和 Cityscapes 这两个公开基准上都有大量可对照的结果。这篇笔记按「Pytorch 环境搭建 → VOC/Cityscapes 数据准备 → 模型搭建 → 训练调参 → 避坑 → 推理导出」的顺序把基于 Pytorch 训练 DeepLabv3 图像分割算法的完整流程讲清楚每一步都能直接复制。想跑通一份带源码的图像分割项目却被 pytorch 安装、CUDA 版本、数据集格式劝退的人这篇就是给你写的熟手可以直接跳到避坑章对参数边界。2. 环境与数据准备把 Pytorch、CUDA、VOC 和 Cityscapes 一次配齐2.1 Anaconda 配置 Pytorch 环境python 和 pytorch 版本对应关系先搞清楚训练分割模型的第一步是环境。很多跑不通 DeepLabv3 源码的人问题不在模型而在 Anaconda 配置 Pytorch 环境时把版本装乱了。Pytorch 的 CPU 版和 GPU 版行为差异很大分割训练没有 GPU 基本是浪费时间所以这一节我们按 GPU 版来配。# 1. 创建独立环境Python 3.10 对绝大多数 Pytorch 版本都稳定 conda create -n deeplab python3.10 -y conda activate deeplab # 2. 先看显卡驱动支持到哪个 CUDA 版本 nvidia-smi # 右上角显示的 CUDA Version: 12.1 是驱动支持的最大版本不是运行时版本 # 3. 用官方索引安装 GPU 版 Pytorch这里以 CUDA 12.1 为例 pip install torch2.3.1 torchvision0.18.1 torchaudio2.3.1 \ --index-url https://download.pytorch.org/whl/cu121 # 4. 验证是否真的在用 GPU python -c import torch; print(torch.__version__, torch.cuda.is_available(), torch.cuda.get_device_name(0))逻辑说明nvidia-smi显示的是驱动支持的 CUDA 上限而 Pytorch 安装包自带完整的 CUDA 运行时两者不必完全一致只要驱动版本不低于 Pytorch 对应的 CUDA 版本即可。python 和 pytorch 版本对应关系以官网 compatibility 矩阵为准PyTorch 2.x 配 Python 3.8~3.11 都不踩雷3.10 最省心。最后一行输出里True和显卡型号都出现才算环境真的配好了。注意如果torch.cuda.is_available()返回False先重装对应 CUDA 版本的 torch别急着重装驱动。这里是最常见翻车点和源码够不够新关系不大。2.2 下载与整理 VOC 2012 和 Cityscapes目录结构、标签语义与 TrainIds 转换VOC 2012 是分类、检测、分割共用的老牌数据集分割任务只用 JPEGImages 和 SegmentationClass 两份目录。常见做法是去 PASCAL VOC 官网下载VOCtrainval_11-May-2012.tar需要更多训练数据就再下 SBD 增强包。解压后目录结构长这样VOCdevkit/ └── VOC2012/ ├── JPEGImages/ # 原图jpg约 1.5w 张 ├── SegmentationClass/ # 语义分割标签png像素值就是 0~20 的类别 id └── ImageSets/ └── Segmentation/ # train.txt / val.txt 存放图片名注意 SegmentationClass 里的 png 是带调色板的 P 模式图像素值直接就是类别索引千万别用convert(RGB)去读标签那会把调色板当颜色算喂给模型后类别全乱。VOC 一共是 21 类0 是背景1~20 分别是飞机、自行车、鸟、船、瓶子这些常见物体类别 id 在文档里写得清清楚楚。Cityscapes 是街景分割的主战场gtFine 子集大约 2975 张训练图、500 张验证图需要去官网申请下载。它和 VOC 最大的区别是原始像素是 0~33 的 labelIds比如道路是 0、人是 24但训练时要映射成 0~18 的 trainIds忽略区域标记成 255。官方工具包里有现成的createTrainIdLabelImgs.py跑一遍就会在同目录生成*_labelTrainIds.pngpython createTrainIdLabelImgs.py # 转换前后像素含义对照 # 原图 labelIds0~33包含“道路0、人24”等原始编号 # 转换后 trainIds0~18用于训练255 表示忽略区域不参与 loss这个转换不做后面 loss 一定会出问题因为 Cityscapes 的很多类别只在标注里不存在于分割目标中直接拿原始编号训练模型会把「忽略区域」学成某个类。2.3 数据加载器与增广策略训练验证划分和类别映射的核心写法自己用 Pytorch 写分割数据加载器最要紧的是保证图像和标签走同一套几何变换而且标签只能用最近邻插值。推荐直接用 Albumentations它对图像和 mask 同步处理省掉手动对齐的麻烦import albumentations as A from albumentations.pytorch import ToTensorV2 def get_transforms(crop_size513): train A.Compose([ A.RandomResizedCrop(heightcrop_size, widthcrop_size, scale(0.5, 2.0)), A.HorizontalFlip(p0.5), A.RandomBrightnessContrast(p0.2), A.Normalize(mean(0.485, 0.456, 0.406), std(0.229, 0.224, 0.225)), ToTensorV2(), ]) val A.Compose([ A.Resize(heightcrop_size, widthcrop_size), A.Normalize(mean(0.485, 0.456, 0.406), std(0.229, 0.224, 0.225)), ToTensorV2(), ]) return train, val参数说明RandomResizedCrop的scale(0.5, 2.0)表示随机裁切后缩放覆盖多尺度mask 默认走 nearest 插值不会凭空生成新类别。Normalize用的是 ImageNet 的 mean/std因为我们后续要加载 ImageNet 预训练权重归一化必须和预训练一致这是很多人忽略的细节。Dataset 的写法也要对应两套数据集的目录差异import os import torch from torch.utils.data import Dataset from PIL import Image import numpy as np class SegmentationDataset(Dataset): def __init__(self, root, splittrain, transformsNone): self.transforms transforms self.split split self.is_voc VOC in root if self.is_voc: self.img_dir os.path.join(root, JPEGImages) self.mask_dir os.path.join(root, SegmentationClass) with open(os.path.join(root, ImageSets, Segmentation, split .txt)) as f: self.ids [line.strip() for line in f] else: self.img_dir os.path.join(root, leftImg8bit, split) self.mask_dir os.path.join(root, gtFine, split) # Cityscapes 的图片按城市分子目录需要递归收集 self.ids [] for city in os.listdir(self.img_dir): for name in os.listdir(os.path.join(self.img_dir, city)): if name.endswith(.png): self.ids.append(os.path.join(city, name[:-4])) def __getitem__(self, i): img Image.open(os.path.join(self.img_dir, self.ids[i] .jpg if self.is_voc else self.ids[i] .png)).convert(RGB) mask Image.open(os.path.join(self.mask_dir, self.ids[i] .png)) if self.transforms: t self.transforms(imageimg, maskmask) img, mask t[image], t[mask] return img, torch.as_tensor(np.array(mask), dtypetorch.long) def __len__(self): return len(self.ids)逻辑说明VOC 的图片是 jpgCityscapes 是 png而且目录按城市分层所以取图路径要区分。mask 用 PIL 读出来后转成torch.long因为交叉熵要求 target 是长整型。很多人拿到源码跑 VOC 没问题、换 Cityscapes 就报错多半就是目录递归这段没适配。3. 从零搭建 DeepLabv3Backbone、ASPP 与 Decoder 逐层拆开3.1 结构拆解空洞卷积和 ASPP 为什么能治多尺度DeepLabv3 的核心思路一句话讲完用带空洞卷积的 ResNet 提取特征保持输出步长output stride在 16再用 ASPP 模块用不同膨胀率抓多尺度上下文最后用 decoder 把低层细节特征融合回来。空洞卷积的意义在于不降低分辨率的前提下扩大感受野这对分割这种逐像素任务是决定性的——它让高层特征图保持足够的分辨率小目标不会因为反复池化直接消失。ASPPAtrous Spatial Pyramid Pooling是这套模型的灵魂。它把同一份高层特征分别送入 3 个膨胀率不同的空洞卷积常见 6、12、18和一个全局平均池化分支最后拼起来再投影。膨胀率 6 看局部纹理18 看大范围上下文全局池化补充整体语义这样一张图里的大楼和小路灯都能被覆盖。decoder 部分则把 stride 8 的低层特征用 1x1 卷积降到 48 通道和 ASPP 输出上采样后拼接补回边界细节这是 DeepLabv3 没有、v3 加上去的关键。Backbone 怎么选是第一个需要拍板的参数。GPU 上默认 ResNet101mIoU 比 ResNet50 高两三个点显存吃紧就换 ResNet50边缘设备考虑 MobileNetV2。VOC 这种物体尺寸跨度大的场景ResNet101 更稳Cityscapes 街景里小目标多通道数大一点对恢复细节有明显帮助。新手直接 ResNet101 起步别在 backbone 上省。3.2 Pytorch 模型代码ASPP 和 Decoder 的最小完整实现先写 ASPP 模块这是整个模型里最值得手写一遍的部分import torch import torch.nn as nn import torch.nn.functional as F class ASPP(nn.Module): def __init__(self, in_ch2048, out_ch256, rates(6, 12, 18)): super().__init__() self.convs nn.ModuleList([ nn.Conv2d(in_ch, out_ch, 1, biasFalse), nn.Conv2d(in_ch, out_ch, 3, paddingrates[0], dilationrates[0], biasFalse), nn.Conv2d(in_ch, out_ch, 3, paddingrates[1], dilationrates[1], biasFalse), nn.Conv2d(in_ch, out_ch, 3, paddingrates[2], dilationrates[2], biasFalse), ]) self.pool nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Conv2d(in_ch, out_ch, 1, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue)) self.project nn.Sequential( nn.Conv2d(out_ch * 5, out_ch, 1, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue)) def forward(self, x): size x.shape[-2:] feats [conv(x) for conv in self.convs] feats.append(F.interpolate(self.pool(x), sizesize, modebilinear, align_cornersFalse)) return self.project(torch.cat(feats, dim1))参数说明rates(6, 12, 18)是工程里最常用的三档膨胀率分别对应小中大三档感受野paddingdilation保证卷积输出尺寸不变。out_ch256统一各分支通道数再拼接避免后续卷积的通道数爆炸。全局池化分支必须插值回输入尺寸才能 concat这里用align_cornersFalse的 bilinear 是 Pytorch 官方语义分割模型的标准做法和用True在最终结果上差零点几个点但保持对齐更稳。完整的 DeepLabv3 主体结构from torchvision.models import resnet50, resnet101 class DeepLabV3Plus(nn.Module): def __init__(self, n_classes21, backboneresnet101): super().__init__() net resnet101(weightsNone, replace_stride_with_dilation[False, False, True]) self.layer0 nn.Sequential(net.conv1, net.bn1, net.relu, net.maxpool) self.layer1 net.layer1 # stride 4 self.layer2 net.layer2 # stride 8低层特征在这里取 self.layer3 net.layer3 # stride 16 self.layer4 net.layer4 # stride 16空洞版 self.aspp ASPP(2048) self.low_proj nn.Sequential( nn.Conv2d(512, 48, 1, biasFalse), nn.BatchNorm2d(48), nn.ReLU(inplaceTrue)) self.decoder nn.Sequential( nn.Conv2d(304, 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, n_classes, 1) def forward(self, x): size x.shape[-2:] x self.layer0(x) x self.layer1(x) low self.layer2(x) # stride 8 的低层特征 x self.layer3(low) x self.layer4(x) # stride 16含空洞 x self.aspp(x) x F.interpolate(x, sizelow.shape[-2:], modebilinear, align_cornersFalse) x torch.cat([x, self.low_proj(low)], dim1) # 256 48 304 x self.decoder(x) x self.classifier(x) return F.interpolate(x, sizesize, modebilinear, align_cornersFalse)逻辑说明replace_stride_with_dilation[False, False, True]是 torchvision 自带参数作用是把 layer4 的 stride 2 换成膨胀卷积让整条特征下采样停在 16 倍而不是 32 倍这是 DeepLab 系列压缩计算量的关键。low_proj把 512 通道的 layer2 特征压到 48避免 decoder 里大通道拼接拖慢训练。assp 输出与low_proj拼接后是 304 通道两个 3x3 卷积负责融合最后是 1x1 分类头。返回时再插值回原图尺寸这样 loss 直接和 label 对齐不用在训练循环里额外 upsample。3.3 预训练权重加载与多卡封装别再让 BN 从头训到尾分割模型很少从随机初始化训ImageNet 预训练的 ResNet 能显著加速收敛、把验证集 mIoU 拉高好几个点。加载方式用strictFalse因为我们的模型没有分类头import torchvision.models as models def build_model(n_classes21, backboneresnet101): model DeepLabV3Plus(n_classesn_classes, backbonebackbone) if backbone resnet101: state models.resnet101(weightsmodels.ResNet101_Weights.IMAGENET1K_V2).state_dict() else: state models.resnet50(weightsmodels.ResNet50_Weights.IMAGENET1K_V2).state_dict() missing model.load_state_dict(state, strictFalse) print(missing keys:, missing.missing_keys) # 只应看到 aspp/decoder/classifier return model说明strictFalse允许缺失键存在缺失的必然是随机初始化的 ASPP 和 decoder 部分这是正常的不需要恐慌。打印missing_keys确认一下如果 ResNet 的卷积层也出现在缺失列表里说明 backbone 名字没对上常见于权重键带module.前缀的情况。多卡训练时如果用nn.DataParallel或 DDP 包了模型保存的 state_dict 键名会多出module.之后单卡推理加载要在 load 前自己 strip 掉或者保存时直接存model.module.state_dict()这个细节放到避坑章细说。4. 训练配置损失函数、学习率策略和 mIoU 评估的完整闭环4.1 损失函数怎么选交叉熵、ignore_index 与 OHEM 的取舍语义分割的默认损失就是逐像素交叉熵VOC 和 Cityscapes 都能直接用。真正要处理的是两个问题类别不平衡和忽略区域。Cityscapes 里道路、建筑这类大区域占像素绝大多数而交通灯、摩托车这种小目标在损失里几乎没存在感VOC 虽然类别比例相对均衡但边缘坏标签也不少。常见做法是加 OHEM在线难例挖掘只回传损失最高的前一部分像素的梯度import torch.nn.functional as F def seg_loss(logits, mask, ignore_index255, ohem_ratioNone): loss F.cross_entropy(logits, mask, ignore_indexignore_index, reductionnone) if ohem_ratio is None: return loss.mean() n_valid (mask ! ignore_index).sum() k int(n_valid * ohem_ratio) flat loss.view(-1) top_k, _ flat.topk(k) return top_k.mean()参数说明ignore_index255让 Cityscapes 的忽略区域完全不参与梯度VOC 没有这类标记但传了也不影响。ohem_ratio0.25是常用值表示只回传损失最高的 25% 像素其余像素梯度为零这样模型注意力会被迫放在难区分的物体边缘和罕见类别上。注意 OHEM 要在reductionnone之后做不能先mean再挑。如果你的训练 loss 曲线掉得很快但验证集 mIoU 纹丝不动第一个要怀疑的就是 loss 被大量易分类背景像素稀释了。4.2 VOC 和 Cityscapes 的训练超参照抄就能跑的参数表两套数据集分辨率、类别数、样本量差异大超参不能一套通吃。下表是我实测下来比较稳的一组起点VOC 用了 SBD 增强数据Cityscapes 用官方 gtFine超参数VOC 2012含 SBDCityscapes输入尺寸513x513 随机裁剪769x769 随机裁剪batch size168总迭代约 80 epoch约 500 epoch或 80k 步初始学习率0.010.01学习率策略polypower0.9polypower0.9优化器SGDmomentum0.9wd1e-4SGDmomentum0.9wd1e-4warmup前 1 epoch前 1 epoch学习率策略用 poly 而不是等间隔下降这是 DeepLab 系的关键传统。poly 的含义是学习率随训练进度按(1 - step/total_steps)^power平滑衰减后期微调更细腻def poly_lr(base_lr, step, total_steps, power0.9): return base_lr * (1 - step / total_steps) ** power完整训练循环加上混合精度避免显存成为瓶颈scaler torch.cuda.amp.GradScaler() optimizer torch.optim.SGD(model.parameters(), lr0.01, momentum0.9, weight_decay1e-4) total_steps len(train_loader) * epochs for epoch in range(start_epoch, epochs): model.train() for i, (images, masks) in enumerate(train_loader): step epoch * len(train_loader) i lr poly_lr(0.01, step, total_steps, 0.9) for g in optimizer.param_groups: g[lr] lr images, masks images.cuda(), masks.cuda() optimizer.zero_grad() with torch.cuda.amp.autocast(): logits model(images) loss seg_loss(logits, masks) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()参数说明每步手动更新param_groups里的学习率比scheduler更直观可控poly 曲线也能精确复现。torch.cuda.amp.autocast配合GradScaler是 Pytorch 官方的混合精度方案前向用 FP16、主权重保持 FP32显存基本省一半Cityscapes 的 769 输入在 16G 卡上也能跑。注意scaler.step和scaler.update的顺序不能换否则梯度会失真。如果你连 batch size 8 都 OOM把scaler去掉、再不行就把输入降到 641 或换 ResNet50别硬上。不嫌工程重的话也可以直接用 mmsegmentation 训练 Cityscapes它把 DeepLabv3 的全套配置封装好了VOC 和 Cityscapes 都内置了数据接口但自己用 Pytorch 写训练循环才能把 OHEM、梯度累积、自定义损失改成顺手的样子这份源码的意义就在这里。4.3 评估指标像素准确率会骗人用 mIoU 说话训练完不能只看 loss分割任务的标准指标是 mIoUmean Intersection over Union也就是每个类别的交并比取平均。像素准确率在这种任务里毫无参考价值因为背景占大头哪怕把所有像素都预测成道路准确率也能到 80% 以上。mIoU 对每个类单独算小类别拉垮一眼就能看出来import numpy as np def compute_miou(pred, mask, n_classes, ignore_index255): pred pred.flatten() mask mask.flatten() valid mask ! ignore_index pred, mask pred[valid], mask[valid] confusion np.zeros((n_classes, n_classes), dtypenp.int64) np.add.at(confusion, (mask, pred), 1) inter np.diag(confusion) union confusion.sum(1) confusion.sum(0) - inter iou inter / np.maximum(union, 1) return iou.mean(), iou逻辑说明np.add.at(confusion, (mask, pred), 1)是在混淆矩阵里按「真实类、预测类」累加比循环快得多也规避了多维索引的坑。n_valid过滤掉 ignore 区域的像素。每个类别的 IoU 单独输出后你才会看到道路 IoU 可能 0.95而交通灯只有 0.3这个差距才是调参的方向。验证集评估时模型输出要插值回原图尺寸再计算别拿 stride 16 的特征直接对齐标签。5. DeepLabv3 训练避坑五个最常见的踩坑记录和对应解法5.1 CUDA out of memory 报错显存不够先动这三处现象刚把 Cityscapes 数据加载器跑起来第一轮迭代就报CUDA out of memory或者训练到一半崩掉。原因分割模型的前向要保留整张特征图用于反向传播输入尺寸每放大一点中间张量体积是平方级上涨。769x769 的输入加 ResNet10111G 显存真的不够。解决按顺序做三件事。第一把torch.cuda.amp.autocast和GradScaler开起来省 40% 显存第二把 batch size 从 16 降到 8、4配合梯度累积第三把随机裁剪尺寸从 769 降到 641。梯度累积的写法是在scaler.step(optimizer)前判断当前 batch 是否凑够了累积步数达不到就scaler.update()后继续下一批效果等于把 batch size 翻倍。如果这三招用尽还 OOM就该换 ResNet50 backbone 了别跟显存硬刚。5.2 训练 loss 一直降验证 mIoU 却在 40% 附近不动现象训练集 loss 从 2.3 降到 0.2看着很健康但每隔几个 epoch 跑一次验证集mIoU 就是上不去。原因最常见的是验证集用了和训练集相同的随机裁剪把大量物体截断了或者验证时Resize把宽高比拉变了标签对不上原图语义。另一个高频原因是数据里有坏标签——SegmentationClass 某些 png 的类别 id 越界loss 里 ignore 没拦住模型在学垃圾。解决验证集固定用A.Compose([A.Resize(crop_size, crop_size), ...])且不换随机种子保证每次验证的是同一批变换跑完一个 epoch 后从验证集随机抽 4~8 张图把预测 mask 叠在原图上可视化肉眼看语义对齐不对齐。这一步比 mIoU 数字更能暴露问题。标签越界的排查是打印torch.unique(mask)发现最大 id 大于等于类别数就人工剔除这几张训练图。5.3 把 VOC 的 21 类和 Cityscapes 的 19 类混在一个训练流程里现象在 VOC 上训得好好的权重换 Cityscapes 数据集继续训练loss 直接变成 nan或者验证 mIoU 跌到个位数。原因VOC 的索引 0 是背景Cityscapes 的索引 0 是道路VOC 没有 ignore 区域Cityscapes 的 255 如果没特判会把-1/255当成类别。两个数据集的语义编号体系完全不同复用同一个分类头和损失函数但没改类别数等于让模型把道路当背景学。解决训练前各写一份类别映射表VOC 用0~20Cityscapes 用0~18ignore 统一处理成 255损失函数的n_classes参数跟随当前数据集。最稳妥的流程是 VOC 和 Cityscapes 分两个独立 Experiment数据集、类别数、分类头各管各的不共享权重文件。如果要做跨数据集的 finetune至少把分类头重新随机初始化绝不能在两个类别体系间直接拷贝。5.4 断点续训翻车checkpoint 里只存了模型权重现象训练中断后重新启动load 了上次的model.pt发现 loss 回到初始值甚至发散前面的训练白跑了。原因只存了model.state_dict()优化器的 momentum、混合精度的 loss scale、当前 epoch 和学习率步数全丢了。SGD 的 momentum 状态丢了以后恢复出来的优化器等于从零开始学习率又从 base_lr 跑刚开始的几百步很容易把已经收敛的权重冲散。解决checkpoint 里把五样东西一起存缺一不可torch.save({ model: model.module.state_dict() if hasattr(model, module) else model.state_dict(), optimizer: optimizer.state_dict(), scaler: scaler.state_dict(), epoch: epoch, miou: best_miou, }, checkpoint.pth)恢复训练时注意两点加载模型权重后先调optimizer.load_state_dict再算total_steps时要把start_epoch传进去否则 poly 学习率会从头衰减相当于给模型突然加了一剂猛药。多卡训练保存时用model.module.state_dict()加载时就无需处理module.前缀这是最简单的防呆方案。5.5 loss 变成 nan三个最隐蔽的原因排查顺序现象训练到第 3~7 个 epochloss 突然变nan然后一路 nan 下去模型输出全是同一种类别。原因概率最高的三个——数据里出现非法标签比如 mask 里有 255 但ignore_index没设混合精度下GradScaler的 loss scale 刷新不及时梯度溢出学习率本身太大SGD 权重更新在 BN 层上炸了。解决按顺序做三件事。先print(torch.unique(mask))检查训练标签确认没有255混进来且最大 id 合法再打开torch.autograd.set_detect_anomaly(True)跑 50 步看 nan 是出现在前向还是反向最后把 base_lr 降到 0.001 试跑 1 个 epoch如果正常了就说明 0.01 对当前 batch size 偏大加 warmup 或降回 0.007。我习惯在训练脚本里写一个nan_guard每 100 步检查 loss 是否有isfinite发现 nan 就自动把学习率减半并跳到上一个 checkpoint而不是直接放弃整轮训练。6. 验证与部署从可视化掩码到 ONNX 导出的最后一步6.1 单图推理与彩色掩码可视化快速判断模型有没有学歪训练完成后第一步是写一个单图推理函数把预测结果叠在图上肉眼检查。mIoU 只是一个数字模型有没有把远处的车识别成人、把天空当成墙只有可视化能看到。VOC 和 Cityscapes 各有一套官方的调色板预测后按像素类别上色def predict_one(model, img_path, devicecuda, n_classes19): model.eval() img Image.open(img_path).convert(RGB) from PIL import Image import torchvision.transforms.functional as TF tensor TF.to_tensor(TF.resize(img, (769, 769))).unsqueeze(0).to(device) tensor TF.normalize(tensor, mean(0.485, 0.456, 0.406), std(0.229, 0.224, 0.225)) with torch.no_grad(): logits model(tensor) return logits.argmax(dim1).squeeze(0).cpu().numpy()说明TF.resize默认最近邻对输入图没问题argmax(dim1)取每个像素概率最大的类得到的是 (769, 769) 的整数掩码。可视化时把掩码映射到颜色表叠加到原图上重点看物体边缘是否贴合、小目标有没有断成两截。6.2 用每类 IoU 而不是平均 mIoU 验证效果小目标类别怎么补验证做到类别粒度才能真正指导下一步。我的习惯是训练结束后跑一次全量验证集把每类 IoU 打印出来排序最低的三个类就是下阶段的优化对象。Cityscapes 上通常是摩托车、交通灯、自行车这类小目标垫底因为它们像素少、标注边缘噪声大、在 769 的裁剪里经常被切一半。对应的补法有几种一是对小目标类别在 loss 里加权重比如按「类别像素占比的倒数」加权交叉熵二是把训练时的随机裁剪从中心式改成偏多采样小目标——先统计哪些图里交通灯面积占比高训练时提高这些图的采样权重三是输入尺寸换到 1024 并配合梯度累积小目标分辨率上来了IoU 会有看得见的提升。切记每改一次只看 mIoU 全局数字那是自欺欺人必须滚动看每类 IoU 变化。6.3 导出 ONNX 并用 onnxruntime 验证给 CPU 和嵌入式设备留后路训练验证全部通过后导出 ONNX 是工程收尾的标准动作。Pytorch 转 ONNX 的坑主要是插值算子和动态尺寸DeepLabv3 里用的 bilinear align_cornersFalse在 opset 11 的支持非常成熟不需要特殊处理model.eval() dummy torch.randn(1, 3, 769, 769).cuda() torch.onnx.export(model, dummy, deeplabv3plus.onnx, input_names[input], output_names[output], opset_version11, do_constant_foldingTrue)导出后用 onnxruntime 跑一遍数值对齐确认 Pytorch 和 ONNX 的输出误差在可接受范围import onnxruntime as ort import numpy as np ort_sess ort.InferenceSession(deeplabv3plus.onnx, providers[CPUExecutionProvider]) with torch.no_grad(): onnx_out ort_sess.run(None, {input: dummy.cpu().numpy()})[0] torch_out model(dummy).cpu().numpy() print(max abs diff:, np.abs(onnx_out - torch_out).max())说明do_constant_foldingTrue会把 BN 和卷积折叠成单个算子运行时更快误差打印出来的最大值只要小于 1e-3 就可以认为导出正确。如果之后要接实时视频流建议把输入尺寸固定避免动态尺寸拉起额外耗时要支持任意分辨率就把dummy换成带-1维度的动态 shape同时加上dynamic_axes{input: {2: h, 3: w}, output: {2: h, 3: w}}。ONNX 出来后 CPU 推理、TensorRT、OpenVINO 都是一条路模型才算真正成了能交付的东西。我自己第一次在 Cityscapes 上折腾 DeepLabv3是拿 1024x1024 输入直接开的训两张卡瞬间 OOM折腾大半晚才意识到先写个小规模冒烟测试再跑全量数据。后来每次动手做分割训练都先把数据加载器和损失函数单拎出来跑 10 张图确认 loss 在掉、mask 颜色对得上才敢扔全量数据然后开自动保存。这份流程踩过的坑都写在上面的避坑章里了值不值得做、投入多大算力照着这份流程跑完一遍你心里会有数。希望帮到你。本文还有配套的精品资源点击获取
返回列表