MMSegmentation 中 Twins 视觉 Transformer 骨干网络全解析GSA/LSA 空间注意力设计与语义分割实战【免费下载链接】mmsegmentationOpenMMLab Semantic Segmentation Toolbox and Benchmark.项目地址: https://gitcode.com/GitHub_Trending/mm/mmsegmentationTwinsTwins-PCPVT 与 Twins-SVT是美团 AutoML 团队提出的视觉 Transformer 骨干网络其核心贡献在于重新审视了空间注意力Spatial Attention的设计仅用简单且高效的注意力机制即可在分类、检测与分割等稠密预测任务上取得出色表现。本文以 MMSegmentation 官方configs/twins目录为线索完整讲解 Twins 两种变体的架构原理、源码级实现、官方预训练权重转换脚本以及 ADE20K 语义分割的完整训练配置与结果帮助你直接将该骨干应用到自己的分割任务中。引言为什么重新设计空间注意力Twins 的论文《Twins: Revisiting the Design of Spatial Attention in Vision Transformers》指出尽管各类面向稠密预测任务的视觉 Transformer 架构层出不穷但空间注意力的设计才是决定其成败的关键。该工作重新审视了这一设计并证明一个精心设计但足够简单的空间注意力机制能够与当时最前沿的方案媲美甚至超越。由此提出了两种架构Twins-PCPVT基于 PVT 改进使用全局子采样注意力Global Sub-sampled Attention, GSATwins-SVT交替使用局部分组自注意力Locally-grouped Self Attention, LSA与 GSA。这两种架构实现简洁、效率高只涉及在现代深度学习框架中被高度优化的矩阵乘法并在图像分类、目标检测、语义分割等多项视觉任务上取得了优秀表现因此非常适合作为通用骨干网络。MMSegmentation 已将其完整集成实现文件位于 mmseg/models/backbones/twins.py官方复现仓库为 Meituan-AutoML/Twins。Twins 的核心架构模块GSA全局子采样注意力Spatial Reduction Attention全局子采样注意力用于减少全局注意力在高分辨率特征图上的计算量。在 MMSegmentation 中GlobalSubsampledAttention 直接继承自mmseg.models.backbones.mit.py中的EfficientMultiheadAttention两者实现并无差异只是按论文命名被重命名为 GSA。其关键超参数如下参数默认值含义embed_dims必填嵌入维度num_heads必填并行注意力头数attn_drop0.0注意力权重上的 Dropoutproj_drop0.0输出投影后的 Dropoutbatch_firstTrue输入形状为 (batch, n, embed_dims)qkv_biasTrueqkv 是否带偏置sr_ratio1GSA 的空间缩减比例即 PVT 中的 spatial reduction ratiosr_ratio越大参与注意力的 key/value 数量越少计算复杂度越低。PCPVT 各阶段的默认sr_ratios[8, 4, 2, 1]即浅层特征图较大时缩减更激进深层逐步恢复为完整注意力。LSA局部分组自注意力LocallyGroupedSelfAttention 将特征图划分为互不重叠的局部窗口在窗口内计算自注意力从而以线性复杂度建模局部关系。其实现要点见forward将 token 序列 reshape 回(b, h, w, c)的空间形式将特征图 padding 到window_size的整数倍生成注意力 mask屏蔽 padding 区域在每个window_size × window_size的窗口内独立计算多头自注意力还原并裁掉 padding 区域。LSA 的关键参数为window_size默认 7窗口越大局部感受野越大。CPE条件位置编码与 ViT 使用固定或可学习的位置编码不同Twins 采用 ConditionalPositionEncoding来自论文Conditional Positional Encodings for Vision Transformers它是一个3×3 的深度可分离卷积groupsembed_dims直接作用在 token 展开后的特征图上以卷积方式注入与输入内容相关的条件位置信息并天然支持任意输入分辨率。从 PCPVT.forward 可以看到每个 stage 的第一个 encoder layer 之后会施加一次 CPE。GSAEncoderLayer 与 LSAEncoderLayer两种编码层结构一致均为经典的 Pre-LN Transformer blockLN → Attention → DropPath 残差 → LN → FFN → DropPath 残差。其中 FFN 采用num_fcs2、GELU 激活、mlp_ratios控制隐藏层宽度具体见 GSAEncoderLayer 与 LSAEncoderLayer 的forwardx x self.drop_path(self.attn(self.norm1(x), hw_shape)) x x self.drop_path(self.ffn(self.norm2(x)))MMSegmentation 中的实现PCPVT 与 SVTPCPVT 骨干PCPVT 通过MODELS.register_module()注册为 MMSegmentation 骨干。它包含 4 个 stage各 stage 由PatchEmbed → GSAEncoderLayer × depths[i] → CPE组成关键构造参数与默认值参数默认值说明in_channels3输入通道数embed_dims[64, 128, 256, 512]各 stage 嵌入维度patch_sizes[4, 2, 2, 2]PatchEmbed 卷积核大小strides[4, 2, 2, 2]PatchEmbed 步长下采样率num_heads[1, 2, 4, 8]各 stage 注意力头数mlp_ratios[4, 4, 4, 4]FFN 隐藏维度相对嵌入维度的倍数depths[3, 4, 6, 3]各 stage 编码层数量sr_ratios[8, 4, 2, 1]各 stage 的 GSA 空间缩减比例out_indices(0, 1, 2, 3)输出哪些 stage 的特征drop_path_rate0.0随机深度衰减率stochastic depthnorm_after_stageFalse每个 stage 结束后是否额外加 LN此外PCPVT 支持pretrained参数新代码建议改用init_cfgdict(typePretrained, checkpoint...)并在init_weights中按 Lineartrunc_normal、Normconstant、Conv2dkaiming三种方式初始化权重。SVT 骨干SVT 继承自 PCPVT区别在于SVT 仅使用 3 个 stage默认depths[4, 4, 4]、sr_ratios[4, 2, 1]、num_heads[1, 2, 4]且在每个 stage 内交替使用 LSA 与 GSA——偶数层i % 2 0被替换为 LSAEncoderLayer奇数层保持 GSAEncoderLayer从而以局部-全局交替的方式覆盖完整感受野。SVT 的windiow_sizes源码中保留的拼写控制 LSA 窗口大小默认[7, 7, 7]norm_after_stage默认 True。官方预训练权重转换twins2mmseg.pyMMSegmentation 已提供从官方仓库转换好的预训练权重如果你希望自行转换官方权重可以使用仓库自带的转换脚本 tools/model_converters/twins2mmseg.pypython tools/model_converters/twins2mmseg.py ${PRETRAIN_PATH} ${STORE_PATH} ${MODEL_TYPE}${PRETRAIN_PATH}官方 pcpvt 或 svt 预训练权重路径或 URL${STORE_PATH}转换后权重的保存路径${MODEL_TYPE}pcpvt或svt决定部分键名映射规则。示例将官方alt_gvt_base.pth转换为 MMSegmentation 风格的 SVT 权重python tools/model_converters/twins2mmseg.py ./alt_gvt_base.pth ./pretrained/alt_gvt_base.pth svt脚本的核心转换逻辑见convert_twins包括丢弃分类头head.*参数patch_embeds.*.proj.*→patch_embeds.*.projection.*blocks.*→layers.*将官方分离的attn.q与attn.kv拼接为 MMSegmentation 的attn.in_proj_权重维度顺序为 q、k、vmlp.fc1→ffn.layers.0.0mlp.fc2→ffn.layers.1pos_block→position_encodings仅对 SVT奇数层LSA 层的attn.proj→attn.out_proj偶数层GSA 层保持不变。脚本支持 timm 风格含state_dict键与直接权重两种输入加载时统一使用 CPU转换结果通过torch.save输出。语义分割配置实战以 ADE20K 为例configs/twins目录共提供 12 个训练配置覆盖 PCPVT-S/B/L 与 SVT-S/B/L 六种骨干 × FPN/UPerNet 两种解码头全部针对 ADE20K150 类。文件名遵循twins_{backbone}_{head}_8xb{2|4}-{80k|160k}_ade20k-512x512.py的命名约定。配置 1SVT-S FPNFPNHead以 twins_svt-s_fpn_fpnhead_8xb4-80k_ade20k-512x512.py 为例它继承twins_pcpvt-s_fpn模型基座、ade20k数据集、默认运行配置与 80k 调度然后覆写骨干与解码头checkpoint https://download.openmmlab.com/mmsegmentation/v0.5/pretrain/twins/alt_gvt_small_20220308-7e1c3695.pth model dict( data_preprocessordata_preprocessor, backbonedict( typeSVT, init_cfgdict(typePretrained, checkpointcheckpoint), embed_dims[64, 128, 256, 512], num_heads[2, 4, 8, 16], mlp_ratios[4, 4, 4, 4], depths[2, 2, 10, 4], windiow_sizes[7, 7, 7, 7], norm_after_stageTrue), neckdict(in_channels[64, 128, 256, 512], out_channels256, num_outs4), decode_headdict(num_classes150), ) optim_wrapper dict( _delete_True, typeOptimWrapper, optimizerdict(typeAdamW, lr0.0001, weight_decay0.0001), clip_gradNone)注意此处 SVT-S 实际使用 4 个 stagedepths[2, 2, 10, 4]的小号配置与论文 SVT 默认的 3 stage 略有差异应以各配置文件实际值为准。配置 2PCPVT-S UPerNetUPerHeadtwins_pcpvt-s_uperhead_8xb4-160k_ade20k-512x512.py 演示了 UPerNet 解码头的完整训练设置包括线性预热 Poly 衰减调度与分组权重衰减optim_wrapper dict( _delete_True, typeOptimWrapper, optimizerdict( typeAdamW, lr0.00006, betas(0.9, 0.999), weight_decay0.01), paramwise_cfgdict(custom_keys{ pos_block: 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, ) ]关键点CPEpos_block与所有 norm 层不参与权重衰减decay_mult0.这是 Transformer 骨干迁移训练中的常见实践前 1500 次迭代使用线性预热避免早期不稳定。模型基座配置详解两种解码头的公共骨干配置分别沉淀在 configs/base/models/twins_pcpvt-s_fpn.py 与 configs/base/models/twins_pcpvt-s_upernet.py 中。以 FPN 基座为例model dict( typeEncoderDecoder, data_preprocessordict( typeSegDataPreProcessor, mean[123.675, 116.28, 103.53], std[58.395, 57.12, 57.375], bgr_to_rgbTrue, pad_val0, seg_pad_val255), backbonedict( typePCPVT, init_cfgdict(typePretrained, checkpointcheckpoint), in_channels3, embed_dims[64, 128, 320, 512], num_heads[1, 2, 5, 8], patch_sizes[4, 2, 2, 2], strides[4, 2, 2, 2], mlp_ratios[8, 8, 4, 4], out_indices(0, 1, 2, 3), qkv_biasTrue, depths[3, 4, 6, 3], sr_ratios[8, 4, 2, 1], drop_path_rate0.2), neckdict( typeFPN, in_channels[64, 128, 320, 512], out_channels256, num_outs4), decode_headdict( typeFPNHead, in_channels[256, 256, 256, 256], in_index[0, 1, 2, 3], feature_strides[4, 8, 16, 32], channels128, dropout_ratio0.1, num_classes150, loss_decodedict( typeCrossEntropyLoss, use_sigmoidFalse, loss_weight1.0)), test_cfgdict(modewhole))UPerNet 基座则使用UPerHeadpool_scales(1, 2, 3, 6)、channels512并额外挂载FCNHead辅助头loss_weight0.4。这套_base_继承机制意味着换骨干如 PCPVT-B/L、SVT时只需在新配置中覆写backbone与neck/decode_head的通道数即可无需重写整份模型定义。训练与测试配置就绪后可沿用 MMSegmentation 的标准流程训练与测试# 8 卡训练80k/160k 配置对应的调度在配置文件名中体现 bash tools/dist_train.sh configs/twins/twins_svt-s_fpn_fpnhead_8xb4-80k_ade20k-512x512.py 8 # 单卡测试并评估 mIoU python tools/test.py configs/twins/twins_svt-s_fpn_fpnhead_8xb4-80k_ade20k-512x512.py ${CHECKPOINT} --eval mIoU其中${CHECKPOINT}可替换为上方表格中对应的权重链接或本地转换得到的权重。ADE20K 实验结果与模型库下表汇总了configs/twins/README.md记录的 ADE20K 分割结果512×512 裁剪、V100 设备msflip为多尺度 翻转测试。完整模型元数据训练数据、批大小、算力资源、权重与日志链接可查阅 configs/twins/metafile.yaml。MethodBackboneCrop SizeLr schdMem (GB)Inf time (fps)DevicemIoUmIoU(msflip)FPNTwins-PCPVT-S512x512800006.6027.15V10043.2644.11UPerNetTwins-PCPVT-S512x5121600009.6714.24V10046.0446.92FPNTwins-PCPVT-B512x512800008.4119.67V10045.6646.48UPerNetTwins-PCPVT-B (8x2)512x5121600006.4612.04V10047.9148.64FPNTwins-PCPVT-L512x5128000010.7814.32V10045.9446.70UPerNetTwins-PCPVT-L (8x2)512x5121600007.8210.70V10049.3550.08FPNTwins-SVT-S512x512800005.8029.79V10044.4745.42UPerNetTwins-SVT-S (8x2)512x5121600004.9315.09V10046.0846.96FPNTwins-SVT-B512x512800008.7521.10V10046.7747.47UPerNetTwins-SVT-B (8x2)512x5121600006.7712.66V10048.0448.87FPNTwins-SVT-L512x5128000011.2017.80V10046.5547.74UPerNetTwins-SVT-L (8x2)512x5121600008.4110.73V10049.6550.63表格阅读要点8x2表示 8 块 GPU、每块 2 个样本Twins 在 ADE20K 上的默认训练配置为 8 卡 × 4 样本即 8xb4。减小每卡 batch 的配置8x2可在显存受限时降低单卡显存占用如 PCPVT-L 的 UPerNet 由默认需求的显存降至 7.82 GB。表中UPerNet与FPN分别对应UPerHead与FPNHead两种解码头官方仓库中的 Twins 模型全部采用UPerHead。上表最后一行 UPerNet SVT-L 对应的实际配置文件为twins_svt-l_uperhead_8xb2-160k_ade20k-512x512.py。从结果趋势看UPerNet160k 迭代整体优于 FPN80k 迭代SVT-L UPerNet 在 ADE20K 上取得了该系列最高 mIoU49.65 / msflip 50.63同时多尺度 翻转测试普遍带来约 1 个点的稳定增益。在自己的数据集上使用 Twins将 Twins 用于自定义数据集只需三步参照 configs/base/models/twins_pcpvt-s_fpn.py 保留骨干配置将decode_head.num_classes与auxiliary_head.num_classes改为你的类别数替换数据集_base_为自定义数据集配置或在配置文件中直接覆写train_dataloader/val_dataloader通过init_cfg加载上述转换后的官方预训练权重即可开始迁移训练。需要注意的是Twins 骨干输出 4 个 stage 的特征图stride 分别为 4/8/16/32因此无论是 FPN 还是 UPerNet 解码头其in_channels都必须与骨干的embed_dims一一对应PCPVT-S 为[64, 128, 320, 512]。引用若在研究中使用了 Twins请按官方文档引用article{chu2021twins, title{Twins: Revisiting spatial attention design in vision transformers}, author{Chu, Xiangxiang and Tian, Zhi and Wang, Yuqing and Zhang, Bo and Ren, Haibing and Wei, Xiaolin and Xia, Huaxia and Shen, Chunhua}, journal{arXiv preprint arXiv:2104.13840}, year{2021} }小结Twins 以简单而高效的空间注意力设计切入视觉 Transformer 骨干通过 GSA 捕获全局上下文、LSA 建模局部细节、CPE 提供条件位置信息在语义分割等稠密预测任务上展现了良好的精度与效率平衡。在 MMSegmentation 中你可以通过 mmseg/models/backbones/twins.py 阅读其完整实现使用 tools/model_converters/twins2mmseg.py 复用官方预训练权重并直接套用configs/twins下的 12 份 ADE20K 配置快速开展实验——无论是替换解码头FPN/UPerNet、调整骨干规格S/B/L还是迁移到自定义数据集Twins 都提供了开箱即用的完整链路。【免费下载链接】mmsegmentationOpenMMLab Semantic Segmentation Toolbox and Benchmark.项目地址: https://gitcode.com/GitHub_Trending/mm/mmsegmentation创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考