MMSegmentation 中的 SETR:以序列到序列视角重新思考语义分割的 Transformer 实现与实战
2026/9/15 19:04:52 网站建设 项目流程

MMSegmentation 中的 SETR:以序列到序列视角重新思考语义分割的 Transformer 实现与实战

【免费下载链接】mmsegmentationOpenMMLab Semantic Segmentation Toolbox and Benchmark.项目地址: https://gitcode.com/GitHub_Trending/mm/mmsegmentation

导读

SETR(SEgmentation TRansformer)是首个以"序列到序列(sequence-to-sequence)"视角重新思考语义分割任务的 Transformer 方法:它用纯 Transformer(无卷积、无分辨率逐步缩减)把整幅图像编码为 patch 序列,再配合极简解码头完成逐像素分类。本文基于 MMSegmentation 仓库中configs/setr/的官方配置与文档,系统讲解 SETR 的三种解码头设计(Naive / PUP / MLA)、完整模型配置解读、预训练权重的转换与加载流程,以及 ADE20K、Cityscapes 两个基准上的官方复现结果,帮助你在 MMSegmentation 中直接复现、微调与部署 SETR。

核心思想:把分割当作序列到序列预测

在 SETR 之前,主流语义分割方法几乎都是"全卷积网络(FCN)+ 编码器-解码器"架构:编码器通过卷积逐级降低空间分辨率,同时借助空洞卷积或注意力模块扩大感受野。SETR 论文提出了一个不同的视角——直接把语义分割建模为序列到序列预测任务

  • 使用一个纯 Transformer 编码器(不含卷积、不降低分辨率)将图像编码为 patch 序列;
  • Transformer 的每一层都建模全局上下文,因此该编码器只需搭配一个简单解码器,就能构成强大的分割模型,即 SETR;
  • 实验表明,SETR 在 ADE20K(50.28% mIoU)、Pascal Context(55.83% mIoU)上刷新当时最优,并在 Cityscapes 上取得有竞争力的结果,其论文提交当日在 ADE20K 测试榜上取得第一名。

注:上述精度为论文报告值;MMSegmentation 官方复现配置下的精度见下文"复现结果"一节,两者口径不同,请以仓库实测为准。

在 MMSegmentation 中,SETR 的官方实现文档为 configs/setr/README.md,模型实现位于mmseg/models/decode_heads/下的setr_up_head.pysetr_mla_head.py

三种解码头:Naive、PUP 与 MLA

SETR 的核心差异在于解码器。MMSegmentation 将 SETR 的两种解码头统一实现为SETRUPHead(mmseg/models/decode_heads/setr_up_head.py),对应论文中的Naive(朴素上采样)PUP(渐进式上采样)两种变体;第三种变体MLA(多级特征聚合,Multi-level feature Aggregation)则单独实现为SETRMLAHead(mmseg/models/decode_heads/setr_mla_head.py)。

Naive 与 PUP:SETRUPHead

SETRUPHead的关键构造参数(来自 setr_up_head.py):

参数默认值说明
norm_layerdict(type='LN', eps=1e-6, requires_grad=True)对 Transformer 输出的序列特征先做 LayerNorm(在(n, h*w, c)维度上)
num_convs1解码卷积的个数,每个卷积后接一次双线性上采样
up_scale4每次上采样的倍率
kernel_size3解码卷积核大小,代码断言只能是 1 或 3
init_cfgLN 常数初始化 +conv_seg正态初始化默认初始化策略

