MMSegmentation 中的 ISANet:交错稀疏自注意力语义分割解码头原理与实战

MMSegmentation 中的 ISANet:交错稀疏自注意力语义分割解码头原理与实战 MMSegmentation 中的 ISANet交错稀疏自注意力语义分割解码头原理与实战【免费下载链接】mmsegmentationOpenMMLab Semantic Segmentation Toolbox and Benchmark.项目地址: https://gitcode.com/GitHub_Trending/mm/mmsegmentation导读ISANetInterlaced Sparse Self-Attention是 MMSegmentation 中一个以自注意力为核心的语义分割算法其核心思路是把稠密的自注意力亲和矩阵分解为两个稀疏亲和矩阵的乘积从而在保持全局感受野的同时显著降低计算与显存开销。本文以仓库中 configs/isanet/README.md 为主线结合 mmseg/models/decode_heads/isa_head.py 的源码实现与 configs/isanet 下的真实配置完整讲解 ISAHead 的原理、配置字段、数据集实验结果并给出在 Cityscapes / ADE20K / VOC 上直接可复制的训练与推理方案。读完后你将掌握 ISA 注意力模块的内部数据流以及如何在 MMSegmentation 中一键复现 ISANet 的全部官方结果。ISANet 算法背景与核心思想ISANet 出自论文Interlaced Sparse Self-Attention for Semantic SegmentationarXiv:1907.12273其作者与 OCNetObject Context Network同源。语义分割任务需要稠密、全局的上下文建模标准的 self-attention 虽能建模任意两个位置之间的关系但对高分辨率特征图而言其亲和矩阵的计算与存储复杂度随像素数平方增长代价过高。ISANet 的主要思想是将稠密亲和矩阵因式分解为两个稀疏亲和矩阵的乘积第一个注意力模块估计“空间间隔距离较长”的位置子集内的亲和关系global relation第二个注意力模块估计“空间间隔距离较短”的位置子集内的亲和关系local relation。两个模块串联设计使得每一个位置都能间接接收到所有其他位置的信息从而在语义上等价于全局自注意力但计算与内存复杂度大幅下降尤其适用于高分辨率特征图。在 configs/isanet/README.md 的摘要中作者在六个语义分割基准上验证了该方法的有效性。MMSegmentation 官方复现了 Cityscapes、ADE20K、Pascal VOC 2012 Aug 三个数据集上的完整实验下文将逐一给出结果表与对应配置。ISAHead 源码结构两个稀疏注意力如何串联MMSegmentation 中的 ISA 实现位于 mmseg/models/decode_heads/isa_head.py通过MODELS.register_module()注册为ISAHead可直接在配置中通过typeISAHead引用。顶层结构in_conv → 双注意力 → out_convISAHead继承自BaseDecodeHead其__init__接收两个关键参数isa_channels (int)ISA 模块内部的通道数即 key/query 变换的输出通道down_factor (tuple[int])local group 的尺寸默认(8, 8)表示在 H 轴和 W 轴上分别以 8 个像素为一个局部分组。其网络结构为in_conv一个3x3卷积含conv_cfg/norm_cfg/act_cfg把in_channels映射到channelsglobal_relation一个SelfAttentionBlock用于在全局分组内计算长距离亲和local_relation另一个SelfAttentionBlock用于在局部分组内计算短距离亲和out_conv一个1x1卷积把channels * 2注意力输出与残差拼接映射回channels最后经cls_seg得到分割 logits。forward 中的交错数据流forward是理解 ISA 的关键对应 isa_head.py其流程可拆解为以下步骤输入特征经in_conv得到x并保存residual x作为残差连接取出loc_h, loc_w self.down_factor局部组尺寸计算全局组数量glb_h ceil(h / loc_h)、glb_w ceil(w / loc_w)若特征尺寸不能被down_factor整除则用F.pad对称补齐全局关系将x变形为(n, c, glb_h, loc_h, glb_w, loc_w)经permute置换为(n, loc_h, loc_w, c, glb_h, glb_w)并 reshape 为(-1, c, glb_h, glb_w)即把空间上相隔loc_h × loc_w的远距离像素聚合到同一批分组中然后送入global_relation做自注意力局部关系将全局注意力的输出按相反方向置换reshape 为(-1, c, loc_h, loc_w)即把每个全局组内相邻的近距离像素聚为分组送入local_relation把像素置换回原始位置裁掉先前补的 padding将注意力输出与residual在通道维拼接经out_conv与cls_seg输出分割结果。正是这两步“先全局、后局部”的交错注意力保证每个位置都能经由中间位置间接获得全图信息而每次矩阵乘法的规模都被控制在分组尺寸内——这是 ISANet 降低复杂度、适合高分辨率输入的根本原因。SelfAttentionBlock通用的自注意力基元ISAHead中两个注意力模块都基于 mmseg/models/utils/self_attention_block.py 中的通用SelfAttentionBlock即 non-local block。该基元通过 key/query/value 三个投影build_project构建的 1x1 卷积序列完成标准注意力计算sim_map matmul(query, key) if matmul_norm: sim_map (channels ** -0.5) * sim_map sim_map softmax(sim_map, dim-1) context matmul(sim_map, value)在isa_head.py中SelfAttentionBlock子类以share_key_queryFalse、key_query_num_convs2、key_query_normTrue、matmul_normTrue、with_outFalse实例化并在外部补一个output_project卷积输出。matmul_normTrue意味着注意力分数会除以sqrt(channels)做缩放对应 Transformer 中的 scaled dot-product attention。配置全解从基础模型到数据集变体基础模型配置configs/_base_/models/isanet_r50-d8.pyISANet 在 MMSegmentation 中以“R-50-D8 / R-101-D8”形式命名D8表示骨干网络使用dilations(1, 1, 2, 4)的膨胀策略。基础配置见 configs/base/models/isanet_r50-d8.py其要点如下骨干ResNetV1cdepth50out_indices(0,1,2,3)输出四个阶段特征预训练权重为open-mmlab://resnet50_v1c解码头typeISAHeadin_channels2048、in_index3取骨干最后阶段、channels512、isa_channels256、down_factor(8, 8)、dropout_ratio0.1、num_classes19Cityscapes辅助头typeFCNHeadin_channels1024、in_index2loss_weight0.4损失主头与辅助头均为CrossEntropyLoss主头loss_weight1.0数据预处理SegDataPreProcessor采用 ImageNet 统计量mean[123.675, 116.28, 103.53]、std[58.395, 57.12, 57.375]bgr_to_rgbTrue。训练计划与运行时schedule_40k/schedule_160kISANet 的实验统一使用 Poly 学习率与迭代式训练。以 configs/base/schedules/schedule_40k.py 为例优化器SGDlr0.01、momentum0.9、weight_decay0.0005经OptimWrapper包装学习率策略PolyLRpower0.9、eta_min1e-4by_epochFalse训练循环IterBasedTrainLoopmax_iters40000val_interval4000钩子CheckpointHook每 4000 次迭代保存一次权重LoggerHook每 50 次迭代打印日志。三个数据集的配置继承方式configs/isanet目录下共 14 个实验配置均通过_base_继承组合仅在必要处覆盖字段例如Cityscapes 512x1024configs/isanet/isanet_r50-d8_4xb2-40k_cityscapes-512x1024.py继承isanet_r50-d8.pycityscapes.pydefault_runtime.pyschedule_40k.py仅覆盖crop_size(512,1024)与data_preprocessor.sizeADE20K 512x512configs/isanet/isanet_r50-d8_4xb4-160k_ade20k-512x512.py继承ade20k.pyschedule_160k.py并把decode_head.num_classes与auxiliary_head.num_classes改为150R-101 变体configs/isanet/isanet_r101-d8_4xb4-160k_ade20k-512x512.py直接以对应 R-50 配置为_base_仅覆盖pretrainedopen-mmlab://resnet101_v1c与backbone.depth101体现 MMSegmentation 配置的高度复用。官方复现结果一览下表完整继承自 configs/isanet/README.md记录了 V100 上的显存占用、推理速度与 mIoU含单尺度与多尺度翻转两种评测协议。评测口径为mIoU与mIoU(msflip)后者使用多尺度输入并叠加水平翻转。Cityscapes19 类骨干Crop Size迭代数Mem (GB)Inf time (fps)mIoUmIoU(msflip)配置文件R-50-D8512x102440k5.8692.9178.4979.44configR-50-D8512x102480k5.8692.9178.6880.25configR-50-D8769x76940k6.7591.5478.7080.28configR-50-D8769x76980k6.7591.5479.2980.53configR-101-D8512x102440k9.4252.3579.5881.05configR-101-D8512x102480k9.4252.3580.3281.58configR-101-D8769x76940k10.8150.9279.6880.95configR-101-D8769x76980k10.8150.9280.6181.59configADE20K150 类骨干Crop Size迭代数Mem (GB)Inf time (fps)mIoUmIoU(msflip)配置文件R-50-D8512x51280k9.022.5541.1242.35configR-50-D8512x512160k9.022.5542.5943.07configR-101-D8512x51280k12.56210.5643.5144.38configR-101-D8512x512160k12.56210.5643.8045.4configPascal VOC 2012 Aug21 类骨干Crop Size迭代数Mem (GB)Inf time (fps)mIoUmIoU(msflip)配置文件R-50-D8512x51220k5.923.0876.7877.79configR-50-D8512x51240k5.923.0876.2077.22configR-101-D8512x51220k9.4657.4278.4679.16configR-101-D8512x51240k9.4657.4278.1279.04config读表要点Cityscapes 上 80k 迭代普遍优于 40kR-101 相比 R-50 在 ADE20K 上有约 1.22.4 个点的提升769x769 的 crop size 相比 512x1024 带来更稳的 mIoU 提升但推理帧率明显下降显存占用也更高——这正是 ISA 稀疏注意力在高分辨率输入上仍可运行的体现。从零复现训练、测试与单图推理安装与数据准备本文所有命令均基于 MMSegmentation 1.x 的tools脚本。请先参照仓库根目录 README.md 安装 mmengine、mmcv 等依赖并按 configs/base/datasets/cityscapes.py、configs/base/datasets/ade20k.py 等数据集配置中data_root指示的目录结构放置数据集。单卡/多卡训练单卡训练 Cityscapes R-50-D8 40k 版本python tools/train.py configs/isanet/isanet_r50-d8_4xb2-40k_cityscapes-512x1024.py多卡分布式训练例如 4 卡注意配置名中的4xb2即 4 卡 × 每卡 2 样本bash tools/dist_train.sh configs/isanet/isanet_r50-d8_4xb4-160k_ade20k-512x512.py 4训练进度与指标会按 configs/base/schedules/schedule_40k.py 中val_interval4000的节奏在验证集上评测并每 4000 次迭代保存 checkpoint。测试与多尺度翻转评测用tools/test.py复现表中mIoU(msflip)口径python tools/test.py configs/isanet/isanet_r50-d8_4xb2-80k_cityscapes-512x1024.py \ /path/to/checkpoint.pth --tta--tta启用 Test Time Augmentation多尺度 水平翻转对应结果表中的mIoU(msflip)列不加--tta则对应mIoU列。单张图片推理使用 demo/image_demo.py 对单张图片做可视化分割python demo/image_demo.py demo/demo.png \ configs/isanet/isanet_r50-d8_4xb2-80k_cityscapes-512x1024.py \ /path/to/checkpoint.pth --device cuda --out-file result.png也可以使用 demo/image_demo_with_inferencer.py 体验 MMSegmentation 1.x 的Inferencer高层推理接口无需手工拼装 pipeline。用测试用例验证 ISAHead 的正确性仓库为 ISAHead 提供了单元测试 tests/test_models/test_heads/test_isa_head.py其验证逻辑是构造输入torch.randn(1, 8, 23, 23)实例化ISAHead(in_channels8, channels4, num_classes19, isa_channels4, down_factor(8, 8))断言输出形状为(1, 19, 23, 23)。这个用例值得注意的两点非整除尺寸的 padding 路径输入宽高为 23而down_factor(8, 8)23 不能被 8 整除因此恰好覆盖了forward中F.pad对称补齐与随后裁剪的代码分支输出空间分辨率保持不变cls_seg采用不改变空间尺寸的卷积上采样align_cornersFalse故输出与输入同分辨率说明 ISA 注意力整体是像素级保形变换可无缝嵌入任意 Encoder-Decoder 分割框架。运行该测试pytest tests/test_models/test_heads/test_isa_head.py -v参数调优指南down_factor 与 isa_channels 的影响down_factor(8, 8)是官方默认的局部组尺寸。从 isa_head.py 可看出它直接决定glb_h, glb_w的取值进而决定全局/局部两个注意力分支各自作用的像素集合大小。增大down_factor会使局部组变大、全局组变小注意力更偏向局部精细建模减小则反之。当特征图尺寸不能被down_factor整除时代码会自动对称 padding不会崩溃。isa_channels256是 key/query 变换的输出通道数即自注意力的“隐维度”。从 self_attention_block.py 看它同时参与matmul_norm的channels ** -0.5缩放直接影响注意力分数分布增大它可提升模型容量代价是更多参数与显存。骨干channels512与in_channels2048需要匹配具体骨干R-50/R-101 的最后阶段输出为 2048 通道若替换其他骨干需同步调整。引用复现或扩展 ISANet 实验时请引用原论文与对应的正式发表版本以下 BibTeX 条目来自 configs/isanet/README.mdarticle{huang2019isa, title{Interlaced Sparse Self-Attention for Semantic Segmentation}, author{Huang, Lang and Yuan, Yuhui and Guo, Jianyuan and Zhang, Chao and Chen, Xilin and Wang, Jingdong}, journal{arXiv preprint arXiv:1907.12273}, year{2019} }该技术报告同时以 OCNet 期刊版形式发表于International Journal of Computer Visionarticle{yuan2021ocnet, title{OCNet: Object Context for Semantic Segmentation}, author{Yuan, Yuhui and Huang, Lang and Guo, Jianyuan and Zhang, Chao and Chen, Xilin and Wang, Jingdong}, journal{International Journal of Computer Vision}, pages{1--24}, year{2021}, publisher{Springer} }小结本文围绕 configs/isanet/README.md 完整展开先从论文视角解释了交错稀疏自注意力“两次稀疏分解、全局与局部串联”的核心思想再从 isa_head.py 源码逐行剖析了 global/local 两个注意力分支的交错数据流与 padding 处理随后给出 configs/isanet 全部 14 个配置的继承关系与关键参数以及 Cityscapes / ADE20K / VOC 的完整官方结果表最后提供了训练、多尺度翻转测试、单图推理与单元测试的实战命令。对于需要在高分辨率输入上做全局上下文建模、又受限于算力的分割任务ISAHead 是兼顾效果与效率的务实选择。【免费下载链接】mmsegmentationOpenMMLab Semantic Segmentation Toolbox and Benchmark.项目地址: https://gitcode.com/GitHub_Trending/mm/mmsegmentation创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考