语义分割解码头:APNB/AFNB 原理、源码实现与配置实战)
MMSegmentation 中的 ANNAsymmetric Non-local Network语义分割解码头APNB/AFNB 原理、源码实现与配置实战【免费下载链接】mmsegmentationOpenMMLab Semantic Segmentation Toolbox and Benchmark.项目地址: https://gitcode.com/GitHub_Trending/mm/mmsegmentation本文围绕 OpenMMLab 语义分割工具箱 MMSegmentation 中内置的 ANNAsymmetric Non-local Neural Networks for Semantic Segmentation算法展开以 configs/ann/README.md 为骨架结合 mmseg/models/decode_heads/ann_head.py 源码、configs/_base_/models/ann_r50-d8.py基础配置以及 17 个开箱即用的训练配置系统讲解 APNB / AFNB 两个核心模块的计算原理、MMSegmentation 中的类实现与参数含义、Cityscapes / ADE20K / Pascal VOC 12Aug 三大基准的实验结果并给出训练与测试的实操命令。读完本文你将能够理解 ANN 解码头的底层计算图独立解读其配置文件并完成复现实验。ANN 算法背景从 Non-local 到非对称非局部Non-local 模块因其对长距离依赖long-range dependency的建模能力在语义分割中非常有效但代价是高昂的计算量和 GPU 显存占用——这正是它被广泛应用时的主要批评点。ANNAsymmetric Non-local Neural Network论文标题为Asymmetric non-local neural networks for semantic segmentationZhu et al., ICCV 2019针对这一痛点提出了两个关键组件APNBAsymmetric Pyramid Non-local Block非对称金字塔非局部块将金字塔采样模块Pyramid Sampling Module引入 non-local 块在不牺牲性能的前提下大幅降低计算与内存消耗。AFNBAsymmetric Fusion Non-local Block非对称融合非局部块由 APNB 改编而来在充分考虑长距离依赖的基础上融合不同层级的特征从而显著提升分割性能。论文摘要报告了如下量化结论引用自原文档所附论文摘要可在 configs/ann/README.md 中查看在 Cityscapes test set 上达到 81.3 mIoU对于 256×128 的输入APNB 在 GPU 上比普通 non-local 块快约 6 倍GPU 运行内存占用小约 28 倍。这些数据是原始论文中报告的实验数字可作为理解 APNB 效率优势的参考。MMSegmentation 中的源码实现三个核心类ANN 解码头在 MMSegmentation 中实现为ANNHead完整代码位于 mmseg/models/decode_heads/ann_head.py。整个实现可以拆解为三个层次通用自注意力块、AFNB/APNB 两个模块、以及顶层ANNHead。1. 通用 SelfAttentionBlock 与 PPMConcatANNHead复用了位于 mmseg/models/utils/self_attention_block.py 的通用SelfAttentionBlock。该模块实现了标准的 key / query / value 自注意力计算流程分别对 query 与 key 特征做 1×1 卷积投影query_project/key_project并对 key 特征做 value 投影按需对 query、key 做下采样query_downsample/key_downsample计算相似度矩阵sim_map query × key^T当matmul_normTrue时按channels^(-0.5)缩放再做 softmax与 value 相乘得到 context 输出。其中share_key_query参数控制 key 与 query 是否共享投影权重with_out控制是否附加输出投影层out_project且该层默认用 0 初始化init_weights中constant_init(self.out_project, 0)保证训练初期融合模块不破坏主干特征。PPMConcatann_head.py是金字塔池化采样的一种轻量实现它仅对不同尺度的池化输出做展平与拼接不做额外的卷积变换。构造时传入pool_scales默认(1, 3, 6, 8)前向时对每个尺度执行nn.AdaptiveAvgPool2d再把各尺度输出在通道维上cat起来作为 key 的下采样模块。2. AFNB跨层级的非对称融合AFNBann_head.py的定位是跨层级特征融合输入为低层特征low_feats作为 key与高层特征high_feats作为 query为每个query_scale构造一个SelfAttentionBlock使用share_key_queryFalse低层与高层通道数不同不能共享投影key 侧挂载PPMConcat(key_pool_scales)做金字塔采样多个尺度的注意力输出求和torch.stack(priors).sum(dim0)后与high_feats拼接经 1×1bottleneck卷积输出。可以看到 AFNB 正是论文所述在长距离依赖基础上融合不同层级特征的实现低层细节特征通过非局部注意力为高层语义特征提供精细的上下文指引。3. APNB金字塔采样降低非局部的开销APNBann_head.py作用于同层特征输入输出通道相同key 与 query 均为同一份feats因此使用share_key_queryTrue共享投影权重参数量进一步压缩key 侧同样挂载金字塔池化采样PPMConcat把参与注意力计算的 key 数量从全图逐像素缩小为193664110个代表点针对 512×512 输入的特征图这正是非对称asymmetric的由来——query 保持全分辨率key 被大幅采样缩减从而显著降低相似度矩阵的尺寸与显存占用各尺度注意力输出求和后与原始feats拼接经 1×1bottleneck输出。4. ANNHead整体组装ANNHeadann_head.py继承BaseDecodeHead使用input_transformmultiple_select接收多个层级特征并断言len(self.in_channels) 2即要求输入恰好是两个层级的特征。其前向流程forward为low_feats, high_feats self._transform_inputs(inputs) output self.fusion(low_feats, high_feats) # AFNB 跨层融合 output self.dropout(output) output self.bottleneck(output) # 3x3 卷积 output self.context(output) # APNB 同层上下文 output self.cls_seg(output) # 分类卷积即先由 AFNB 融合低层与高层特征 → dropout → 3×3 bottleneck 降维 → APNB 提炼长距离上下文 → 最终分类。该实现与论文架构一一对应且ANNHead已通过MODELS.register_module()注册进 MMSegmentation 模型注册表可直接在配置中以typeANNHead引用。配置文件逐项解析ANN 的全部训练配置位于 configs/ann 目录共 17 个文件覆盖 R-50 / R-101 骨干 × Cityscapes / ADE20K / Pascal VOC 12Aug 数据集 × 多档学习率调度。所有配置均通过继承_base_的方式组织最核心的是共享模型配置 configs/base/models/ann_r50-d8.py。基础模型配置ann_r50-d8.py该文件定义了 ANN ResNet50 的完整模型结构关键部分如下配置块参数取值与含义data_preprocessortypeSegDataPreProcessor输入归一化mean[123.675, 116.28, 103.53]、std[58.395, 57.12, 57.375]bgr_to_rgbTruepad_val0seg_pad_val255backbonetypeResNetV1c,depth504 个 stageout_indices(0,1,2,3)输出全部层级dilations(1,1,2,4)表示 stage3/4 使用空洞卷积d8 即 dilation 8contract_dilationTruedecode_headtypeANNHeadin_channels[1024, 2048]、in_index[2, 3]取 stage3 与 stage4 两个层级与源码len(in_channels)2的断言对应channels512head 内部通道数project_channels256non-local 投影通道数query_scales(1,)key_pool_scales(1,3,6,8)金字塔采样尺度dropout_ratio0.1num_classes19Cityscapesalign_cornersFalse损失为CrossEntropyLossuse_sigmoidFalse,loss_weight1.0auxiliary_headtypeFCNHead辅助头接在in_channels1024stage3channels256num_convs1concat_inputFalseloss_weight0.4与主头共同监督训练train_cfg/test_cfg-train_cfgdict()test_cfgdict(modewhole)即整图推理不做滑窗值得注意的对应关系decode_head的in_channels[1024, 2048]直接决定了ann_head.py中low_in_channels1024与high_in_channels2048的取值进而决定 AFNB 的 key/query 输入维度project_channels256则作为 key/query/value 投影的目标通道数传入两个模块。数据集级配置继承 覆盖以 Cityscapes 40k 配置 configs/ann/ann_r50-d8_4xb2-40k_cityscapes-512x1024.py 为例_base_ [ ../_base_/models/ann_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)通过_base_继承模型、数据集、运行时与调度四类基础配置只需用crop_size与data_preprocessor指定裁剪尺寸并同步给预处理器换数据集时只需额外覆盖类别数例如 ADE20K 配置 configs/ann/ann_r50-d8_4xb4-80k_ade20k-512x512.py 中model dict( data_preprocessordata_preprocessor, decode_headdict(num_classes150), auxiliary_headdict(num_classes150))将主头与辅助头的num_classes从 19 覆盖为 150VOC 12Aug 配置则将num_classes覆盖为 21并换用schedule_20k/schedule_40k与voc12aug.py数据集基础配置。17 个配置一览数据集骨干分辨率迭代数配置文件CityscapesR-50-D8512×102440k / 80kann_r50-d8_4xb2-40k_cityscapes-512x1024.py、ann_r50-d8_4xb2-80k_cityscapes-512x1024.pyCityscapesR-50-D8769×76940k / 80kann_r50-d8_4xb2-40k_cityscapes-769x769.py、ann_r50-d8_4xb2-80k_cityscapes-769x769.pyCityscapesR-101-D8512×1024 / 769×76940k / 80kann_r101-d8_4xb2-{40k,80k}_cityscapes-{512x1024,769x769}.pyADE20KR-50-D8512×51280k / 160kann_r50-d8_4xb4-80k_ade20k-512x512.py、ann_r50-d8_4xb4-160k_ade20k-512x512.pyADE20KR-101-D8512×51280k / 160kann_r101-d8_4xb4-80k_ade20k-512x512.py、ann_r101-d8_4xb4-160k_ade20k-512x512.pyVOC 12AugR-50-D8512×51220k / 40kann_r50-d8_4xb4-20k_voc12aug-512x512.py、ann_r50-d8_4xb4-40k_voc12aug-512x512.pyVOC 12AugR-101-D8512×51220k / 40kann_r101-d8_4xb4-20k_voc12aug-512x512.py、ann_r101-d8_4xb4-40k_voc12aug-512x512.py配置命名遵循 MMSegmentation 惯例ann_{backbone}_4xb{batch}_{iters}k_{dataset}-{crop_size}.py其中d8表示骨干网络 dilation8 的空洞卷积设置。基准实验结果以下结果表格完整摘自 configs/ann/README.md为官方在 V100 上的复现数据mIoU 为单尺度测试与 msflip多尺度 水平翻转两种评测方式的结果。Cityscapes19 类MethodBackboneCrop SizeLr schdMem (GB)Inf time (fps)DevicemIoUmIoU(msflip)ANNR-50-D8512×10244000063.71V10077.4078.57ANNR-101-D8512×1024400009.52.55V10076.5578.85ANNR-50-D8769×769400006.81.70V10078.8980.46ANNR-101-D8769×7694000010.71.15V10079.3280.94ANNR-50-D8512×102480000--V10077.3478.65ANNR-101-D8512×102480000--V10077.1478.81ANNR-50-D8769×76980000--V10078.8880.57ANNR-101-D8769×76980000--V10078.8080.34ADE20K150 类MethodBackboneCrop SizeLr schdMem (GB)Inf time (fps)DevicemIoUmIoU(msflip)ANNR-50-D8512×512800009.121.01V10041.0142.30ANNR-101-D8512×5128000012.514.12V10042.9444.18ANNR-50-D8512×512160000--V10041.7442.62ANNR-101-D8512×512160000--V10042.9444.06Pascal VOC 2012 Aug21 类MethodBackboneCrop SizeLr schdMem (GB)Inf time (fps)DevicemIoUmIoU(msflip)ANNR-50-D8512×51220000620.92V10074.8676.13ANNR-101-D8512×512200009.513.94V10077.4778.70ANNR-50-D8512×51240000--V10076.5677.51ANNR-101-D8512×51240000--V10076.7078.06几点可观察到的规律Cityscapes 上 769×769 大裁剪尺寸一致优于 512×1024约 1.5 mIoUR-101 在 512×1024 下略低于 R-50但在 769×769 与 ADE20K/VOC 上反超同配置下 40k 与 80k 迭代结果相当说明 40k 已基本收敛。训练与测试单卡训练python tools/train.py configs/ann/ann_r50-d8_4xb2-40k_cityscapes-512x1024.py多卡分布式训练bash tools/dist_train.sh configs/ann/ann_r50-d8_4xb2-40k_cityscapes-512x1024.py 8其中8为 GPU 卡数Slurm 集群环境可改用 tools/slurm_train.sh。测试与评测python tools/test.py configs/ann/ann_r50-d8_4xb2-40k_cityscapes-512x1024.py /path/to/checkpoint.pth测试时模型使用配置中的test_cfgdict(modewhole)即整图推理模式如需切换为滑窗slide模式可在配置或命令行中覆盖test_cfg。单元测试验证MMSegmentation 为ANNHead提供了单元测试 tests/test_models/test_heads/test_ann_head.py用随机输入验证前向逻辑与输出形状inputs [torch.randn(1, 4, 45, 45), torch.randn(1, 8, 21, 21)] head ANNHead( in_channels[4, 8], channels2, num_classes19, in_index[-2, -1], project_channels8) outputs head(inputs) assert outputs.shape (1, head.num_classes, 21, 21)该测试从两个关键侧面印证了实现约束其一in_channels必须是长度为 2 的列表对应源码中assert len(self.in_channels) 2的断言其二输出分辨率与第二个输入高层特征21×21一致说明 AFNB/APNB 的注意力计算不会改变特征图的空间尺寸最终分类卷积输出num_classes个通道的分割 logits。引用文献如果要在论文或技术报告中引用 ANN 算法可使用如下 BibTeX摘自 configs/ann/README.mdinproceedings{zhu2019asymmetric, title{Asymmetric non-local neural networks for semantic segmentation}, author{Zhu, Zhen and Xu, Mengde and Bai, Song and Huang, Tengteng and Bai, Xiang}, booktitle{Proceedings of the IEEE/CVF International Conference on Computer Vision}, pages{593--602}, year{2019} }小结ANN 通过金字塔采样 key 全分辨率 query的非对称设计把 non-local 模块从贵而准变为快而准并在 MMSegmentation 中落地为结构清晰、可直接复用的ANNHeadAFNB负责跨层级融合、APNB负责同层上下文二者共享同一个通用SelfAttentionBlock。配合 configs/ann 下 17 个覆盖三大数据集与两种骨干的现成配置读者可以零门槛复现论文在 Cityscapes最高 79.32 mIoU 769×769、ADE20K最高 42.94 mIoU与 VOC 12Aug最高 77.47 mIoU上的官方结果并在此基础上进一步调参或改进。【免费下载链接】mmsegmentationOpenMMLab Semantic Segmentation Toolbox and Benchmark.项目地址: https://gitcode.com/GitHub_Trending/mm/mmsegmentation创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考