前向过程很直观(setr_up_head.py):

  1. 将 Transformer 输出的序列(n, c, h*w)reshape 回空间特征(n, c, h, w)
  2. 在序列维度上做 LayerNorm;
  3. 依次执行num_convs个"卷积 + 双线性上采样"模块;
  4. 最后通过cls_seg(1x1 卷积)输出逐像素 logits。
  • Naive 变体:解码只用1 个 1x1 卷积(kernel_size=1将 1024 维特征降到channels,再一次性up_scale=4上采样到原图 1/4,语义过于"朴素";
  • PUP 变体:解码采用4 个 3x3 卷积(kernel_size=3、每次up_scale=2渐进式上采样,逐级恢复分辨率,精度通常略高于 Naive。

MLA:SETRMLAHead

SETRMLAHead走的是另一条路线(setr_mla_head.py):

  • 从 ViT 的多个中间层(如第 5、11、17、23 层)抽取多级特征,输入通道数不再是单个 int,而是多级列表;
  • 每个输入先经过两组 3x3 卷积(降到mla_channels,默认 128)再 4 倍上采样,得到分辨率一致的多级特征;
  • 最后将多级特征在通道维torch.cat,再经cls_seg分类;
  • 代码约束channels == num_inputs * mla_channels,因此默认channels=512(4 级 × 128)。

同时,MLA 变体在 ViT 编码器输出端多了一个MLANeck(见 configs/base/models/setr_mla.py),先把四级1024通道特征统一压到 256,再喂给SETRMLAHead;辅助头则复用普通FCNHeadnum_convs=0concat_input=False,即纯 1x1 卷积分类)。

解码头与辅助头的分工

三种变体都采用"一个主解码头 + 三个辅助头"的结构,且共用同一个损失函数配置:

loss_decode=dict(type='CrossEntropyLoss', use_sigmoid=False, loss_weight=0.4)

主头loss_weight=1.0,三个辅助头各0.4;辅助头分别接在 ViT 的早期层输出上(in_index=0/1/2),为浅层特征提供监督信号,帮助训练收敛。

模型配置逐项解读

仓库在configs/setr/下提供了 7 个官方配置,覆盖两种数据集、三种变体。所有配置都通过_base_继承公共模块:

配置文件变体数据集输入尺寸迭代数
setr_vit-l_naive_8xb2-160k_ade20k-512x512.pyNaiveADE20K512×512160k
setr_vit-l_pup_8xb2-160k_ade20k-512x512.pyPUPADE20K512×512160k
setr_vit-l-mla_8xb1-160k_ade20k-512x512.pyMLAADE20K512×512160k
setr_vit-l_mla_8xb2-160k_ade20k-512x512.pyMLA(大 batch)ADE20K512×512160k
setr_vit-l_naive_8xb1-80k_cityscapes-768x768.pyNaiveCityscapes768×76880k
setr_vit-l_pup_8xb1-80k_cityscapes-768x768.pyPUPCityscapes768×76880k
setr_vit-l_mla_8xb1-80k_cityscapes-768x768.pyMLACityscapes768×76880k

以 PUP/ADE20K 为例,配置核心结构如下(setr_vit-l_pup_8xb2-160k_ade20k-512x512.py):

_base_ = [ '../_base_/models/setr_pup.py', '../_base_/datasets/ade20k.py', '../_base_/default_runtime.py', '../_base_/schedules/schedule_160k.py' ] crop_size = (512, 512) data_preprocessor = dict(size=crop_size) norm_cfg = dict(type='SyncBN', requires_grad=True) model = dict( data_preprocessor=data_preprocessor, pretrained=None, backbone=dict( img_size=(512, 512), drop_rate=0., init_cfg=dict(type='Pretrained', checkpoint='pretrain/vit_large_p16.pth')), decode_head=dict(num_classes=150), auxiliary_head=[...三个 SETRUPHead...], test_cfg=dict(mode='slide', crop_size=(512, 512), stride=(341, 341)), ) optimizer = dict(lr=0.001, weight_decay=0.0) optim_wrapper = dict( type='OptimWrapper', optimizer=optimizer, paramwise_cfg=dict(custom_keys={'head': dict(lr_mult=10.)})) train_dataloader = dict(batch_size=2) # 8 卡 × 2 = 16

几个值得注意的设计点:

  1. Backbone 是标准VisionTransformer(ViT-Large),公共模型定义见 configs/base/models/setr_pup.py:patch_size=16embed_dims=1024num_layers=24num_heads=16interpolate_mode='bilinear'。Naive/PUP 取out_indices=(9, 14, 19, 23)with_cls_token=True;MLA 取out_indices=(5, 11, 17, 23)with_cls_token=False(见 setr_mla.py)。
  2. Backbone 归一化用 LayerNorm,解码头归一化用 SyncBN:前者符合 Transformer 惯例,后者保证分割头在 BN 统计上的稳定性。
  3. 预训练权重通过init_cfg指定本地路径pretrain/vit_large_p16.pth),这就是下文要讲解的转换产物。
  4. 优化器差异化学习率paramwise_cfg=dict(custom_keys={'head': dict(lr_mult=10.)})使头部学习率是 backbone 的 10 倍(如 backbone lr=0.001 时 head 实际 lr=0.01),且weight_decay=0.0
  5. 测试采用滑窗(slide)模式:ADE20K 配置用crop_size=(512,512)stride=(341,341);Cityscapes 配置用 768×768 滑窗、stride=512,见 setr_vit-l_mla_8xb1-80k_cityscapes-768x768.py。这是因为 ViT 没有平移等变性,整图推理在超大图上效果欠佳,滑窗可显著提升精度。

