MMSegmentation 中的 NonLocal Net长距离依赖语义分割网络原理、源码解析与实战指南【免费下载链接】mmsegmentationOpenMMLab Semantic Segmentation Toolbox and Benchmark.项目地址: https://gitcode.com/GitHub_Trending/mm/mmsegmentation本文以 MMSegmentation 仓库中 configs/nonlocal_net/README.md 为主体结合 NLHead 源码、基础模型配置、单元测试 与 训练入口 进行纵深解读帮助读者掌握 NonLocal Net 在语义分割中的完整落地路径算法原理、源码实现、配置撰写、训练与评测实操。算法背景为什么需要非局部建模局部操作的天然局限卷积与循环操作本质上都是逐局部邻域处理的基本构建块卷积核每次只感知一个固定大小的感受野循环操作沿时间轴逐步传递信息。对于语义分割这类需要全局上下文的任务仅靠堆叠卷积层来扩大感受野往往带来参数与计算开销的急剧上升且远距离依赖仍难以高效建模。非局部操作的核心思想Non-local Neural NetworksNon-local Neural NetworksWang 等CVPR 2018提出了一族通用的非局部操作构建块用于直接捕获长距离依赖。其灵感来自计算机视觉中经典的非局部均值non-local means方法非局部操作计算某个位置上的响应时将其定义为所有位置上特征的加权和。形式化地讲对于一个位置上的输出响应其计算公式可概括为y_i (1 / C(x)) * Σ_j f(x_i, x_j) * g(x_j)其中f用于计算位置i与位置j之间的两两相似度pairwise functiong是位置j处特征的变换函数C(x)是归一化因子。由于求和遍历全部位置任意两个位置无论空间距离多远都能直接交互从而一次性捕获全图的长距离依赖。通用性与应用场景非局部块是即插即用的通用组件可以嵌入到多种视觉架构中。原文在视频分类任务Kinetics、Charades与静态图像任务COCO 上的目标检测、分割、姿态估计上均验证了其有效性。在 MMSegmentation 中NonLocal Net 被实现为一个解码头decode head用于语义分割以 ResNet 等骨干网络提取多尺度特征在解码阶段通过非局部块聚合全局上下文再输出逐像素分类结果。源码实现剖析NLHead 与 NonLocal2dNLHead 类定义在 MMSegmentation 中NonLocal Net 的解码头实现位于 mmseg/models/decode_heads/nl_head.py核心类为NLHead它通过MODELS.register_module()注册到模型注册表因此可以直接在配置文件中以typeNLHead使用。从源码看NLHead继承自FCNHead并在其基础上插入一个非局部块构造函数接收三个关键参数reductionint默认2投影变换的降维因子用于降低非局部块中通道投影的计算量use_scalebool默认True是否将两两相似度权重乘以sqrt(1/inter_channels)进行缩放以稳定训练modestr默认embedded_gaussian非局部模式可选embedded_gaussian、dot_product此外 mmcv 的NonLocal2d还支持gaussian、concatenation等模式。内部通过mmcv.cnn.NonLocal2d构造非局部块self.nl_block其in_channels取解码头的self.channels并将conv_cfg、norm_cfg一并透传。前向计算流程NLHead.forward的计算顺序清晰体现了卷积 - 非局部聚合 - 卷积 - 分类的典型结构_transform_inputs(inputs)从骨干网络多尺度输出中取出对应索引in_index的特征self.convs0第一组卷积对输入特征做通道变换self.nl_block(output)非局部块在特征图上执行全图两两交互聚合长距离上下文self.convs1第二组卷积进一步处理聚合后的特征若concat_inputTrue继承自FCNHead的默认行为通过conv_cat将输入与输出拼接self.cls_seg(output)分类卷积层输出逐像素 logits。因此NLHead是一个FCN 骨架 非局部模块的组合体两个卷积负责通道变换与局部细节中间的非局部块负责全局建模。单元测试 tests/test_models/test_heads/test_nl_head.py 验证了该结构的核心不变量head NLHead(in_channels8, channels4, num_classes19) assert len(head.convs) 2 assert hasattr(head, nl_block) # 输入 (1, 8, 23, 23)输出 (1, num_classes, 23, 23)空间尺寸保持 outputs head(inputs) assert outputs.shape (1, head.num_classes, 23, 23)测试同时确认了NLHead在 GPU 可用时会自动迁移到 CUDA 执行。另外同目录的DNLHeadDisentangled Non-Local是 NLHead 的后续变体可参考 test_dnl_head.py 了解其在gaussian、dot_product、concatenation等模式下的行为。配置文件逐项解析基础模型配置 nonlocal_r50-d8.pyNonLocal Net 的完整模型配置位于 configs/base/models/nonlocal_r50-d8.py它定义了ResNetV1c 骨干 NLHead 解码头 FCNHead 辅助头的整体结构关键点如下数据预处理器SegDataPreProcessor使用 ImageNet 统计的均值[123.675, 116.28, 103.53]与标准差[58.395, 57.12, 57.375]bgr_to_rgbTruepad_val0seg_pad_val255骨干网络ResNetV1cdepth50out_indices(0, 1, 2, 3)空洞卷积配置dilations(1, 1, 2, 4)、strides(1, 2, 1, 1)即 D8 系列输出 stride 为 8预训练权重open-mmlab://resnet50_v1ccontract_dilationTrue解码头typeNLHeadin_channels2048、in_index3取骨干最后一级特征、channels512、dropout_ratio0.1非局部参数为reduction2、use_scaleTrue、modeembedded_gaussiannum_classes19Cityscapes 类别数损失为CrossEntropyLossuse_sigmoidFalseloss_weight1.0辅助头typeFCNHeadin_channels1024、in_index2、channels256、num_convs1、concat_inputFalse、num_classes19损失权重0.4辅助监督帮助训练收敛推理设置test_cfgdict(modewhole)即整图推理不做滑动窗口切块。一个完整的 Cityscapes 训练配置以 configs/nonlocal_net/nonlocal_r50-d8_4xb2-40k_cityscapes-512x1024.py 为例它通过_base_继承四份基础配置_base_ [ ../_base_/models/nonlocal_r50-d8.py, ../_base_/datasets/cityscapes.py, ../_base_/default_runtime.py, ../_base_/schedules/schedule_40k.py ] crop_size (512, 1024) data_preprocessor dict(sizecrop_size) model dict(data_preprocessordata_preprocessor)4xb2表示 4 张 GPU、每张 batch size 为 2总 batch size 840k表示训练 40000 次迭代调度配置见 configs/base/schedules/schedule_40k.pySGD 优化器lr0.01、momentum0.9、weight_decay0.0005PolyLR学习率策略power0.9、eta_min1e-4IterBasedTrainLoop每 4000 次迭代验证一次并保存 checkpointcrop_size通过data_preprocessor.size传递给训练与评测管线保证输入尺寸一致。ADE20K 与 VOC 配置的差异点ADE20Knonlocal_r50-d8_4xb4-80k_ade20k-512x512.py 将decode_head与auxiliary_head的num_classes改为150ADE20K 类别数crop size 为512x512采用schedule_80kR-101 骨干nonlocal_r101-d8_4xb4-80k_ade20k-512x512.py 仅以 R-50 配置为基底替换pretrainedopen-mmlab://resnet101_v1c与backbonedict(depth101)体现了 MMSegmentation 配置继承机制的简洁性Pascal VOC 2012 Augnonlocal_r50-d8_4xb4-20k_voc12aug-512x512.py等配置面向voc12aug数据集训练迭代数为 20k/40k。实验结果与模型库原文档按数据集给出了完整的 benchmark 结果以下表格完整保留各项指标。模型权重与训练日志的索引可查看 configs/nonlocal_net/metafile.yaml其中登记了每个模型的 Weights 与 Training log 下载地址以及 4×V100 的训练资源配置。所有结果均为 V100 GPU 上测得。CityscapesMethodBackboneCrop SizeLr schdMem (GB)Inf time (fps)DevicemIoUmIoU(msflip)configNonLocalNetR-50-D8512x1024400007.42.72V10078.24-configNonLocalNetR-101-D8512x10244000010.91.95V10078.66-configNonLocalNetR-50-D8769x769400008.91.52V10078.3379.92configNonLocalNetR-101-D8769x7694000012.81.05V10078.5780.29configNonLocalNetR-50-D8512x102480000--V10078.01-configNonLocalNetR-101-D8512x102480000--V10078.93-configNonLocalNetR-50-D8769x76980000--V10079.0580.68configNonLocalNetR-101-D8769x76980000--V10079.4080.85configADE20KMethodBackboneCrop SizeLr schdMem (GB)Inf time (fps)DevicemIoUmIoU(msflip)configNonLocalNetR-50-D8512x512800009.121.37V10040.7542.05configNonLocalNetR-101-D8512x5128000012.613.97V10042.9044.27configNonLocalNetR-50-D8512x512160000--V10042.0343.04configNonLocalNetR-101-D8512x512160000--V10044.6345.79configPascal VOC 2012 AugMethodBackboneCrop SizeLr schdMem (GB)Inf time (fps)DevicemIoUmIoU(msflip)configNonLocalNetR-50-D8512x512200006.421.21V10076.2077.12configNonLocalNetR-101-D8512x512200009.814.01V10078.1578.86configNonLocalNetR-50-D8512x51240000--V10076.6577.47configNonLocalNetR-101-D8512x51240000--V10078.2779.12config从结果可以观察到的规律基于表格数据的事实归纳更高分辨率的 crop size769x769 对比 512x1024普遍带来 mIoU 提升多尺度 翻转测试msflip较单尺度进一步提升约 12 个点R-101 骨干一致优于 R-50ADE20K 上 160k 迭代的 R-101 配置取得该数据集内的最优结果 44.63。训练、测试与推理实操训练使用 tools/train.py 启动训练其支持--work-dir保存日志与模型、--resume从 work_dir 最新 checkpoint 恢复、--amp混合精度训练、--cfg-options命令行覆盖配置等参数。单机单卡示例python tools/train.py configs/nonlocal_net/nonlocal_r50-d8_4xb2-40k_cityscapes-512x1024.py多卡分布式训练4 卡与配置命名4xb2的假设一致bash tools/dist_train.sh configs/nonlocal_net/nonlocal_r50-d8_4xb2-40k_cityscapes-512x1024.py 4也可以通过--cfg-options临时调整超参数例如修改学习率python tools/train.py configs/nonlocal_net/nonlocal_r50-d8_4xb2-40k_cityscapes-512x1024.py --cfg-options optim_wrapper.optimizer.lr0.005测试与推理测试使用 tools/test.py配合 dist_test.sh 可做多卡评测。例如加载 Cityscapes 40k 模型权重进行整图评测test_cfg.modewholepython tools/test.py configs/nonlocal_net/nonlocal_r50-d8_4xb2-40k_cityscapes-512x1024.py /path/to/checkpoint.pth --eval mIoU对单张图片做可视化推理可使用 demo/image_demo.pypython demo/image_demo.py demo/demo.png \ configs/nonlocal_net/nonlocal_r50-d8_4xb2-40k_cityscapes-512x1024.py \ /path/to/checkpoint.pth --out-file result.png需要说明的适用前提上述命令均基于当前仓库MMSegmentation 1.x / MMEngine 运行体系编写权重文件需按 metafile.yaml 中登记的链接下载后填入本地路径。总结NonLocal Net 在 MMSegmentation 中的落地展示了如何把一个通用的非局部数学算子全图位置两两加权求和封装成即插即用的解码头NLHead以两个卷积包夹一个NonLocal2d块在保留 FCN 简洁结构的同时显著增强全局上下文建模能力。通过继承式配置同一个模型可以极低成本迁移到 Cityscapes、ADE20K、Pascal VOC 等数据集与 R-50/R-101 骨干组合配合丰富的 benchmark 与完整工具链是理解长距离依赖建模与 MMSegmentation 解码头机制的上佳范例。Citation如果本模型对你的研究有所帮助请引用原论文inproceedings{wang2018non, title{Non-local neural networks}, author{Wang, Xiaolong and Girshick, Ross and Gupta, Abhinav and He, Kaiming}, booktitle{Proceedings of the IEEE conference on computer vision and pattern recognition}, pages{7794--7803}, year{2018} }【免费下载链接】mmsegmentationOpenMMLab Semantic Segmentation Toolbox and Benchmark.项目地址: https://gitcode.com/GitHub_Trending/mm/mmsegmentation创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
