Contents — find the section you need

在 SegFormer 出现之前,已有尝试将 Transformer 引入语义分割领域的研究,但许多方法依赖于位置编码和复杂的解码器。因此,它们面临两个问题:当测试分辨率与训练分辨率不同时,准确率可能会下降;计算成本也可能很高。SegFormer(Xie、Wang、Yu、Anandkumar、Alvarez 和 Luo,NeurIPS 2021)通过一个极其简洁的设计解决了这两个问题:一个不依赖位置编码的分层编码器和一个完全不使用卷积的全 MLP 解码器。本文基于原始论文(arXiv:2105.15203),逐步分析了该设计奏效的原因。

本文基于原始论文“SegFormer:基于Transformer的语义分割的简单高效设计”(Xie等人,NeurIPS 2021,arXiv:2105.15203)。

0. 你将学到什么

  • SegFormer旨在解决早期基于Transformer的分割模型中存在的问题

  • 分层MiT(混合Transformer)编码器如何在没有位置编码的情况下工作

  • 为什么仅由MLP构成、完全没有卷积的解码器也能达到高精度

  • 在Cityscapes和ADE20K数据集上测量B0-B5的mIoU、参数数量和计算成本

  • 如何在实践中选择模型,包括对图像损坏的鲁棒性

1. 什么是SegFormer?

SegFormer 是一种语义分割模型,它将不进行位置编码的分层 Transformer 编码器 (MiT) 和一个完全不使用卷积的轻量级全 MLP 解码器相结合。其核心思想是重新设计编码器和解码器之间的分工:编码器生成多分辨率特征后,解码器可以大大简化。

2. 为什么需要这种设计?

直接使用 Vision Transformer (ViT) 作为语义分割的编码器会带来两个问题。首先,ViT 生成的特征只有单一分辨率。同一幅图像中的物体可能具有多种不同尺寸,而传统的基于 CNN 的分割模型(例如 FCN 和 U-Net)通过组合多分辨率特征来解决这个问题。仅生成单一分辨率特征的 ViT 则失去了这种多尺度表征能力。

第二个问题是对位置编码的依赖性。 ViT 将图像分割成多个图像块并将其输入到 Transformer 中,但 Transformer 本身并不知道每个图像块在图像中的具体位置。因此,必须单独添加固定的或学习到的位置编码。由于这种编码是针对训练时使用的输入分辨率构建的,因此推理时如果图像分辨率不同,则需要对位置编码进行插值,这可能会降低精度。图像分割通常需要比检测更高分辨率的输入,这使得这种分辨率依赖性尤为棘手。

SegFormer 通过使用分层编码器并完全消除位置编码,从根本上解决了这两个问题。

3. 输入是什么?

与其他分割模型一样,输入是 RGB 图像 x \in \mathbb{R}^{H \times W \times 3}。SegFormer 在高分辨率图像(1024×2048)上进行评估,用于 Cityscapes 数据集;在 ADE20K 数据集上使用约 512×512 的图像进行评估。 ViT 使用粗糙的 16×16 图像块进行下采样,而 SegFormer 的编码器则从 4×4 的小图像块开始,在特征提取初期保留更多细节。

4. 预测结果是什么?

输出结果是一个分辨率为输入分辨率 1/4 的分割图,每个像素的类别用 \hat{y} \in \mathbb{R}^{\frac{H}{4} \times \frac{W}{4} \times N_{cls}} 表示(N_{cls} 表示类别数)。然后使用最近邻插值或类似方法将其还原到原始分辨率。其核心思想是,如果编码器生成足够丰富的多尺度特征,解码器只需组合一些简单的线性层即可将这些特征整合并转换为精确的掩码。

5. 基本架构

SegFormer 由两个主要模块组成:一个 MiT(混合变换器)编码器,用于提取四种分辨率的特征;以及一个 全 MLP 解码器,用于组合这些特征并输出分割掩码。

Diagram 1 · Use the button to switch views
SegFormer basic pipeline A pipeline in which an input image passes through MiT encoder Stages 1–4, reducing its resolution from 1/4 to 1/32 while extracting features, and then applies an All-MLP decoder. Input image MiT encoder (no positional encoding) Stage 1 (1/4) Stage 2 (1/8) Stage 3 (1/16) Stage 4 (1/32) Each stage: Overlap Patch Merging Efficient Self-Attention Mix-FFN (3×3 Conv injects position implicitly) All-MLP decoder 1. Unify channels with MLPs 2. Upsample to 1/4 resolution 3. Concatenate 4. Fuse with MLP (4C→C) 5. Predict classes with MLP Mask

图 1 — MiT 编码器生成四种分辨率(1/4、1/8、1/16 和 1/32)的特征,All-MLP 解码器仅使用线性层将它们组合起来以创建分割掩码。唯一的卷积位于编码器的 Mix-FFN 中;解码器实际上完全由线性层构成。

由于编码器预先构建了四种分辨率的特征,解码器只需极少的处理即可获得精确的掩码:对齐特征并将它们合并。这种分工是 SegFormer 参数总数保持较低的主要原因。

6. 组件的技术细节