三个变体的关键差异速查

变体解码头主头卷积上采样方式辅助头ViT 输出层cls token
NaiveSETRUPHead1×1 conv ×1一次 4× 双线性SETRUPHead(1×1, up×4)(9,14,19,23)保留
PUPSETRUPHead3×3 conv ×4四次 2× 渐进双线性SETRUPHead(3×3, up×4)(9,14,19,23)保留
MLASETRMLAHead3×3 conv ×2/级每级 4× 后 catFCNHead(1×1)(5,11,17,23)去掉

预训练权重准备:vit2mmseg.py 转换与加载

SETR 的 ViT-Large 骨干需要从 ImageNet 预训练权重初始化。官方文档给出的完整流程(configs/setr/README.md):

  1. 下载 timm 仓库发布的 ViT-Large 预训练权重(文件名jx_vit_large_p16_384-b3be5167.pth);
  2. 用转换脚本vit2mmseg.py将 timm 的权重键名转换为 MMSegmentation 风格:
python tools/model_converters/vit2mmseg.py ${PRETRAIN_PATH} ${STORE_PATH}

例如:

python tools/model_converters/vit2mmseg.py \ jx_vit_large_p16_384-b3be5167.pth pretrain/vit_large_p16.pth

该脚本读取PRETRAIN_PATH的权重并保存到STORE_PATH(tools/model_converters/vit2mmseg.py)。由于dst路径会被mmengine.mkdir_or_exist自动创建父目录,建议把转换产物统一放到pretrain/目录下,与各配置中init_cfg.checkpoint='pretrain/vit_large_p16.pth'的约定保持一致。

键名映射规则

转换核心逻辑在convert_vit函数(vit2mmseg.py),它把 timm/deit 风格的键名映射到 MMSegmentation 的VisionTransformer实现:

timm 键名片段MMSeg 键名片段
norm.(顶层)ln1.
patch_embed.projpatch_embed.projection
blocks.layers.
blocks.*.normlayers.*.ln
blocks.*.mlp.fc1layers.*.ffn.layers.0.0
blocks.*.mlp.fc2layers.*.ffn.layers.1
blocks.*.attn.qkvlayers.*.attn.attn.in_proj_
blocks.*.attn.projlayers.*.attn.attn.out_proj
head.*(分类头)直接丢弃

脚本同时兼容三种 checkpoint 格式:timm 的state_dict字段、deit 的model字段,以及直接就是权重字典的情况(vit2mmseg.py)。

转换完成后,配置中通过pretrained=None关闭模型默认预训练路径,改用backbone.init_cfg显式加载本地权重,从而避免与_base_中默认的远程预训练路径冲突。

复现结果与资源占用

ADE20K(150 类)

以下为 MMSegmentation 官方在 V100 上的复现结果(来自 configs/setr/README.md,mIoU 为单尺度 / 多尺度+翻转两种评测口径):

