去年接手一个边缘设备上的实时分割任务要在车载级别的芯片上处理 1080p 输入精度还不能比主流方法差太多。翻了大半圈论文最后选定了 ArXiv 2022 上的 PIDNet 作为基线。为什么是它因为 PIDNet 把自动控制里的 PID 控制器思想直接搬进了语义分割网络设计里用三个分支分别处理细节、上下文和边界在高分辨率城市场景数据集上做到了精度和速度的平衡。这篇文章不聊论文复述只讲我实际配置环境、跑通训练、评估部署以及在这套模型上做轻量化改进创新的完整过程。无论你是刚入门的同学还是准备拿 PIDNet 做基线刷点的工程师后面这些内容都能直接参考。1. PIDNet 到底在解决什么问题1.1 PID 控制器怎么变成了一个网络结构传统语义分割网络最典型的是编码器-解码器结构低层分支保细节高层分支保语义最后通过特征融合把两类信息拼在一起。这种结构的问题在于细节分支感受野小、容易受纹理干扰语义分支分辨率低、物体边缘容易糊成一团。PIDNet 的出发点非常直观把经典 PID 控制器里的比例、积分、微分三项映射到三个分支。P 分支负责保留高频细节类似比例项看到的是“现在正在发生什么”提供锐利的空间结构信息。I 分支通过大感受野聚合上下文类似积分项看到的是“过去一段时间的累积状态”负责消除噪声和歧义。D 分支专门提取物体边界类似微分项关注的是信号变化的剧烈程度用来补强边缘响应。三个分支最后通过融合模块做加权聚合而不是简单把特征拼接起来。实际效果就是在 Cityscapes 这类高分辨率数据集上PIDNet 在保持实时推理速度的同时mIoU 能达到一个很有竞争力的水平。官方报告的单卡 2080 Ti 上测速数据我记得在 90 FPS 上下具体数值和你输入分辨率、TensorRT 优化都有关系但整体确实属于“能实时跑”的模型。1.2 复现之前的三个准备数据、显卡、心理预期很多人上来就直接拉代码、装环境然后跑训练跑到一半发现数据集组织不对或者显卡爆显存再回头折腾。我的建议是先做三件事。第一确认自己的数据源。PIDNet 官方训练指标基本都是在 Cityscapes 上测的你至少要拿到该数据集的训练集和验证集。如果你只有自己的业务数据那需要提前把标注格式对齐成 19 类城市场景的模式或者重新设计分类器头这个不是简单换个数据加载路径就能解决的问题。第二确认显存和迭代规模。PIDNet 不算重模型但 1024x2048 级别的输入对显存仍然有压力。我用单张 11GB 的卡试过 batch size 2训练比较勉强换成 24GB 的卡可以放到 batch size 4 到 6。如果你只有 8GB 显存最好把 crop size 调小或者直接用梯度累积模拟大 batch不要硬上大分辨率。第三明确复现目标。是要跑出官方论文指标还是只要在自己的数据上跑通流程如果是前者加载官方提供的预训练权重再做少量迭代验证是比较快的路径如果是后者那你必须接受从头训练可能带来的指标波动。PIDNet 的配置文件和训练次数对学习率非常敏感不是随便跑几个 epoch 就能刷出高分的。2. 环境配置把官方代码真正跑起来的步骤2.1 拉代码和创建虚拟环境官方代码库的结构和 mmsegmentation 这类工程框架非常接近目录里主要是 configs、tools、model 等模块并且使用 PyTorch 的分布式训练封装。因为代码发布较早它对最新版 PyTorch 的兼容性并不好所以不要直接装一个最新的 PyTorch 2.x 就开跑除非你做好了自己修算子的准备。我实测比较稳的一套组合是 Python 3.8、CUDA 11.3、PyTorch 1.10.1、torchvision 0.11.1。命令如下conda create -n pidnet python3.8 -y conda activate pidnet pip install torch1.10.1cu113 torchvision0.11.1cu113 --extra-index-url https://download.pytorch.org/whl/cu113进入项目目录后再装剩余依赖。如果官方提供 requirements.txt直接执行即可没有的话手动装这几项mmcv 相关基础库、tensorboard、yaml、pyyaml、rich、tqdm。这里最怕的是 mmcv 版本和 PyTorch 版本对不上我建议先用官方指定版本不要盲目升级。pip install -r requirements.txt python -c import torch; print(torch.__version__, torch.cuda.is_available())2.2 第一次跑代码前要处理的几个编译坑代码里如果有自定义 CUDA 算子一般会提供一个 compile 脚本或者依赖 mmcv 的编译链路。在跑训练之前先花五分钟把算子编译掉不然等你启动训练时才发现 import 失败会浪费很长时间。其次如果项目里有 apex 相关的依赖我建议直接把它卸掉改用 PyTorch 自带的混合精度接口。老项目经常出现 apex 和最新显卡驱动不兼容的情况表现是训练到某个 iteration 突然报错或者 loss 变成 nan。还有就是同步批归一化。PIDNet 这类实时模型理论上单卡也能跑但代码里默认可能开了 SyncBN。如果只用单卡训练SyncBN 退化成普通 BN不会报错可一旦开多卡某些卡显存不足或 GPU 利用率不均衡的问题就会被放大。建议在单卡调试阶段先把分布式相关配置注释掉确认前向和反向都能跑通再谈多卡。提示拿一张土测试卡先把训练脚本完整跑 20 个 iteration比直接上全量训练更省时间。环境坑基本都是在这个阶段暴露的。2.3 准备数据目录和标签格式Cityscapes 数据下载后原始目录是这样组织的data/ ├── leftImg8bit/ │ ├── train/ │ └── val/ └── gtFine/ ├── train/ └── val/官方代码一般会要求你把数据软链接到项目下的 data 目录。标签不是直接使用 RGB 图像而是用 gtFine 下的 labelIds 语义标签。Cityscapes 本身的像素级标注有 30 多类但模型训练通常只用其中 19 类参与评测和 loss 计算其余类别会被归到 ignore 索引 255。我犯过的错误是把所有类别都当成有效类别去算 loss结果 mIoU 直接掉到底。正确的做法是在配置里明确num_classes19和ignore_index255同时检查数据增强阶段是否做了 class map 映射不要直接用原始 34 类的 id。3. 动手训练一轮从配置到实际迭代记录3.1 配置文件里必须改的四个位置PIDNet 每个实验配置都会指定 backbone 类型、训练分辨率、batch size、学习率、迭代总数等参数。第一次跑不要追求精确复现官方指标先只改下面四个位置。第一数据路径。检查data_root是否指向你的 Cityscapes 目录ann_file是否生成了对应的 train.txt 和 val.txt。很多复现跑不起来不是模型问题是文件列表里的路径写错了。第二分辨率。官方训练分辨率通常是 1024x512 的 crop验证时是 1024x2048 全图。如果你的显卡显存不够把训练分辨率降到 896x448验证分辨率保持在 2048 沿上打这样精度损失不会太大但训练速度会快很多。第三学习率。PIDNet 使用 poly 学习率策略基础学习率一般配 0.02 或 0.01配合 warmup。如果你把 batch size 从 8 改成 2学习率最好跟着降否则容易在开头几千个 iteration 里出现 loss 震荡。第四迭代总数和验证间隔。官方从零训练大概需要 120K 到 160K 个 iteration但在复现阶段你可以先设 5K 个 iteration 验证流程再把总数拉回去不用一上来就 120K 起步。3.2 云上的代跑训练没有本地卡怎么把迭代跑完很多人找“代跑训练”其实是因为自己没有合适显卡或者本地环境不干净想直接在云端把完整实验跑完。我的经验是与其把日志给别人再看结果不如自己掌握一套远程跑实验的最小操作流程。第一步租一台带有 NVIDIA GPU 的云实例显存建议 24GB 起步至少 12GB。实例里装好驱动和 CUDA 工具包再按上一节方式创建 conda 环境。第二步把代码、数据和预训练权重传上去。这一步容易漏的是预训练权重PIDNet 的骨干网络需要 ImageNet 预训练没有它前面几千个 iteration 的 loss 会下降得非常慢甚至直接发散。第三步用 nohup 或 tmux 把训练放到后台确保断开 SSH 后训练不中断。这里我推荐 tmux因为它可以随时切回去看实时日志比 nohup 更直观。tmux new -s pidnet_train CUDA_VISIBLE_DEVICES0 python tools/train.py --cfg configs/pidnet_s_cityscapes.py训练启动后第一件事不是等着看 mIoU而是先观察 loss 曲线和 GPU 利用率。如果 GPU 利用率低于 90%问题多半出在数据加载上检查 DataLoader 的 num_workers 和 prefetch 配置。如果 loss 在下降说明环境没问题可以让它跑着每隔一段时间打开验证集做评测。3.3 单卡、多卡和混合精度的实际操作单卡训练命令最简单直接设置CUDA_VISIBLE_DEVICES0启动。多卡训练使用 PyTorch 官方分布式启动器需要注意--master_port不要和别的进程撞了。CUDA_VISIBLE_DEVICES0,1,2,3 python -m torch.distributed.launch \ --nproc_per_node4 \ --master_port29511 \ tools/train.py \ --cfg configs/pidnet_m_cityscapes.py多卡训练里最常见的坑是 batch size 和 lr 的换算关系。如果原来单卡 batch size 是 4四卡训练时总 batch size 变成 16学习率可以按线性缩放调到原来的 2 到 4 倍但不要直接调 4 倍容易爆炸。我在实际实验里一般只放大 2.5 倍到 3 倍再配合更长的 warmup。混合精度方面开启 AMP 后显存占用能降低 20% 到 30%训练速度也会提升。但如果出现了 loss 持续抖动、mIoU 波动大的情况优先检查是不是某些算子在低精度下不稳定把相关层强制回落到 FP32 即可。3.4 断点续训和日志判断长时间训练最怕的是中途断掉。我习惯每隔 5000 个 iteration 保存一个 checkpoint并且把权重和优化器状态都保存下来。这样即使训练中断也能从上一个 checkpoint 继续恢复不需要重新从头跑。日志阶段需要重点看两个信号。一个是 loss 下降曲线是不是平滑如果出现反复的尖峰说明学习率偏高或数据增强策略过于激进。另一个是验证集 mIoU 是否还在上升如果连续 10 个验证周期没有提升可以提前停止训练或者切换到更小的学习率做微调。4. 评估与部署别只盯着 mIoU 这一列数字4.1 验证脚本和速度测试的正确姿势训练结束后评估验证集指标通常是调用tools/test.py或者tools/eval.py。这里有个很容易被忽略的点验证时会默认开启多尺度测试也就是对输入图片做多种尺度的缩放再把结果融合。多尺度测试确实会提升 mIoU但它和实时部署没有任何关系因为实际推理时不可能做多尺度输入。如果你要用 PIDNet 做边缘部署一定要单独跑一遍单尺度、batch size 为 1 的速度测试。记录稳定速度之前先让模型跑 10 次以上预热确保 GPU 已经进入最高频率状态否则你测出来的 FPS 会偏低 20% 以上。4.2 mIoU、类别 IoU 和速度的权衡评估报告里最容易骗人的就是平均 mIoU。Cityscapes 的 19 个类别里有的类别样本极多但容易识别比如道路、建筑有的类别样本少且形态多变比如骑车人、卡车。平均 mIoU 高不代表所有类别都可用。我在部署前的习惯是单看三个关键类别行人、骑车人、上方的交通标志。这三个类别对边缘设备来说是最难搞的如果它们在单尺度推理下 IoU 低于 60%那实际产品体验基本不可用。实时语义分割部署还需要关注参数量和计算量。PIDNet 系列里 S 版本和 M 版本参数量差异不大但速度差距明显因为高分辨率分支的通道数不同。如果你的目标是 1080p 实时我建议直接选 S 版本起步不要一上来就用大版本再去做剪枝后面改进空间会更大。4.3 导出到轻量推理后端的几个注意点常见的部署流程是把 PyTorch 模型导出成 ONNX再转成对应的推理引擎。PIDNet 结构里有几个自定义融合模块导出时最容易遇到的问题是动态 shape 支持不佳。第一个建议是固定输入尺寸。实时部署通常都固定分辨率导出时直接把动态维度锁死很多算子兼容性问题会消失。第二个建议是检查融合模块里的某些重整形操作。这类操作在 PyTorch 中顺理成章但转成 ONNX 后可能会产生低效的拷贝算子导致推理速度反而比 PyTorch 慢。遇到这种情况不要全局改只对瓶颈子图做重写即可。第三个建议是量化。INT8 量化的 mIoU 下降通常在 1 到 3 个点左右如果下降超过 5 个点优先检查边界分支因为边界分支的输出对量化误差非常敏感建议保留为 FP16 或 FP32。5. 模型改进创新从换主干到加注意力模块5.1 轻量化主干替换和通道裁剪PIDNet 的原生主干在精度和速度上已经平衡得不错但如果要跑到更低功耗的设备上最常见的做法是更换更轻量的骨干网络比如 MobileNetV3 或者轻量自注意力结构。替换主干最麻烦的是对齐下采样倍数和通道数。PIDNet 的三个分支各自在不同层提取特征新主干的 stem 和 stage 输出步长必须能和原分支接上。我的建议是保留原来 PIDNet 的融合模块和头部分支只替换主干部分这样改造成本最小。通道裁剪则是另一个思路。原模型的 P 分支和 D 分支在推理时并不是所有通道都同等重要通过通道稀疏化计算找到一个合理的剪枝比例把冗余通道裁掉后做几轮微调。这个过程需要结合速度测试不要只看参数量减小了实际 FPS 可能没变化。5.2 给中间分支加注意力模块不加选择地到处加注意力是很多改进工作掉进坑里的原因。PIDNet 最大的特点是三分支结构每个分支承担的任务不同所以注意力加的时机和位置很关键。我的经验是在 I 分支融合上下文后加坐标注意力因为在低分辨率语义分支中长距离依赖的建模对消除歧义最有帮助。坐标注意力的优势是轻量还能保留空间位置信息比起直接上全局自注意力更高效安全。class CoordAtt(nn.Module): def __init__(self, inp, oup, reduction32): super().__init__() self.pool_h nn.AdaptiveAvgPool2d((None, 1)) self.pool_w nn.AdaptiveAvgPool2d((1, None)) self.conv1 nn.Conv2d(inp, inp // reduction, kernel_size1) self.conv_h nn.Conv2d(inp // reduction, oup, kernel_size1) self.conv_w nn.Conv2d(inp // reduction, oup, kernel_size1) def forward(self, x): b, c, h, w x.size() x_h self.pool_h(x) x_w self.pool_w(x).permute(0, 1, 3, 2) y torch.cat([x_h, x_w], dim2) y self.conv1(y) # 激活后拆回两个方向再乘回原特征 return x加完注意力之后必须立刻做一次速度回归测试。很多注意力模块虽然加了几个点 mIoU但推理耗时增加超过 15%这对于实时模型来说基本等于白加。5.3 边界损失增强让 D 分支真正发挥作用PIDNet 的 D 分支本身提取边界但很多人训练时只是隐式地让边界信息参与融合没有显式计算边界损失。改进方向之一就是在训练阶段给 D 分支加一个边界监督信号。边界标签可以从语义分割标签里的类别交界处求梯度得到。先对 one-hot 标签做拉普拉斯滤波取绝对值大于阈值的像素当成边界区域。然后让 D 分支通过 1x1 卷积输出单通道响应用带权重的二元交叉熵计算损失。整体损失可以写成total_loss seg_loss 0.3 * boundary_loss这个 0.3 的权重不是固定的。边界损失太大会导致 D 分支只顾边缘、不看整体结构太小则起不到监督作用。先小规模试 0.1 到 0.5 之间的几个值观察验证集边界相关类别的 IoU 变化再确定权重。5.4 蒸馏、微调与消融实验设计如果你想得到一个比 S 版本更小的模型但又不愿意降低太多精度我建议用蒸馏路线先训练一个大的 PIDNet M 版本再让 S 版本去学它的软化输出。这里有两个细节比较重要。第一个是温度系数。蒸馏温度太高模型会过于关注模糊的类别关系反而忽略真实标签温度太低蒸馏退化成普通交叉熵。我一般从 4 开始调。第二个是只蒸馏语义分支不蒸馏边界分支。边界分支本身噪声比较大强行蒸馏会把不稳定的边缘响应传给小模型最终边界质量反而变差。所有的改进实验建议用下面这张表来记录不然做几个版本后你就不知道该往哪个方向走了。实验版本改动内容mIoU单尺度FPS参数量PIDNet-S 基线无77.5938.1M加坐标注意力I分支后加CoordAtt78.3888.4M加边界损失D分支显式监督78.9918.1M轻量主干MobileNetV3替换Backbone74.81264.3M蒸馏M蒸馏到S78.61125.2M这个表格是模拟数据但思路足够说明问题。每一项改进都要同时看精度和速度只追求 mIoU 上涨没有意义。6. 问题排查与避坑实录6.1 训练阶段常见问题现象可能原因解决办法前几个iteration loss直接变成nan学习率过大 / 混合精度含非稳定算子降低初始学习率关闭AMP训练日志mIoU长期不涨未加载ImageNet预训练检查pretrained路径多卡训练GPU利用率不均衡DataLoader num_workers太小提高workers开启prefetch验证时图像全黑label id映射不正确检查class map和ignore_indexloss快速下降到0但mIoU很低类不平衡被忽略模型一直在学背景检查数据加载的类别重映射最让我印象深刻的是一次 loss 稳步下降但验证集 mIoU 只有 20 的情况。排查了一整天才发现是数据增强环节把标签图片做了线性插值导致 255 这个忽略索引被污染成了 254、253 等数值。这个问题在代码里不会报错只能通过对比增强前后标签的像素分布才能发现。6.2 代跑训练中远程环境问题远程实例上最常遇到的两个问题是中断和权限。训练中断了不要慌只要 checkpoint 保存完整续训时把配置里的--resume参数指向最新权重文件即可。但是一定要确保优化器状态也被保存否则学习率会和迭代次数对不上后续 loss 曲线会出现一个明显的台阶。权限问题主要体现在数据集读取上。有些云实例默认挂载的数据目录是只读的训练过程中如果要写入缓存文件就会在随机的地方报错。我建议把数据软链接到实例的高性能临时盘训练完成后手动同步回存储盘。6.3 推理和部署阶段问题现象可能原因解决办法TensorRT导出速度比PyTorch慢动态shape或低效reshape算子固定输入分辨率重写瓶颈子图INT8量化后边界效果突然恶化边界分支对量化敏感保留边界分支为FP16单尺度FPS比论文低很多没有预热 / 显卡降频固定频率预热20次再测部署后目标边缘出现白边后处理未对边界进行膨胀腐蚀在概率图上做小半径形态学操作线上部署有一点很值得记住模型推理出来的边缘大概率比人工标注的边缘更“毛糙”。不要急着改模型结构先检查是不是概率图阈值定得太高以及有没有把 255 ignore 区域当成了真实背景。很多时候后处理上一个 3x3 的中值滤波就能解决代价是边缘可能稍微变粗需要根据业务需要取舍。7. 一点个人经验收尾我在实际跑完这套完整流程后最深的体会是实时语义分割实验的周期比普通分割长得多因为每一项改进都要同时看精度和速度。如果你也准备拿 PIDNet 做基线我建议先搭好一套完整的“训练 验证 单尺度速度回归”流程再动手加模块。否则很可能出现改进后 mIoU 涨了 1.2但 FPS 掉了 4 帧的情况而你根本不知道瓶颈到底在哪个分支。另外想提一个小技巧复制 PIDNet 实验时不要一开始就在大模型权重上调试新想法先用 S 版本把改进思路验证完再迁移到 M 或 L 版本。S 版本迭代速度快显存占用低一个晚上能跑完多个消融实验对判断思路是否有效更有利。等思路确定了再用大版本拉精度节省的时间非常可观。PIDNet 三分支的设计在实时语义分割方向里是一个很适合做工程改造的底座后续无论是轻量化、蒸馏还是边界增强都还有很大的折腾空间。希望这份手把手的复现过程能帮你少踩几个坑。