重叠块合并 — 四阶段层次结构

与 ViT 仅将图像分割一次不同,MiT 编码器分四个阶段逐步降低分辨率。输出分辨率分别为输入分辨率的 1/4、1/8、1/16 和 1/32,重现了 ResNet 等 CNN 骨干网络所使用的层次结构。

块分割方法也不同于 ViT 的非重叠块:相邻块在合并时会重叠。第一阶段使用内核大小为 K=7、步长为 S=4 和填充大小为 P=3 的参数;后续阶段使用 K=3, S=2, P=1 的参数。重叠部分保留了相邻图像块边界处的局部连续性(即相邻图像块共享的信息),同时特征图也得到了缩减。

高效的自注意力机制

普通的多头自注意力机制对于长度为 N 的序列需要进行 O(N^2) 次计算。由于高分辨率分割时 N 会变得非常大,这是一个严重的瓶颈。SegFormer 引入了序列缩减机制,在计算注意力之前,将键值序列压缩 R 倍。

\text{Attention}(Q, K, V) = \text{Softmax}\!\left(\frac{QK^{\top}}{\sqrt{d_{head}}}\right)V

K, V 序列在使用前会从长度 N 缩减到 N/R。第 1 至 4 阶段的缩减比例 R 设置为 [64, 16, 4, 1]:浅层阶段具有高分辨率和长序列,因此压缩率很高;深层阶段分辨率较低、序列较短,因此几乎不压缩。这使得计算成本从 O(N^2) 降低到 O(N^2/R),即使对于高分辨率输入,自注意力机制也切实可行。

Mix-FFN — 如何消除位置编码

SegFormer 可以完全避免位置编码,因为它将一个 3×3 卷积混合到前馈网络 (FFN) 中,从而创建了 Mix-FFN。

x_{out} = \text{MLP}\big(\text{GELU}(\text{Conv}_{3\times3}(\text{MLP}(x_{in})))\big) + x_{in}

3×3 卷积在图像边界处使用零填充。这种填充边界间接地将“我在图像中的位置”这一信息传递给了卷积层,因此位置信息只需通过 Mix-FFN 即可融入特征中,而无需显式的位置编码。因此,在与训练分辨率不同的分辨率下进行测试时,无需插值位置编码,否则可能会降低精度。

全 MLP 解码器 — 为什么不需要卷积

解码器由以下四个步骤组成,实际上只使用了线性层(MLP)。

  1. 统一通道:使用独立的线性层将每个阶段的特征 F_i 映射到相同通道数的 C。

  2. 上采样并连接:将每个特征上采样到 1/4 分辨率,并沿通道维度连接。

  3. 融合:使用线性层将拼接后的 4C 维特征映射到 C 维。

  4. 预测:使用最后一个线性层输出一个分割掩码,每个类别对应一个通道。

\hat{M} = \text{Linear}\Big(\text{Concat}\big(\text{Upsample}(\text{Linear}(F_1)), \dots, \text{Upsample}(\text{Linear}(F_4))\big)\Big)

原始论文通过对有效感受野 (ERF)的分析来支持此设计。基于 CNN 的解码器(例如 DeepLabv3+)需要复杂的机制——空洞卷积和多尺度池化——才能获得宽广的感受野并实现高精度。相比之下,MiT 编码器通过自注意力机制,在浅层阶段自然地获取局部 ERF,在深层阶段获取覆盖整幅图像的非局部 ERF,而无需额外的机制。由于编码器本身已经构建了包含从局部到全局尺度信息的特征,因此解码器不需要复杂的感受野扩展机制。这就是简单的线性层组合就足以实现的理论原因。

训练方案

原始论文采用了一种简单的训练设置,没有使用任何特殊技巧。MiT 编码器在 ImageNet-1K 数据集上进行预训练,解码器则从随机初始化开始训练。优化器采用 AdamW,学习率从初始值 0.00006 开始,并按照幂律衰减(系数为 1.0)进行衰减。在 ADE20K 和 Cityscapes 数据集上迭代次数为 16 万次,在 COCO-Stuff 数据集上迭代次数为 8 万次。训练数据增强仅限于随机缩放(缩放比例为 0.5–2.0)、随机水平翻转和随机裁剪(ADE20K 数据集为 512×512,Cityscapes 数据集为 1024×1024,ADE20K 数据集上的 B5 部分为 640×640)。论文明确指出,它没有使用诸如 OHEM(在线难例挖掘)、辅助损失或类别平衡损失等广泛应用的技巧。目标是证明 MiT 编码器和 All-MLP 解码器无需添加这些组件即可达到较高的 mIoU。

7. 模型变体比较

SegFormer 提供六种尺寸,从 MiT-B0 到 MiT-B5。以下是原始论文表 2 中报告的 Cityscapes 和 ADE20K 测试结果。