方法BackboneCrop SizeBatchLr schdMem (GB)Inf time (fps)mIoUmIoU(ms+flip)
SETR NaiveViT-L512×5121616000018.404.7248.2849.56
SETR PUPViT-L512×5121616000019.544.5048.2449.99
SETR MLAViT-L512×512816000010.96-47.3449.05
SETR MLAViT-L512×5121616000017.305.2547.3949.37

对应配置:Naive、PUP、MLA(bs=8)、MLA(bs=16)。MLA(bs=16) 仅是在 MLA(bs=8) 基础上把单卡 batch 从 1 提到 2(train_dataloader = dict(batch_size=2)),见 setr_vit-l_mla_8xb2-160k_ade20k-512x512.py。

Cityscapes(19 类)

方法BackboneCrop SizeBatchLr schdMem (GB)Inf time (fps)mIoUmIoU(ms+flip)
SETR NaiveViT-L768×76888000024.060.3978.1080.22
SETR PUPViT-L768×76888000027.960.3779.2181.02
SETR MLAViT-L768×76888000024.100.4177.0079.59

对应配置:Naive、PUP、MLA。Cityscapes 三个配置均以 8 卡 × 单卡 batch=1 训练,且学习率策略切换到schedule_80k

从表格可观察到两个规律:一是768×768 高分辨率对 Cityscapes 至关重要(ViT 无金字塔结构,输入分辨率直接决定特征密度);二是PUP 的渐进式上采样在两种数据集上都是最稳的选择,MLA 则凭借更少的内存占用(ADE20K 上 10.96GB vs 19.54GB)适合显存受限场景。

如何在 MMSegmentation 中运行 SETR

仓库为标准 OpenMMLab 工程,训练与测试可直接复用tools/下的通用脚本(需先安装依赖并准备好对应数据集):

# 单机多卡训练(以 PUP/ADE20K 为例) bash tools/dist_train.sh \ configs/setr/setr_vit-l_pup_8xb2-160k_ade20k-512x512.py 8 # 测试 bash tools/dist_test.sh \ configs/setr/setr_vit-l_pup_8xb2-160k_ade20k-512x512.py \ ${CHECKPOINT} 8

运行前需确保:

  1. 已按上文完成预训练权重转换,并把pretrain/vit_large_p16.pth放到工作目录(与配置中的init_cfg.checkpoint相对路径一致);
  2. ADE20K / Cityscapes 数据集按 configs/base/datasets/ade20k.py 或 cityscapes_768x768.py 约定的目录结构组织好;
  3. 显存满足上表需求(ViT-L 体量较大,单卡 batch 建议按配置保持 1~2)。

另外,仓库提供了模型仓库索引 metafile.yaml,记录了每个官方权重的下载地址与评测结果,可配合tools/下的推理脚本直接加载权重做单图分割演示。

小结

SETR 在 MMSegmentation 中的落地体现了三个可复用的工程范式:以纯 Transformer 作为分割骨干VisionTransformer+ LayerNorm)、多种解码头的可插拔设计(Naive / PUP 共用SETRUPHead,MLA 独立为SETRMLAHead)、以及timm → OpenMMLab 的权重键名转换工具链(vit2mmseg.py)。如果你希望在自己的数据上微调 SETR,最经济的路径是:基于 PUP 配置替换数据集与num_classes,复用 ADE20K 训练配置中的滑窗测试策略,即可快速得到基线。

引用

若在研究中使用了 SETR,请引用原论文:

@article{zheng2020rethinking, title={Rethinking Semantic Segmentation from a Sequence-to-Sequence Perspective with Transformers}, author={Zheng, Sixiao and Lu, Jiachen and Zhao, Hengshuang and Zhu, Xiatian and Luo, Zekun and Wang, Yabiao and Fu, Yanwei and Feng, Jianfeng and Xiang, Tao and Torr, Philip HS and others}, journal={arXiv preprint arXiv:2012.15840}, year={2020} }

【免费下载链接】mmsegmentationOpenMMLab Semantic Segmentation Toolbox and Benchmark.项目地址: https://gitcode.com/GitHub_Trending/mm/mmsegmentation

创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询