MMSegmentation 中的 STDC 实时语义分割:STDCNet 主干、Detail Aggregation 模块与 Cityscapes 训练配置全解析
MMSegmentation 中的 STDC 实时语义分割STDCNet 主干、Detail Aggregation 模块与 Cityscapes 训练配置全解析【免费下载链接】mmsegmentationOpenMMLab Semantic Segmentation Toolbox and Benchmark.项目地址: https://gitcode.com/GitHub_Trending/mm/mmsegmentationSTDCShort-Term Dense Concatenate Network短时密集拼接网络是一种专为实时语义分割设计的单流高效网络通过去除 BiSeNet 双流结构中的冗余来在精度与速度之间取得平衡。本文以 configs/stdc/README.md 为骨架结合 MMSegmentation 中 STDC 的源码、配置与测试系统讲解 STDCNet 主干的结构设计、STDCHead 的 Detail Aggregation 损失原理、Cityscapes 四套训练配置的参数细节以及官方预训练权重的转换与复现方法帮助读者在 MMSegmentation 中完整落地 STDC 实时分割方案。STDC 算法背景与设计动机STDC 系列算法出自论文Rethinking BiSeNet For Real-time Semantic SegmentationCVPR 2021论文链接见 configs/stdc/README.md 顶部其核心动机直指 BiSeNet 双流结构的两大问题结构冗余导致推理耗时BiSeNet 通过额外引入一条空间信息编码路径Spatial Path来补充细节但这会显著增加计算量STDC 主张以单流方式处理去掉这条冗余路径。分类预训练主干不适配分割任务从图像分类任务借用的主干缺乏针对分割的任务级设计效率不高。STDC 的解决方案分为两步编码端逐步降低特征图通道数并利用各阶段特征图的聚合来表征图像形成 STDC 的基本模块解码端提出 Detail Aggregation细节聚合模块以单流方式把空间信息的学习下沉到低层网络层最后将低层特征与深层特征融合输出最终分割结果。论文在 Cityscapes 与 CamVid 上验证了精度与速度的权衡在 NVIDIA GTX 1080Ti 上Cityscapes 测试集以 250.4 FPS 达到 71.9% mIoU比当时最新方法快 45.2%在更高分辨率输入下以 97.0 FPS 达到 76.8% mIoU。这些数据来自原 README 中引用的论文 Abstract仅作为算法背景参考。仓库内 STDC 相关文件布局MMSegmentation 中 STDC 的实现分布在配置、源码、转换工具与测试四个层面先建立全局文件地图类别路径作用算法文档configs/stdc/README.md算法介绍、用法、结果表与引用模型元数据configs/stdc/metafile.yaml模型集合的注册信息与权重链接训练配置configs/stdc/ 下 4 个.pySTDC1/2 在 Cityscapes 上的 4 套训练配置基础模型配置configs/base/models/stdc.pySTDC 完整网络主干解码头辅助头定义主干源码mmseg/models/backbones/stdc.pySTDCModule、FeatureFusionModule、STDCNet、STDCContextPathNet解码头源码mmseg/models/decode_heads/stdc_head.pySTDCHead及 Detail Aggregation 损失权重转换工具tools/model_converters/stdc2mmseg.py官方 STDC1/2 权重键名转换脚本单元测试tests/test_models/test_backbones/test_stdc.py主干与模块的结构约束验证其中主干源码文档字符串注明了实现来源为官方仓库 MichaelFan01/STDC-Segmmseg/models/backbones/stdc.py 文件头注释代码主体是官方实现的 MMSegmentation 风格改写。核心组件源码解析从 STDCModule 到 STDCContextPathNetSTDCModule短时密集拼接的基本单元STDCModule定义在 mmseg/models/backbones/stdc.py 中是 STDC 网络的基本构建块。其核心思想是逐步降维 密集拼接一个模块内串联多个卷积层每层通道数逐级减半最后把各层输出沿通道维拼接实现低层细节与高层语义的密集融合。从源码看STDCModule的关键构造参数包括in_channels/out_channels输入通道数 / 缩放前的输出通道数stride首层卷积的步长stride2时模块内部自动带下采样分支num_convs模块内卷积层数默认 4且断言必须大于 1fusion_type融合方式仅支持add或catmmseg/models/backbones/stdc.py 中有assert fusion_type in [add, cat]约束。模块内的通道调度规律是第i层输出通道为out_channels // 2**i最后一层为out_channels // 2**(num_convs-1)即逐层减半。以测试 tests/test_models/test_backbones/test_stdc.py 中的x_stdc STDCModule(in_channels32, out_channels32, stride4)为例其layers[0].conv.in_channels 32、layers[3].conv.out_channels 4正好对应四层卷积 32→16→8→4 的降维链。forward中把各层输出torch.cat后再与下采样后的输入相加即短时密集拼接 残差的完整形态。FeatureFusionModule空间细节与语义上下文的融合FeatureFusionModule与 BiSeNetV1 中的同名模块不同它使用两个ConvModule组成self.attention中间通道由scale_factor控制默认 4而 BiSeNetV1 只用一个conv_attenmmseg/models/backbones/stdc.py 源码注释明确说明了这一区别。其前向过程为拼接空间与上下文特征 → 1x1 卷积压缩 → 自适应全局池化 两级 1x1 卷积 Sigmoid 生成注意力权重 → 加权残差输出x_attn x。这一模块在STDCContextPathNet中负责融合浅层空间特征与经过 Attention Refinement Module 的上下文特征。STDCNet两种主干结构与通道配置STDCNet注册名STDCNet通过stdc_type区分两种主干其arch_settings定义了两者的阶段结构mmseg/models/backbones/stdc.pyarch_settings { STDCNet1: [(2, 1), (2, 1), (2, 1)], STDCNet2: [(2, 1, 1, 1), (2, 1, 1, 1, 1), (2, 1, 1)] }STDCNet1FLOPs 约 813M每个阶段由 2 个 STDCModule 组成STDCNet2FLOPs 约 1446M阶段内模块数更多源码 docstring 给出了两个 FLOPs 数值。主干默认通道配置为channels (32, 64, 256, 512, 1024)前两个阶段是普通 3x3 卷积下采样num_shallow_features 2从第三阶段起才进入 STDCModule 堆叠。前向时前两个浅层阶段的特征被裁剪outs outs[self.num_shallow_features:]仅输出 256、512、1024 三个深层次特征这与论文 Figure 4 的浅层特征不参与后续 ARM/FFM设计一致源码注释亦有说明。STDCContextPathNet上下文路径网络STDCContextPathNet将STDCNet主干与上下文路径Context Path整合为分割用的完整主干主干输出 → 全局平均池化 1x1 卷积 → 与 Attention Refinement Module 输出逐级上采样相加 → 最后经FeatureFusionModule融合浅层细节特征输出 4 个特征图outs[0]最大分辨率浅层特征送入STDCHead辅助头计算 Detail Loss两个arms_out特征送入两个FCNHead辅助头feat_fuse融合特征送入主解码头。从 configs/base/models/stdc.py 可以看到该主干的默认参数last_in_channels(1024, 512)、out_channels128、ffm_cfgdict(in_channels384, out_channels256, scale_factor4)。STDCHead 与 Detail Aggregation Loss 原理STDCHead定义在 mmseg/models/decode_heads/stdc_head.py继承自FCNHead是 Detail Aggregation 模块在解码端的具体实现。其核心是用拉普拉斯算子从分割标签中自动提取边界真值boundary target再以边界二分类作为辅助监督。实现要点均可从源码确认注册一个 3x3 拉普拉斯核[-1,-1,-1,-1,8,-1,-1,-1,-1]对 GT 分割标签做F.conv2d卷积得到边界响应图以boundary_threshold默认 0.1二值化响应大于阈值置 1否则置 0得到边界目标分别以stride1/2/4对标签卷积并上采样回原尺寸构建边界金字塔boundary pyramid金字塔沿通道堆叠后用fusion_kernel3x1x1权重固定为[[6/10],[3/10],[1/10]]加权聚合再次阈值化得到最终边界标签最终以边界二分类标签为监督计算CrossEntropyLossuse_sigmoidTrue与DiceLoss的组合损失。README 特别提示了一个重要细节STDCHead中的fusion_kernel是不可学习的。虽然论文将其描述为可训练的 1x1 卷积参数但在官方仓库及其他代码库中由于阈值操作会使它脱离计算图实际是常数。源码也通过torch.nn.Parameter(..., requires_gradFalse)明确将其冻结并在注释中说明paper claimsfusion_kernelis a trainable 1x1 conv parameters. However, it is a constant in original repo。若读者自行对照官方仓库复现可打印原始仓库模型参数验证这一点。配置文件深度解读Cityscapes 四套训练配置基础模型配置configs/base/models/stdc.py基础配置将整网定义为EncoderDecoder包含数据预处理器、主干、解码头与三个辅助头数据预处理器SegDataPreProcessorImageNet 均值[123.675, 116.28, 103.53]、方差[58.395, 57.12, 57.375]bgr_to_rgbTrue主干STDCContextPathNet内部backbone_cfg为STDCNetstdc_typeSTDCNet1、channels(32, 64, 256, 512, 1024)、bottleneck_typecat、num_convs4解码头FCNHead输入in_channels256输出 19 类Cityscapes使用OHEMPixelSamplerthresh0.7, min_kept10000在线困难样本挖掘损失为CrossEntropyLoss三个辅助头两个FCNHead各取一个arms_out特征加一个STDCHeadin_channels256, num_classes2, boundary_threshold0.1其中STDCHead的损失为 CE Dice 组合。训练配置与调度策略configs/stdc/stdc1_4xb12-80k_cityscapes-512x1024.py 是 STDC1 的基础训练配置继承_base_下的模型、数据集Cityscapes、默认运行与 80k 调度四份基础配置并作出以下关键设置crop_size (512, 1024) data_preprocessor dict(sizecrop_size) param_scheduler [ dict(typeLinearLR, by_epochFalse, start_factor0.1, begin0, end1000), dict(typePolyLR, eta_min1e-4, power0.9, begin1000, end80000, by_epochFalse) ] train_dataloader dict(batch_size12, num_workers4) val_dataloader dict(batch_size1, num_workers4) test_dataloader val_dataloader输入尺寸裁剪到 512x1024学习率调度前 1000 次迭代用LinearLR线性预热起始因子 0.1之后切换PolyLR多项式衰减power0.9, eta_min1e-4共 80000 次迭代批大小训练每卡 12 张即 README 中所注4 GPUs with 12 samples per GPU总 batch 为 48测试模式基础配置中test_cfgdict(modewhole)即整图推理不切片。configs/stdc/stdc2_4xb12-80k_cityscapes-512x1024.py 仅一行改动将backbone_cfg的stdc_type替换为STDCNet2其余全部复用。预训练权重的挂载方式两份in1k-pre配置通过仅修改init_cfg的方式挂载 ImageNet 预训练权重网络结构本身不动。例如 configs/stdc/stdc1_in1k-pre_4xb12-80k_cityscapes-512x1024.pycheckpoint https://download.openmmlab.com/mmsegmentation/v0.5/pretrain/stdc/stdc1_20220308-5368626c.pth _base_ ./stdc1_4xb12-80k_cityscapes-512x1024.py model dict( backbonedict( backbone_cfgdict( init_cfgdict(typePretrained, checkpointcheckpoint))))注意挂载点是backbone.backbone_cfg即STDCNet本身的init_cfg因为预训练权重对应的是纯STDCNet主干而非完整的STDCContextPathNet。configs/stdc/stdc2_in1k-pre_4xb12-80k_cityscapes-512x1024.py 结构完全相同仅权重链接与类型不同。这两份配置文件与 configs/stdc/metafile.yaml 中记录的权重下载地址相互印证。官方预训练权重转换工具stdc2mmseg.py由于官方仓库的权重键名与 MMSegmentation 不同仓库提供了转换脚本 tools/model_converters/stdc2mmseg.py。其用法为python tools/model_converters/stdc2mmseg.py ${PRETRAIN_PATH} ${STORE_PATH} ${STDC_TYPE}其中${PRETRAIN_PATH}为官方权重路径${STORE_PATH}为转换后权重保存路径${STDC_TYPE}只能是STDC1或STDC2脚本内有断言约束。README 给出的示例python tools/model_converters/stdc2mmseg.py ./STDCNet813M_73.91.tar ./pretrained/stdc1.pth STDC1 python tools/model_converters/stdc2mmseg.py ./STDCNet1446_76.47.tar ./pretrained/stdc2.pth STDC2脚本的核心是convert_stdc函数做以下键名映射源码逐行可查删除cp.前缀将features.num_layer.*映射为stages.stage_idx.*其中 STDC1 的阶段索引为[0,1,2.0,2.1,3.0,3.1,4.0,4.1]STDC2 的索引序列更长2.0~2.3、3.0~3.4、4.0~4.2与主干arch_settings的结构一一对应将conv_list重命名为layers将avd_layer.0/avd_layer.1分别映射为downsample.conv/downsample.bn。脚本在加载官方 checkpoint 时会依次尝试state_dict、model键若都不存在则直接视为裸状态字典最终用torch.save写出转换结果。读者若想自行转换官方仓库的预训练模型按上述命令执行即可若希望直接使用转换好的权重也可使用 metafile 中列出的现成链接。复现结果与注意事项官方复现指标Cityscapes 验证集README 的结果表完整记录了四套配置的复现指标V100 环境、512x1024 输入MethodBackboneCrop SizeLr schdMem (GB)Inf time (fps)DevicemIoUmIoU(msflip)ConfigSTDCSTDC1 (No Pretrain)512x1024800007.1523.06V10071.8273.89configSTDCSTDC1512x102480000--V10074.9476.97configSTDCSTDC2 (No Pretrain)512x1024800008.2723.71V10073.1576.13configSTDCSTDC2512x102480000--V10076.6778.67config对应权重与训练日志下载链接见 configs/stdc/README.md 与 configs/stdc/metafile.yaml。复现时需注意的四个要点默认训练环境Cityscapes 上 STDC 默认是 4 卡 × 每卡 12 样本即上述配置中的batch_size12复现时应保持一致。No Pretrain含义表示从零开始训练pretrainedNone不加载任何 ImageNet 权重。FPS 数值仅供参考README 明确说明论文中的 FPS 环境与本仓库不同——论文输入尺寸为 512x1024 与 768x1536分别相当于本仓库输入尺寸的 50% 与 75%且使用了 TensorRT本仓库结果表为 PyTorch 下的 V100 实测两者不可直接对比。fusion_kernel不可学习如上文源码分析STDCHead.fusion_kernel在 MMSegmentation 中被冻结为常数官方仓库在训练时设置了find_unused_parametersTrueREADME 中给出了官方训练脚本的对应位置读者在分布式训练对比时需要注意这一差异。如何快速验证 STDC 实现仓库提供了针对 STDC 各模块的单元测试 tests/test_models/test_backbones/test_stdc.py覆盖了STDCContextPathNet的标准前向与输出结构STDCNet的约束断言非法类型如STDCNet3、非法bottleneck_type、非法channels长度均会触发assertSTDCModule的通道配置与前向。读者可通过pytest tests/test_models/test_backbones/test_stdc.py运行这些测试快速验证本地环境与实现的正确性。在 mmseg/models/backbones/stdc.py 的STDCNetdocstring 中还提供了最小示例以 1x3x1024x2048 输入前向输出三个特征图尺寸分别为[1, 256, 128, 256]、[1, 512, 64, 128]、[1, 1024, 32, 64]可作为理解主干输出形状的参考。引用如果您的论文或项目使用了 STDC 算法可按 README 给出的 BibTeX 引用inproceedings{fan2021rethinking, title{Rethinking BiSeNet For Real-time Semantic Segmentation}, author{Fan, Mingyuan and Lai, Shenqi and Huang, Junshi and Wei, Xiaoming and Chai, Zhenhua and Luo, Junfeng and Wei, Xiaolin}, booktitle{Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition}, pages{9716--9725}, year{2021} }STDC 在 MMSegmentation 中的完整落地路径——从STDCModule的密集拼接、STDCContextPathNet的上下文路径整合、STDCHead的边界监督到四套可直接运行的训练配置与权重转换脚本——构成了一个从论文到工程复现的闭环。本文所有代码与数值均可回溯到 configs/stdc/ 与 mmseg/models/backbones/stdc.py 等仓库文件读者可据此在 Cityscapes 上开展实时分割的训练与实验。【免费下载链接】mmsegmentationOpenMMLab Semantic Segmentation Toolbox and Benchmark.项目地址: https://gitcode.com/GitHub_Trending/mm/mmsegmentation创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考