模型 参数 Cityscapes 浮点运算次数 Cityscapes mIoU (MS) ADE20K 浮点运算次数 ADE20K mIoU
SegFormer-B0 3.8M 125.5G 76.2 8.4G 37.4
SegFormer-B1 13.1M 243.7G 78.5 15.9G 42.2
SegFormer-B2 24.2M 717.1G 81.0 62.4G 46.5
SegFormer-B3 44.0M 962.9G 81.7 79.0G 49.4
SegFormer-B4 64.1M 1240.6G 82.3 95.7G 50.3
SegFormer-B5 84.7M 1460.4G 84.0 183.3G 51.0

来源:原始论文(Xie et al., NeurIPS 2021)表2。Cityscapes mIoU采用多尺度和左右翻转(MS)测试进行测量。参数数量包含完整的编码器和解码器;即使是 B5,解码器也仅占总参数的约 4%。作为参考,B5 在单尺度 (SS) 测试中达到 82.4 mIoU,在多尺度 (MS) 测试中达到 84.0 mIoU。

该论文将这些结果与当时最先进的方法进行了比较,并指出 SegFormer-B4 在 ADE20K 数据集上以 6410 万个参数达到了 50.3% 的 mIoU,比当时最佳方法高出 2.2 个百分点,而模型规模却只有其五分之一左右。B0 仅使用 380 万个参数就保持了 76.2% 的实用 Cityscapes mIoU,这些数据支持其面向边缘的设计。

8. 它的不足之处

SegFormer 的大多数弱点也是基于 Transformer 的模型普遍存在的。首先,较大的模型(B3 及以上)具有较高的训练和推理成本。 Cityscapes 数据集的浮点运算量 (FLOPs) 从 B0 的 125.5G 增长到 B5 的 1460.4G,增长超过十倍。对于高分辨率分割而言,这种增长会使实时操作变得困难。

对大规模预训练数据的依赖也至关重要。MiT 编码器在 ImageNet-1K 数据集上进行预训练,因此微调数据的质量和数量会极大地影响其迁移到与预训练数据差异显著的领域(例如医学或卫星图像)的效果。

**对于 SegFormer 而言,处理非常小的物体和细小的结构(例如电线和路标)也并非易事。重叠块合并即使在第一阶段也会进行 4×4 的下采样,因此只有几个像素宽的结构在进入更深的阶段之前就可能丢失信息。

与此同时,原始论文中报告的 Cityscapes-C 数据集上的鲁棒性评估结果是一个明显的优势。对于添加了高斯噪声的受损图像,SegFormer-B5 相对于基于 Xception-71 骨干网络的 DeepLabv3+,最大相对 mIoU 提升了 588%;对于类似雪景的受损图像,最大相对提升为 295%。这些测量结果表明,与普通的基于 CNN 的模型相比,SegFormer-B5 对图像受损具有更高的容忍度。这种鲁棒性被认为是由于自注意力机制不易受到局部噪声的影响。

9. 如何在实践中选择模型

对于边缘设备和实时处理,SegFormer-B0 或 B1 是切实可行的选择。B0 在精度和紧凑性之间取得了平衡,拥有 380 万个参数,在 Cityscapes 数据集上实现了 76.2% 的 mIoU,使其能够部署在嵌入式设备和无人机上。

对于自动驾驶中的可行驶区域识别等需要高分辨率图像和稳定精度的应用,B2 或 B3 通常是一个合理的折衷方案。 B4 和 B5 的精度最高,但它们在 Cityscapes 数据集上的浮点运算量超过 1200G,因此在车载 GPU 上进行实时推理通常需要进行优化,例如量化或 TensorRT 转换。

对于医疗或卫星图像的离线精确分析,精度优先于实时性能。B4 或 B5 较为合适,并可根据需要使用特定领域的数据进行微调。

对于经常受到恶劣天气或光照变化影响的户外环境,例如农业机器人和户外监控摄像头,在 Cityscapes-C 数据集上测得的鲁棒性是一项实际优势。与精度相近的基于 CNN 的模型相比,SegFormer 在图像损坏的情况下精度损失预计更小。

要从基础开始学习目标检测和语义分割,请参阅“目标检测和语义分割简介”。关于分割领域的最新趋势,请参阅“语义分割技术趋势”。

10. 三行总结

  • SegFormer 结合了分层编码器 (MiT) 和 All-MLP 解码器。MiT 编码器无需位置编码即可生成多尺度特征,而 All-MLP 解码器则无需卷积。

  • Mix-FFN 中的 3×3 卷积隐式地泄露了位置信息,从而无需位置编码;而 Transformer 的宽广有效感受野使得简单的 MLP 解码器成为可能。

  • 该系列模型的参数范围从 B0(380 万参数,在 Cityscapes 数据集上的 IoU 为 76.2 mIoU)到 B5(8470 万参数,IoU 为 84.0 mIoU),测量结果也支持其对图像损坏的鲁棒性。

参考资料

检查你的理解
图像分类和分割的输出有何不同?

分类预测整幅图像的标签,而分割则预测图像的局部属性。

预测每个像素的标签。不应仅以图像级精度来评估小物体和边界。

What to read next

Review the backgroundYOLO11 深度解析——C3k2 和 C2PSA 与 YOLOv8 相比发生了哪些变化?Continue the series机器学习入门:姿态估计Explore another aspect of this field基准测试本地LLM:首次响应、生成率和内存