MMSegmentation 中的 Swin Transformer分层窗口注意力骨干网络配置、权重转换与语义分割实践【免费下载链接】mmsegmentationOpenMMLab Semantic Segmentation Toolbox and Benchmark.项目地址: https://gitcode.com/GitHub_Trending/mm/mmsegmentationSwin Transformer 是一种采用移位窗口Shifted Windows方案的分层视觉 Transformer可作为通用骨干网络用于图像分类与密集预测任务。本文以 MMSegmentation 仓库中的 configs/swin/README.md 为骨架结合 SwinTransformer 骨干实现、权重转换脚本 与各变体配置文件系统讲解其在语义分割任务中的架构原理、预训练权重获取、UPerNet 解码头配置、优化器与学习率调度细节以及训练、测试与结果复现方法。一、Swin Transformer 核心思想与论文要点Swin Transformer 由 Liu 等人提出论文Swin Transformer: Hierarchical Vision Transformer using Shifted WindowsarXiv:2103.14030其核心动机是解决 Transformer 从语言迁移到视觉时面临的两大差异视觉实体的尺度变化巨大以及图像像素分辨率远高于文本词元。针对上述问题论文提出了两大设计层次化架构Hierarchical Architecture与 ViT 输出单一分辨率的特征不同Swin Transformer 通过逐级 Patch Merging 下采样产出多尺度特征图这与 CNN 骨干如 ResNet的 1/4、1/8、1/16、1/32 多级特征完全兼容因此可以无缝嵌入 FPN、UPerNet 等需要多级特征的密集预测框架。移位窗口注意力Shifted Window Attention将自注意力计算限制在互不重叠的局部窗口内同时通过相邻层间窗口的周期性移位建立跨窗口连接。该方案将自注意力的计算复杂度从图像尺寸的二次方降为线性并显著提升了计算效率。论文在 ImageNet-1K 分类top-1 87.3、COCO 目标检测test-dev 上 58.7 box AP 与 51.1 mask AP以及 ADE20K 语义分割val 上 53.5 mIoU等任务上取得了当时领先的结果并证明了分层设计与移位窗口方案对全 MLP 架构同样有益。仓库中 configs/swin/metafile.yaml 收录了 MMSegmentation 官方复现的 Swin 系列分割模型记录可在model-index.yml与dataset-index.yml中检索对应权重。二、源码级解析MMSegmentation 中的 SwinTransformer 骨干实现MMSegmentation 在 mmseg/models/backbones/swin.py 中完整实现了 Swin Transformer 骨干并在MODELS.register_module()下注册为SwinTransformer。整个实现由以下模块层层组装而成1. WindowMSA窗口内多头自注意力WindowMSA 实现了窗口内的多头自注意力W-MSA关键设计包括相对位置偏置表relative_position_bias_table是一个形状为(2*Wh-1) * (2*Ww-1), num_heads的可学习参数表通过relative_position_index索引后加到注意力分数上使模型具备平移等变性qkv 合并线性层通过单个nn.Linear(embed_dims, embed_dims * 3)一次性计算 q、k、v再reshape与permute拆分为多头缩放因子self.scale qk_scale or head_embed_dims**-0.5即默认按 head 维度的负二分之一次方缩放。2. ShiftWindowMSA移位窗口自注意力ShiftWindowMSA 是 Swin 的核心创新点其forward过程清晰展示了移位窗口的完整流程Padding将特征图补齐到窗口大小的整数倍pad_r、pad_bCyclic Shift使用torch.roll沿 H、W 方向平移-shift_size实现窗口的循环移位注意力掩码生成当shift_size 0时构造 3×3 区域的img_mask进而生成形状为(nW, window_size*window_size, window_size*window_size)的attn_mask非同一区域的注意力被masked_fill为 -100.0窗口划分与 W-MSA通过window_partition将特征划分为nW*B个窗口送入WindowMSAshift_size0时退化为普通窗口注意力掩码为None窗口还原与反向移位window_reverse恢复特征图再torch.roll移回原位最后裁掉 padding。在 SwinBlock 中ShiftWindowMSA的shift_size被设定为window_size // 2偶数块 shiftTrue 时奇数块则使用普通窗口shiftFalse两者交替出现形成跨窗口信息流动。块内结构为 Norm → W-MSA/SW-MSA → 残差 → Norm → FFN → 残差FFN 通过FFN组件构建默认隐藏层维度为mlp_ratio * embed_dims默认 4 倍。3. SwinBlockSequence 与 Patch Merging四阶段层次化下采样SwinBlockSequence 代表一个 stage内部串行堆叠depth个SwinBlock偶数索引块关闭移位、奇数索引块开启移位每个 stage 末尾可选挂接downsample模块最后一个 stage 不设置。SwinTransformer主干在__init__中依次构建PatchEmbed使用kernel_sizepatch_size、stridestrides[0]默认 4的无重叠卷积进行 Patch Embedding要求strides[0] patch_size测试用例 test_swin.py 专门断言了这一点四层 stage深度depths(2, 2, 6, 2)注意力头数num_heads(3, 6, 12, 24)通道数依次为 96/192/384/768逐级通过PatchMerging减半分辨率并加倍通道随机深度Stochastic Depthdrop_path_rate按torch.linspace(0, drop_path_rate, total_depth)在全部 12 个块上线性递增分配输出层out_indices默认(0, 1, 2, 3)每个输出 stage 之后接一个 LayerNorm输出形状为(B, C, H/stride, W/stride)的多级特征预训练权重加载init_weights支持Pretrained初始化自动处理backbone./module.前缀剥离、绝对位置编码 reshape以及相对位置偏置表在窗口尺寸不同时的 bicubic 插值重采样见 swin.py。测试文件 tests/test_models/test_backbones/test_swin.py 覆盖了 SwinBlock 与完整骨干的前向形状如 256×256 输入输出(1,96,64,64)、(1,192,32,32)、(1,384,16,16)、(1,768,8,8)、非规整尺寸输入、with_cp梯度检查点、frozen_stages冻结以及绝对位置编码冻结等场景。三、预训练权重与官方权重转换脚本MMSegmentation 已提供由 Swin Transformer 官方仓库转换而来的预训练权重。同时仓库在tools/model_converters目录下提供了转换脚本 swin2mmseg.py允许用户自行将官方仓库权重转换为 MMSegmentation 风格。1. 脚本用法python tools/model_converters/swin2mmseg.py ${PRETRAIN_PATH} ${STORE_PATH}其中PRETRAIN_PATH为源模型路径或 URLSTORE_PATH为转换后权重的保存路径。例如python tools/model_converters/swin2mmseg.py https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window7_224.pth pretrain/swin_base_patch4_window7_224.pth2. 转换脚本的键名映射逻辑从源码看convert_swin 完成了以下几类键名映射丢弃分类头k.startswith(head)的键直接跳过注意力attn.→attn.w_msa.FFNmlp.fc1.→ffn.layers.0.0.mlp.fc2.→ffn.layers.1.Patch Embeddingproj→projectionstage 命名layers→stages仅首次出现处替换关键数值重排Patch Merging 的reduction权重与norm偏置需要按通道顺序重排correct_unfold_reduction_order与correct_unfold_norm_order将第 2、3 通道交换顺序这正是官方实现与 MMSegmentation 实现之间窗口划分顺序差异导致的必须转换才能正确加载。脚本入口通过CheckpointLoader.load_checkpoint加载支持 URL兼容state_dict、model键包裹的 checkpoint最后以torch.save存储到目标路径。3. 官方预训练权重对照表默认设置下以下预训练权重与官方原始模型一一对应下载后建议存放于pretrain/目录转换后权重MMSegmentation官方原始权重SwinTransformer storagepretrain/swin_tiny_patch4_window7_224.pthswin_tiny_patch4_window7_224.pthpretrain/swin_small_patch4_window7_224.pthswin_small_patch4_window7_224.pthpretrain/swin_base_patch4_window7_224.pthswin_base_patch4_window7_224.pthpretrain/swin_base_patch4_window7_224_22k.pthswin_base_patch4_window7_224_22k.pthpretrain/swin_base_patch4_window12_384.pthswin_base_patch4_window12_384.pthpretrain/swin_base_patch4_window12_384_22k.pthswin_base_patch4_window12_384_22k.pth注意命名中的patch4表示 Patch 大小为 4window7/window12表示窗口大小为 7/12224/384为预训练输入分辨率22k表示在 ImageNet-22K 上预训练。四、配置文件详解Swin UPerNet 的语义分割实践configs/swin/目录下提供了 9 个基于 UPerNet 解码头的 ADE20K 训练配置与 1 个 LEVIR 变化检测配置。以基础配置 swin-tiny-patch4-window7-in1k-pre_upernet_8xb2-160k_ade20k-512x512.py 为例其组成与关键参数如下1. 配置继承结构_base_ [ ../_base_/models/upernet_swin.py, ../_base_/datasets/ade20k.py, ../_base_/default_runtime.py, ../_base_/schedules/schedule_160k.py ]即复用 upernet_swin.py 的模型定义、ADE20K 数据集定义、默认运行配置与 160k 迭代训练计划。2. 骨干网络与解码头配置crop_size (512, 512) data_preprocessor dict(sizecrop_size) checkpoint_file https://download.openmmlab.com/mmsegmentation/v0.5/pretrain/swin/swin_tiny_patch4_window7_224_20220317-1cdeb081.pth model dict( data_preprocessordata_preprocessor, backbonedict( init_cfgdict(typePretrained, checkpointcheckpoint_file), embed_dims96, depths[2, 2, 6, 2], num_heads[3, 6, 12, 24], window_size7, use_abs_pos_embedFalse, drop_path_rate0.3, patch_normTrue), decode_headdict(in_channels[96, 192, 384, 768], num_classes150), auxiliary_headdict(in_channels384, num_classes150))对照 upernet_swin.py 中的完整默认值各参数含义如下pretrain_img_size224预训练输入分辨率用于绝对位置编码的 patch 数计算patch_size4与strides(4, 2, 2, 2)Patch Embedding 与各 stage 下采样步长window_size7窗口尺寸决定相对位置偏置表大小mlp_ratio4FFN 隐藏层为嵌入维度的 4 倍out_indices(0, 1, 2, 3)输出四个 stage 的多尺度特征qkv_biasTrue、qk_scaleNone注意力 qkv 偏置与缩放patch_normTruePatch Embedding/Merging 后接 LayerNormuse_abs_pos_embedFalseSwin 默认不使用绝对位置编码相对位置偏置已足够drop_path_rate0.3随机深度率在 12 个块上线性递增分配解码头UPerHead接收in_channels[96, 192, 384, 768]四层特征pool_scales(1, 2, 3, 6)构成金字塔池化辅助头FCNHead取in_index2第 3 层 384 通道特征ADE20K 数据集共 150 类因此num_classes150。3. 优化器与学习率调度Swin 的标准训练配方optim_wrapper dict( _delete_True, typeOptimWrapper, optimizerdict( typeAdamW, lr0.00006, betas(0.9, 0.999), weight_decay0.01), paramwise_cfgdict( custom_keys{ absolute_pos_embed: dict(decay_mult0.), relative_position_bias_table: dict(decay_mult0.), norm: dict(decay_mult0.) })) param_scheduler [ dict( typeLinearLR, start_factor1e-6, by_epochFalse, begin0, end1500), dict( typePolyLR, eta_min0.0, power1.0, begin1500, end160000, by_epochFalse, ) ]训练配方要点优化器为AdamW学习率6e-5权重衰减 0.01_delete_True表示覆盖基配置中的 SGD 优化器参数分组位置编码absolute_pos_embed、相对位置偏置表relative_position_bias_table与所有norm层权重不参与权重衰减decay_mult0.这与 BERT 等 Transformer 的常见做法一致调度策略前 1500 次迭代用LinearLR线性预热起始因子1e-6之后切换为PolyLRpower1.0即线性衰减直至 160000 次迭代结束训练采用 8 卡 × 每卡 2 张图train_dataloader dict(batch_size2)总 batch size 16验证/测试每卡 1 张。4. 各变体配置的差异目录下其余配置通过继承与局部覆盖实现差异化Swin-S / Swin-B在 upernet_swin.py 基础上仅覆盖embed_dims、depths、num_heads与decode_head/auxiliary_head的in_channelsImageNet-22K 预训练如 swin-base-patch4-window7-in22k-pre_upernet_8xb2-160k_ade20k-512x512.py仅覆盖checkpoint_file为 22K 版本权重window12 / 384×384 预训练如 swin-base-patch4-window12-in22k-384x384-pre_upernet_8xb2-160k_ade20k-512x512.py窗口尺寸与预训练分辨率不同权重加载时相对位置偏置表会自动插值Swin-Largeconfigs/swin/下还提供swin-large-patch4-window7-in22k-pre与swin-large-patch4-window12-in22k-384x384-pre两个大模型配置遥感变化检测示例 swin-tiny-patch4-window7_upernet_1xb8-20k_levir-256x256.py 展示了 Swin 在 LEVIR 双时相遥感数据上的扩展——in_channels6两期影像拼接、bgr_to_rgbFalse、num_classes2说明该骨干可灵活适配多输入通道的专属任务。五、ADE20K 语义分割结果官方在 ADE20K val 上的复现结果UPerNet Swin均基于 8×V100、160k 迭代训练如下表中配置与权重均可从 configs/swin/metafile.yaml 及model-index.yml检索方法骨干Crop Size预训练预训练分辨率Batch SizeLr schdMem (GB)Inf time (fps)DevicemIoUmIoU(msflip)UPerNetSwin-T512×512ImageNet-1K224×224161600005.0221.06V10044.4145.79UPerNetSwin-S512×512ImageNet-1K224×224161600006.1714.72V10047.7249.24UPerNetSwin-B512×512ImageNet-1K224×224161600007.6112.65V10047.9949.57UPerNetSwin-B512×512ImageNet-22K224×22416160000--V10050.1351.9UPerNetSwin-B512×512ImageNet-1K384×384161600008.5212.10V10048.3549.65UPerNetSwin-B512×512ImageNet-22K384×38416160000--V10050.7652.4从上表可以清晰看到两个规律其一骨干规模从 Swin-T 到 Swin-B 带来 mIoU 稳步提升44.41 → 47.99其二使用 ImageNet-22K 预训练或更高分辨率384×384预训练权重可带来显著的精度增益最高 50.76 mIoU / 52.4 msflip。注意表中的显存与推理速度均为 V100 上的实测参考值实际表现会随硬件、batch size 与输入尺寸变化。六、训练与评估实践在 MMSegmentation 中基于上述配置训练与评估的标准流程如下# 单机多卡训练 bash tools/dist_train.sh ${CONFIG_FILE} ${GPU_NUM} [optional arguments] # 例如使用 8 卡训练 Swin-T UPerNet bash tools/dist_train.sh configs/swin/swin-tiny-patch4-window7-in1k-pre_upernet_8xb2-160k_ade20k-512x512.py 8 # 测试需指定权重文件 bash tools/dist_test.sh ${CONFIG_FILE} ${CHECKPOINT_FILE} ${GPU_NUM} [--out ${RESULT_FILE}] [--eval mIoU]单卡环境可分别使用 tools/train.py 与 tools/test.py 直接运行python tools/train.py configs/swin/swin-tiny-patch4-window7-in1k-pre_upernet_8xb2-160k_ade20k-512x512.py python tools/test.py configs/swin/swin-tiny-patch4-window7-in1k-pre_upernet_8xb2-160k_ade20k-512x512.py ${CHECKPOINT_FILE} --eval mIoU几点实操提示训练前请先在pretrain/目录准备对应预训练权重参见第三节对照表或直接使用配置中内嵌的 OpenMMLab 权重 URL 自动下载由于默认配置面向 8 卡环境单卡训练需相应下调 batch size并可能按线性缩放规则调整学习率测试时可参考表中mIoU(msflip)的多尺度 翻转测试设置通常需要在配置或命令行中启用--tta相关选项以复现该指标若需复现论文中更大的 Swin-L 模型或更高的 640×640 输入可参考 upernet 与 swin 目录下的相关配置进行扩展。七、引用若在研究中使用了 Swin Transformer 或本仓库实现建议引用article{liu2021Swin, title{Swin Transformer: Hierarchical Vision Transformer using Shifted Windows}, author{Liu, Ze and Lin, Yutong and Cao, Yue and Hu, Han and Wei, Yixuan and Zhang, Zheng and Lin, Stephen and Guo, Baining}, journal{arXiv preprint arXiv:2103.14030}, year{2021} }结语本文围绕 MMSegmentation 中 Swin Transformer 的完整落地链路展开从移位窗口注意力的源码实现WindowMSA/ShiftWindowMSA/ 四阶段层次化下采样到官方权重转换脚本 swin2mmseg.py 的键名映射与通道重排细节再到 UPerNet 解码头、AdamW 优化器、参数分组与预热 Poly 衰减调度等训练配方最后给出 ADE20K 复现结果与训练评估命令。无论你是要复现论文结果还是将 Swin 作为骨干接入自己的分割任务如 LEVIR 变化检测示例都可以直接以 configs/swin 下的配置文件为起点进行修改与扩展。【免费下载链接】mmsegmentationOpenMMLab Semantic Segmentation Toolbox and Benchmark.项目地址: https://gitcode.com/GitHub_Trending/mm/mmsegmentation创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考