MMSegmentation 中的 NonLocal Net：长距离依赖语义分割网络原理、源码解析与实战指南

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 Networks（Non-local Neural Networks，Wang 等，CVPR 2018）提出了一族通用的非局部操作构建块，用于直接捕获长距离依赖。其灵感来自计算机视觉中经典的非局部均值（non-local means）方法：

非局部操作计算某个位置上的响应时，将其定义为所有位置上特征的加权和。

形式化地讲，对于一个位置上的输出响应，其计算公式可概括为：

y_i = (1 / C(x)) * Σ_j f(x_i, x_j) * g(x_j)

其中f用于计算位置i与位置j之间的两两相似度（pairwise function），g是位置j处特征的变换函数，C(x)是归一化因子。由于求和遍历全部位置，任意两个位置（无论空间距离多远）都能直接交互，从而一次性捕获全图的长距离依赖。

通用性与应用场景

非局部块是即插即用的通用组件，可以嵌入到多种视觉架构中。原文在视频分类任务（Kinetics、Charades）与静态图像任务（COCO 上的目标检测、分割、姿态估计）上均验证了其有效性。在 MMSegmentation 中，NonLocal Net 被实现为一个解码头（decode head），用于语义分割：以 ResNet 等骨干网络提取多尺度特征，在解码阶段通过非局部块聚合全局上下文，再输出逐像素分类结果。

源码实现剖析：NLHead 与 NonLocal2d

NLHead 类定义

在 MMSegmentation 中，NonLocal Net 的解码头实现位于 mmseg/models/decode_heads/nl_head.py，核心类为NLHead，它通过@MODELS.register_module()注册到模型注册表，因此可以直接在配置文件中以type='NLHead'使用。

从源码看，NLHead继承自FCNHead，并在其基础上插入一个非局部块：

• 构造函数接收三个关键参数：
  • reduction（int，默认2）：投影变换的降维因子，用于降低非局部块中通道投影的计算量；
  • use_scale（bool，默认True）：是否将两两相似度权重乘以sqrt(1/inter_channels)进行缩放，以稳定训练；
  • mode（str，默认'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的计算顺序清晰体现了"卷积 - 非局部聚合 - 卷积 - 分类"的典型结构：

1. _transform_inputs(inputs)：从骨干网络多尺度输出中取出对应索引（in_index）的特征；
2. self.convs0：第一组卷积对输入特征做通道变换；
3. self.nl_block(output)：非局部块在特征图上执行全图两两交互，聚合长距离上下文；
4. self.convs1：第二组卷积进一步处理聚合后的特征；
5. 若concat_input=True（继承自FCNHead的默认行为），通过conv_cat将输入与输出拼接；
6. self.cls_seg(output)：分类卷积层输出逐像素 logits。

因此NLHead是一个"FCN 骨架 + 非局部模块"的组合体：两个卷积负责通道变换与局部细节，中间的非局部块负责全局建模。单元测试 tests/test_models/test_heads/test_nl_head.py 验证了该结构的核心不变量：

head = NLHead(in_channels=8, channels=4, num_classes=19) 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 执行。另外，同目录的DNLHead（Disentangled Non-Local）是 NLHead 的后续变体，可参考 test_dnl_head.py 了解其在gaussian、dot_product、concatenation等模式下的行为。

配置文件逐项解析

基础模型配置 nonlocal_r50-d8.py

NonLocal 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_rgb=True，pad_val=0，seg_pad_val=255；
• 骨干网络：ResNetV1c，depth=50，out_indices=(0, 1, 2, 3)，空洞卷积配置dilations=(1, 1, 2, 4)、strides=(1, 2, 1, 1)（即 D8 系列，输出 stride 为 8），预训练权重open-mmlab://resnet50_v1c，contract_dilation=True；
• 解码头：type='NLHead'，in_channels=2048、in_index=3（取骨干最后一级特征）、channels=512、dropout_ratio=0.1，非局部参数为reduction=2、use_scale=True、mode='embedded_gaussian'，num_classes=19（Cityscapes 类别数），损失为CrossEntropyLoss（use_sigmoid=False，loss_weight=1.0）；
• 辅助头：type='FCNHead'，in_channels=1024、in_index=2、channels=256、num_convs=1、concat_input=False、num_classes=19，损失权重0.4，辅助监督帮助训练收敛；
• 推理设置：test_cfg=dict(mode='whole')，即整图推理（不做滑动窗口切块）。

一个完整的 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(size=crop_size) model = dict(data_preprocessor=data_preprocessor)

• 4xb2表示 4 张 GPU、每张 batch size 为 2（总 batch size 8）；
• 40k表示训练 40000 次迭代，调度配置见 configs/base/schedules/schedule_40k.py：SGD 优化器（lr=0.01、momentum=0.9、weight_decay=0.0005），PolyLR学习率策略（power=0.9、eta_min=1e-4），IterBasedTrainLoop每 4000 次迭代验证一次并保存 checkpoint；
• crop_size通过data_preprocessor.size传递给训练与评测管线，保证输入尺寸一致。

ADE20K 与 VOC 配置的差异点

• ADE20K：nonlocal_r50-d8_4xb4-80k_ade20k-512x512.py 将decode_head与auxiliary_head的num_classes改为150（ADE20K 类别数），crop size 为512x512，采用schedule_80k；
• R-101 骨干：nonlocal_r101-d8_4xb4-80k_ade20k-512x512.py 仅以 R-50 配置为基底，替换pretrained='open-mmlab://resnet101_v1c'与backbone=dict(depth=101)，体现了 MMSegmentation 配置继承机制的简洁性；
• Pascal VOC 2012 + Aug：nonlocal_r50-d8_4xb4-20k_voc12aug-512x512.py等配置面向voc12aug数据集，训练迭代数为 20k/40k。

实验结果与模型库

原文档按数据集给出了完整的 benchmark 结果，以下表格完整保留各项指标。模型权重与训练日志的索引可查看 configs/nonlocal_net/metafile.yaml（其中登记了每个模型的 Weights 与 Training log 下载地址，以及 4×V100 的训练资源配置）。所有结果均为 V100 GPU 上测得。

Cityscapes

Method	Backbone	Crop Size	Lr schd	Mem (GB)	Inf time (fps)	Device	mIoU	mIoU(ms+flip)	config
NonLocalNet	R-50-D8	512x1024	40000	7.4	2.72	V100	78.24	-	config
NonLocalNet	R-101-D8	512x1024	40000	10.9	1.95	V100	78.66	-	config
NonLocalNet	R-50-D8	769x769	40000	8.9	1.52	V100	78.33	79.92	config
NonLocalNet	R-101-D8	769x769	40000	12.8	1.05	V100	78.57	80.29	config
NonLocalNet	R-50-D8	512x1024	80000	-	-	V100	78.01	-	config
NonLocalNet	R-101-D8	512x1024	80000	-	-	V100	78.93	-	config
NonLocalNet	R-50-D8	769x769	80000	-	-	V100	79.05	80.68	config
NonLocalNet	R-101-D8	769x769	80000	-	-	V100	79.40	80.85	config

ADE20K

Method	Backbone	Crop Size	Lr schd	Mem (GB)	Inf time (fps)	Device	mIoU	mIoU(ms+flip)	config
NonLocalNet	R-50-D8	512x512	80000	9.1	21.37	V100	40.75	42.05	config
NonLocalNet	R-101-D8	512x512	80000	12.6	13.97	V100	42.90	44.27	config
NonLocalNet	R-50-D8	512x512	160000	-	-	V100	42.03	43.04	config
NonLocalNet	R-101-D8	512x512	160000	-	-	V100	44.63	45.79	config

Pascal VOC 2012 + Aug

Method	Backbone	Crop Size	Lr schd	Mem (GB)	Inf time (fps)	Device	mIoU	mIoU(ms+flip)	config
NonLocalNet	R-50-D8	512x512	20000	6.4	21.21	V100	76.20	77.12	config
NonLocalNet	R-101-D8	512x512	20000	9.8	14.01	V100	78.15	78.86	config
NonLocalNet	R-50-D8	512x512	40000	-	-	V100	76.65	77.47	config
NonLocalNet	R-101-D8	512x512	40000	-	-	V100	78.27	79.12	config

从结果可以观察到的规律（基于表格数据的事实归纳）：更高分辨率的 crop size（769x769 对比 512x1024）普遍带来 mIoU 提升，多尺度 + 翻转测试（ms+flip）较单尺度进一步提升约 1～2 个点；R-101 骨干一致优于 R-50；ADE20K 上 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.lr=0.005

测试与推理

测试使用 tools/test.py，配合 dist_test.sh 可做多卡评测。例如加载 Cityscapes 40k 模型权重进行整图评测（test_cfg.mode='whole'）：

python tools/test.py configs/nonlocal_net/nonlocal_r50-d8_4xb2-40k_cityscapes-512x1024.py /path/to/checkpoint.pth --eval mIoU

对单张图片做可视化推理，可使用 demo/image_demo.py：

python 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