原文链接:https://github.com/Autumnair007/DeepLearning-notes
参考资料:第一章 SegFormer(语义分割篇)——SegFormer: 简单高效的基于Transformer的语义分割设计-CSDN博客

SegFormer: 轻量级语义分割Transformer学习与实现 - Brisling - 博客园


SegFormer的全称是 “SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers”,其核心目标是构建一个既简单高效又性能强大的语义分割框架。它通过摒弃传统视觉Transformer(ViT)中一些复杂且计算昂贵的设计,例如位置编码(Positional Encoding)和复杂的解码器,成功地在保持高精度的同时,大幅提升了计算效率和模型对不同输入分辨率的鲁棒性。
该模型主要由两大创新部分组成:一个分层的Transformer编码器 (Hierarchical Transformer Encoder) 和一个轻量级的全MLP解码器 (Lightweight All-MLP Decoder)

模型架构图与数据流

在这里插入图片描述

这张图清晰地展示了 SegFormer 的两大核心组件:一个分层的 Transformer 编码器 (Encoder) 和一个轻量级的全 MLP 解码器 (Decoder)

整体数据流概览

假设我们输入一张分辨率为 512 x 512 的图像,类别数 N_cls 为 19(例如 Cityscapes 数据集)。

第一部分:MiT-B2 编码器 (Encoder)

编码器的任务是分层提取特征,生成一个包含四个不同尺度特征图的特征金字塔。

Stage 1: 提取浅层细节特征
  1. 输入: 原始图像,尺寸 512 x 512 x 3

  2. 操作: Overlapping Patch Embedding (重叠块嵌入)

    • 目的: 将 2D 图像转换为 Transformer 能处理的块嵌入序列,并完成第一次下采样。
    • 实现: 这是一个带步长的 2D 卷积层。
    • MiT-B2 参数:
      • 卷积核大小 (Kernel Size, K 1 K_1 K1): 7x7
      • 步长 (Stride, S 1 S_1 S1): 4
      • 填充 (Padding, P 1 P_1 P1): 3
      • 输出通道数 (Channels, C 1 C_1 C1): 64
    • 数据流: 输入的 512x512x3 图像经过这个卷积层,输出尺寸变为 (512-7+2*3)/4 + 1 = 128。在深度学习框架(如 PyTorch, TensorFlow)中,计算卷积层输出尺寸的完整公式通常包含一个 向下取整(floor) 的操作,所以输出尺寸的计算没有问题。
    • 输出: 得到第一个特征图 F 1 ′ F_1' F1, 尺寸为 128 x 128 x 64
  3. 操作: Transformer Encoder Block 1

    • 目的: 对 F 1 ′ F_1' F1 进行特征学习,捕捉局部和全局的依赖关系。
    • MiT-B2 参数:
      • 堆叠层数 ( L 1 L_1 L1): 3 层。这意味着内部的 (高效自注意力 -> Mix-FFN) 结构会重复 3 次。
      • 注意力头数 ( N 1 N_1 N1): 1 个头。
      • FFN 扩展比 ( E 1 E_1 E1): 8。在 Mix-FFN 中,通道数会从 64 扩展到 64 * 8 = 512,然后再缩减回 64。
      • 序列缩减比 ( R 1 R_1 R1): 8。在高效自注意力模块中,Key 和 Value 序列的长度会从 128*128=16384 缩减为 16384 / 8 = 2048,极大地降低了计算复杂度。
    • 数据流: 128x128x64 的特征图经过 3 层 Transformer 模块的变换,尺寸保持不变。
    • 输出: 得到 Stage 1 的最终特征图 F 1 F_1 F1,尺寸仍为 128 x 128 x 64。这是送往解码器的第一个特征。
Stage 2: 进一步下采样和特征提取
  1. 输入: 来自 Stage 1 的特征图 F 1 F_1 F1,尺寸 128 x 128 x 64

  2. 操作: Overlapping Patch Merging (下采样)

    • 实现: 同样是一个带步长的 2D 卷积。
    • MiT-B2 参数: K 2 = 3 , S 2 = 2 , P 2 = 1 K_2=3, S_2=2, P_2=1 K2=3,S2=2,P2=1, 输出通道数 C 2 = 128 C_2=128 C2=128
    • 数据流: (128-3+2*1)/2 + 1 = 64
    • 输出: 得到下采样后的特征图,尺寸为 64 x 64 x 128
  3. 操作: Transformer Encoder Block 2

    • MiT-B2 参数: L 2 = 3 , N 2 = 2 , E 2 = 8 , R 2 = 4 L_2=3, N_2=2, E_2=8, R_2=4 L2=3,N2=2,E2=8,R2=4
    • 数据流: 64x64x128 的特征图经过 3 层 Transformer 模块(这次有 2 个注意力头,序列缩减比为 4),尺寸保持不变。
    • 输出: 得到 Stage 2 的最终特征图 F 2 F_2 F2,尺寸为 64 x 64 x 128。这是送往解码器的第二个特征。
Stage 3: 提取更深层的语义特征
  1. 输入: 来自 Stage 2 的特征图 F 2 F_2 F2,尺寸 64 x 64 x 128

  2. 操作: Overlapping Patch Merging (下采样)

    • MiT-B2 参数: K 3 = 3 , S 3 = 2 , P 3 = 1 K_3=3, S_3=2, P_3=1 K3=3,S3=2,P3=1, 输出通道数 C 3 = 320 C_3=320 C3=320
    • 数据流: (64-3+2*1)/2 + 1 = 32
    • 输出: 得到下采样后的特征图,尺寸为 32 x 32 x 320
  3. 操作: Transformer Encoder Block 3

    • MiT-B2 参数: L 3 = 6 , N 3 = 5 , E 3 = 4 , R 3 = 2 L_3=6, N_3=5, E_3=4, R_3=2 L3=6,N3=5,E3=4,R3=2
    • 数据流: 32x32x320 的特征图经过 6 层 Transformer 模块(5 个注意力头,FFN扩展比为4,序列缩减比为2),尺寸保持不变。
    • 输出: 得到 Stage 3 的最终特征图 F 3 F_3 F3,尺寸为 32 x 32 x 320。这是送往解码器的第三个特征。
Stage 4: 提取最高层的全局上下文
  1. 输入: 来自 Stage 3 的特征图 F 3 F_3 F3,尺寸 32 x 32 x 320

  2. 操作: Overlapping Patch Merging (下采样)

    • MiT-B2 参数: K 4 = 3 , S 4 = 2 , P 4 = 1 K_4=3, S_4=2, P_4=1 K4=3,S4=2,P4=1, 输出通道数 C 4 = 512 C_4=512 C4=512
    • 数据流: (32-3+2*1)/2 + 1 = 16
    • 输出: 得到下采样后的特征图,尺寸为 16 x 16 x 512
  3. 操作: Transformer Encoder Block 4

    • MiT-B2 参数: L 4 = 3 , N 4 = 8 , E 4 = 4 , R 4 = 1 L_4=3, N_4=8, E_4=4, R_4=1 L4=3,N4=8,E4=4,R4=1
    • 数据流: 16x16x512 的特征图经过 3 层 Transformer 模块(8 个注意力头,FFN扩展比为4,序列缩减比为1,即不缩减),尺寸保持不变。
    • 输出: 得到 Stage 4 的最终特征图 F 4 F_4 F4,尺寸为 16 x 16 x 512。这是送往解码器的第四个特征。

至此,编码器部分完成,我们获得了四个关键的特征图:

  • F 1 F_1 F1: 128 x 128 x 64
  • F 2 F_2 F2: 64 x 64 x 128
  • F 3 F_3 F3: 32 x 32 x 320
  • F 4 F_4 F4: 16 x 16 x 512

下面是各个类型的SegFormer模型的具体配置表格:

在这里插入图片描述

分层的 Transformer 编码器 (Hierarchical Transformer Encoder) 详解

SegFormer 的编码器,也称为 MiT (Mix-Transformer),是整个框架的核心,其主要作用是特征提取。与传统的视觉 Transformer (ViT) 只能输出单一分辨率的特征图不同,SegFormer 的编码器借鉴了卷积神经网络 (CNN) 中特征金字塔网络 (FPN) 的思想,设计了一个分层(Hierarchical)结构。

这种设计的核心优势在于,它能从原始图像中提取出一个包含多尺度特征的金字塔。这个金字塔既包含了高分辨率的粗粒度特征(保留了丰富的空间细节和边缘信息),也包含了低分辨率的细粒度特征(捕捉了高级的语义上下文),这对于需要同时理解“哪里是物体”和“这是什么物体”的语义分割任务至关重要。

编码器主要由两个关键构建模块组成:重叠块嵌入 (Overlapped Patch Merging)Transformer 模块 (Transformer Block)。下面我们按照数据流动的顺序,深入剖析其工作原理。

第一步:初始块嵌入 (Initial Overlap Patch Embeddings)

这是数据进入 Transformer 架构的第一站,负责将输入的 2D 图像转换为一系列的 1D 序列(或称为“块嵌入”),作为后续 Transformer 模块的输入。

  • 功能:

    1. 分块 (Patching): 将连续的图像数据分割成块。
    2. 嵌入 (Embedding): 将每个块映射到一个高维特征向量。
    3. 下采样 (Downsampling): 完成第一次分辨率降低。
  • 实现细节:
    此过程通过一个带步长的 2D 卷积操作实现。与 ViT 使用不重叠的块不同,SegFormer 巧妙地让卷积的步长 (Stride, S S S) 小于卷积核大小 (Kernel Size, K K K),从而实现了重叠的块划分。
    x o u t = Conv2d ( x i n , kernel_size = K , stride = S , padding = P ) x_{out} = \text{Conv2d}(x_{in}, \text{kernel\_size}=K, \text{stride}=S, \text{padding}=P) xout=Conv2d(xin,kernel_size=K,stride=S,padding=P)
    这种重叠设计的核心优势在于能够保留块边界周围的局部连续性信息,有效避免了 ViT 中因硬性分割而导致的局部上下文信息丢失问题,对像素级别的密集预测任务(如分割)尤其有利。

  • 数据流:

    • 输入:一张 H x W x 3 的图像。
    • 过程:经过一个 K=7, S=4, P=3 的卷积操作。
    • 输出:一个尺寸为 H/4 x W/4 x C1 的特征图。此时,图像的空间尺寸被缩小了 4 倍,而通道维度从 3 扩展到了 C1。这个特征图可以被看作是一系列 (H/4 * W/4) 个特征向量,每个向量维度为 C1
第二步到第五步:四个分层阶段 (Stage 1-4)

在初始嵌入之后,数据会流经四个连续的处理阶段。每个阶段都由一个或多个 Transformer Block 组成,并在阶段之间通过 Overlap Patch Merging 模块进行下采样,从而逐步降低分辨率、提取更深层的语义特征。

Transformer 模块 (Transformer Block) 内部结构

这是编码器的核心构建单元,如图中放大的橙色框所示,每个分层阶段都由 N 个这样的模块堆叠而成。每个模块包含两个核心子组件:

  1. 高效自注意力 (Efficient Self-Attention)

    问题背景

    标准多头自注意力机制的计算复杂度与序列长度(即图像块的数量 N = H × W N=H \times W N=H×W)的平方成正比,即 O ( N 2 ) O(N^2) O(N2)。在处理高分辨率图像时,巨大的序列长度会导致难以承受的计算和内存开销。

    例如,在 Stage 1 中,输入特征图尺寸为 128x128,序列长度 N = 16384 N = 16384 N=16384。此时 N 2 N^2 N2 是一个天文数字,使得标准自注意力变得不切实际。

    SegFormer 解决方案

    为了解决这一瓶颈,SegFormer 提出了一种高效的自注意力机制。其核心思想是在计算注意力矩阵之前,仅对键 (Key, K) 和值 (Value, V) 的序列进行降维,而保持查询 (Query, Q) 的序列长度不变。这样既能保留完整的空间细节信息(由 Q 承载),又能大幅降低计算复杂度。

    详细实现流程与原理

    1. Q, K, V 的来源

    在 MiT 编码器每个 Stage 的 Transformer Block 内部:

    1. 输入: 来自上一层的特征图,尺寸为 H x W x C
    2. 展平: 特征图被展平为一个序列,尺寸为 N x C (其中 N = H * W)。
    3. 生成 Q, K, V: 该序列通过三个独立的线性层(Linear Layers)分别生成 Q, K, V。在进行序列缩减之前,它们的尺寸都是 N x C
    2. 序列缩减 (Sequence Reduction)

    这是高效注意力的关键。我们以 键 (K) 为例,根据论文中的公式进行说明(对 V 的操作完全相同):

    K ^ = Reshape ( N R , C ⋅ R ) ( K ) K ′ = Linear ( C ⋅ R , C ) ( K ^ ) \hat{K}=\text{Reshape}(\frac{N}{R}, C \cdot R)(K) \\ K'=\text{Linear}(C \cdot R, C)(\hat{K}) K^=Reshape(RN,CR)(K)K=Linear(CR,C)(K^)

    这个过程分为两步:

    • 步骤一: Reshape

      • 输入: 原始的键序列 K,尺寸为 N x C
      • 操作: 这里的 Reshape 是一个概念性描述,并非简单的对 N x C 序列进行维度重组。其实际操作必须回到 2D 空间维度 (H x W x C) 进行,以保留像素的空间局部性。其核心是:将原特征图上 R 个空间相邻的像素块,合并成一个新的“超级像素”,并将它们的特征沿通道维度拼接在一起。
        • 实现细节:
          1. 还原: 将输入序列 K (N x C) 还原为其 2D 形状 (H x W x C)。
          2. 块化: 将 H x W 的空间平面划分为 N/R 个不重叠的块,每个块包含 R 个像素。
          3. 拼接: 对每个块,将其内部所有 R 个像素的特征向量(每个维度为 C)沿通道维度进行拼接,形成一个维度为 C*R 的新特征向量。
      • 输出: 经过上述操作,得到张量 hat(K),其尺寸变为 (N/R) x (C*R)
      • 举例 (Stage 1):
        • 输入 K 的 2D 形状为 128 x 128 x 64。缩减比例 R=64
        • 块化: 128x128 的平面被划分为 16384/64 = 256 个块。这通常通过划分 8x8 (sqrt(R) x sqrt(R)) 的像素块来实现。
        • 拼接: 对每个 8x8 的块,将内部 64 个像素的 C=64 维特征向量全部拼接,得到一个新特征向量,其维度为 R * C = 64 * 64 = 4096
        • 结果: 最终得到 hat(K),其尺寸为 256 x 4096
    • 步骤二: Linear

      • 输入: 重塑后的张量 hat(K),尺寸为 (N/R) x (C*R)
      • 操作: 应用一个线性层,将 C*R 的维度映射回 C。这个线性层学习如何将拼接后的长特征向量,压缩回一个具有代表性的、原始通道维度的特征。
      • 输出: 得到最终被缩减的键序列 K',其尺寸为 (N/R) x C
      • 举例 (Stage 1): hat(K) (256 x 4096) 经过线性层后,得到 K',尺寸为 256 x 64

    经过这个过程,KV 的序列长度都从 N 成功缩减到了 N/R


    一个简单的例子:4x4 矩阵
    • 输入: 一个 4x4 的单通道特征图 (C=1)。
    • 尺寸: H=4, W=4, C=1
    • 序列长度: N = H * W = 16
    • 缩减比例: R = 2

    我们的目标是把 K (尺寸 16 x 1) 变成 K' (尺寸 (16/2) x 1 = 8 x 1)。根据论文公式,中间步骤 hat(K) 的尺寸应为 (N/R) x (C*R),即 (16/2) x (1*2) = 8 x 2

    步骤一: Reshape (实际操作的本质)

    这个操作的本质是:将原特征图上 R 个空间相邻的像素块,合并成一个新的“超级像素”,并将它们的特征(通道)拼接在一起。

    1. 初始状态:2D 特征图

    首先,我们要明确矩阵中每个数字的含义。它代表一个像素的特征向量。因为 C=1,所以这个向量只有一个元素。

    • Pixel(0,0) 的特征向量是 [0]
    • Pixel(0,1) 的特征向量是 [1]
    K_2D = [[ [0], [1], [2], [3] ],
            [ [4], [5], [6], [7] ],
            [ [8], [9],[10],[11] ],
            [[12],[13],[14],[15] ]]
    

    (为了清晰,用 [v] 表示每个像素的特征向量)

    2. 进行“块化” (Patching/Unfolding)

    我们将特征图在空间上划分为不重叠的 1x2 的块。我们只分析 块1

    • 块1 覆盖了 Pixel(0,0)Pixel(0,1) 这两个空间位置。
    • 因此,块1 包含了两个独立的特征向量:[0][1]
    • 此时,您可以将 块1 理解为一个 1x2x1 (H x W x C) 的小张量。
    3. 特征拼接 (Concatenation) - 核心操作

    现在,我们要将这个 1x2 空间块内的信息,压缩成一个 1x1 的“超级像素”,并通过拼接来保留所有信息。

    • 操作: 沿着通道维度 (Channel Axis),将块内所有像素的特征向量拼接起来。
    • 对象: Pixel(0,0) 的特征 [0]Pixel(0,1) 的特征 [1]
    • 过程: Concatenate([0], [1]) 沿通道维度进行。
    • 结果: 产生一个新的、更长的特征向量 [0, 1]。这个新向量的维度是 2 (C*R)。

    这个新的特征向量 [0, 1] 代表了一个新的“超级像素”,它在 1x1 的空间位置上,同时包含了原先 Pixel(0,0)Pixel(0,1) 的全部信息。

    4. 形成 hat(K)

    对所有8个块都执行同样的操作,然后将这些新生成的“超级像素”的特征向量堆叠起来,形成最终的序列 hat(K)

    hat(K) = [ [0, 1],   // 新的 "超级像素" 1 的特征 (C=2)
               [2, 3],   // 新的 "超级像素" 2 的特征 (C=2)
               [4, 5],
               [6, 7],
               [8, 9],
               [10, 11],
               [12, 13],
               [14, 15] ]
    

    这个 hat(K) 的尺寸是 8 x 2,完美符合 (N/R) x (C*R) 的要求。

    步骤二: Linear

    现在 hat(K) (8 x 2) 经过一个线性层 Linear(in_features=C*R, out_features=C),即 Linear(2, 1)。这个线性层会学习如何将这个拼接后的、维度为2的特征向量,重新压缩回一个最有代表性的、维度为1的特征向量。

    K' = Linear(hat(K)),最终 K' 的尺寸会是 8 x 1


    3. 高效注意力计算

    现在,我们使用新的 K'V' 来计算注意力:

    • Q 尺寸: N x C (例如: 16384 x 64)
    • K' 尺寸: (N/R) x C (例如: 256 x 64)
    • V' 尺寸: (N/R) x C (例如: 256 x 64)

    计算过程如下:
    Efficient-Attention ( Q , K ′ , V ′ ) = Softmax ( Q K ′ T d h e a d ) V ′ \text{Efficient-Attention}(Q, K', V') = \text{Softmax}\left(\frac{QK'^T}{\sqrt{d_{head}}}\right)V' Efficient-Attention(Q,K,V)=Softmax(dhead QKT)V

    1. 计算注意力分数: Q 乘以 K' 的转置。

      • Q @ K'^T 的维度: (N x C) @ (C x N/R) -> N x (N/R)
      • 复杂度: 这里的计算复杂度为 O ( N ⋅ ( N / R ) ) = O ( N 2 / R ) O(N \cdot (N/R)) = O(N^2/R) O(N(N/R))=O(N2/R)
    2. 加权求和: 将归一化后的注意力分数乘以 V'

      • Softmax(...) @ V' 的维度: (N x N/R) @ ((N/R) x C) -> N x C

    最终输出的序列尺寸为 N x C,与输入 Q 的尺寸保持一致,确保了信息流的连续性。

    效果与优势

    • 复杂度显著降低: 通过序列缩减,自注意力的计算复杂度从 O ( N 2 ) O(N^2) O(N2) 大幅降低到 O ( N 2 / R ) O(N^2/R) O(N2/R)。在 Stage 1 中,当 R=64 时,计算量减少了 64 倍
    • 分层缩减: SegFormer 在不同阶段使用不同的缩减比例 R,分别为 [64, 16, 4, 1]。在早期层使用大的 R 来处理高分辨率特征图,在深层使用小的 R 来保留更丰富的上下文信息。
    • 实现高分辨率处理: 这种高效设计使得 Transformer 编码器能够首次在语义分割任务中直接处理高分辨率特征图,而不会产生过高的计算成本。
  2. 混合前馈网络 (Mix-Feed-Forward Network, Mix-FFN)
    • 问题背景: 传统 ViT 为了弥补自注意力机制缺乏位置感的缺陷,引入了固定的或可学习的位置编码 (Positional Encoding, PE)。然而,PE 的分辨率是固定的。当输入图像分辨率变化时,这些位置编码需要进行插值,这通常会导致模型精度下降,降低了模型的泛化能力和鲁棒性。

    • SegFormer 解决方案: SegFormer 巧妙地抛弃了所有显式的位置编码。它认为 PE 对于语义分割并非必需,并提出了一种更灵活、更有效的替代方案。其核心思想是在标准的前馈网络 (FFN) 中,通过引入一个 3x3 的深度可分离卷积来隐式地注入位置信息

    • 实现与原理:
      x o u t = MLP ( GELU ( Conv 3 × 3 ( MLP ( x i n ) ) ) ) + x i n x_{out} = \text{MLP}(\text{GELU}(\text{Conv}_{3\times3}(\text{MLP}(x_{in})))) + x_{in} xout=MLP(GELU(Conv3×3(MLP(xin))))+xin

      • 核心机制: 卷积操作本身具有强大的局部感知能力,能捕捉像素间的相对位置关系。但 SegFormer 更进一步,利用了卷积操作中的 零填充 (Zero-Padding) 来“泄露”绝对位置信息。当 3x3 的卷积核在特征图的边缘或角落移动时,其感受野会包含一部分零填充区域。这种独特的填充模式为模型提供了一个强大的信号,使其能够感知到自己处于特征图的绝对位置(例如“在左上角”、“在底部边缘”),从而实现了数据驱动的、无需显式编码的位置感知。

      • 计算流程与维度变化:

        1. 输入: x_in 是来自自注意力模块的输出,其尺寸为序列形态的 N x C (例如 16384 x 64)。
        2. 升维: 通过第一个 MLP (通常是一个线性层),对特征进行升维。尺寸变为 N x C_expand
        3. Reshape (关键步骤): 为了让卷积能够作用,必须将序列还原为 2D 空间形态。尺寸从 N x C_expand 变回 H x W x C_expand (例如 128 x 128 x C_expand)。
        4. 卷积注入位置信息: 在 H x W x C_expand 的特征图上应用一个 3x3 的深度可分离卷积 (Conv_3x3)。这一步在不改变特征图尺寸的同时,通过局部邻域计算和零填充效应,将位置信息融入特征中。选用深度可分离卷积是为了大幅减少参数量,提升模型效率。
        5. 激活与降维: 结果通过 GELU 激活函数,然后送入第二个 MLP (线性层)将维度从 C_expand 降回到 C。此时尺寸仍为 H x W x C
        6. Reshape & 残差连接: 将 H x W x C 的特征图重新展平为 N x C,与原始输入 x_in 进行残差连接。

      这个设计简洁而有效,成功地将对分辨率变化鲁棒的位置信息无缝地融合到了 Transformer 的 FFN 模块中。

  3. 深度可分离卷积讲解

    深度可分离卷积就是两步:先分通道做空间卷积,再用1x1卷积组合通道

    我们通常说的“一张特征图”,其实是一个具有 高度 x 宽度 x 通道数 的三维 “特征体 (Feature Volume)”。一个标准的 3x3 卷积核,其真实尺寸是 3 x 3 x 输入通道数,它也是一个小的三维“卷积体”。卷积操作就是这个“卷积体”在输入的“特征体”上滑动,每次覆盖输入体的所有通道,进行加权求和,最终产生输出“特征体”中的一个像素点。

    理解了这一点,深度可分离卷积就清晰了:

    1. 深度卷积 (Depthwise Convolution): 这一步只负责空间。它把标准卷积“大一统”的操作给拆开了。它不再使用一个3x3x输入通道数的“胖”卷积核,而是为输入的每一个通道都分配一个独立的、3x3x1的“薄”卷积核。你可以想象成,有C个输入通道,我们就用C个“平面”卷积核,分别在各自的通道平面上独立地滑动计算。这一步只在各自通道内部寻找空间特征(比如边缘),完全不关心通道之间的关系。所以,输入一个H x W x C的特征体,输出的还是一个H x W x C的特征体,只是每个通道平面都被“润色”过了。

    2. 逐点卷积 (Pointwise Convolution): 这一步只负责通道。它使用1x1的卷积核,而这个核的真实尺寸是1x1x输入通道数——这正是您领悟到的那个“三维卷积体”最纯粹的样子!它在空间上只看1x1的区域(一个像素),但它会同时观察这个像素位置上所有通道的值。它的工作就是对这些来自不同通道的值进行加权组合,将C_in个通道的信息“混合”成C_out个新的通道。

    所以,SegFormer 用它,就是先用“深度卷积”在每个特征通道上独立地寻找局部位置信息(3x3的范围),然后再用“逐点卷积”高效地混合这些信息并调整通道维度,整个过程计算量远小于一个标准3x3卷积,但同样达到了融合局部信息的效果。

跨阶段的数据流
  • Stage 1: 输入 H/4 x W/4 x C1 的特征图,经过 L1 个 Transformer 模块处理。处理后的特征图再经过一个 Overlap Patch Merging 模块(同样是带步长的卷积),进行下采样。输出一个尺寸为 H/8 x W/8 x C2 的特征图。

  • Stage 2: 输入 H/8 x W/8 x C2 的特征图,经过 L2 个 Transformer 模块处理,再通过 Overlap Patch Merging 下采样。输出一个尺寸为 H/16 x W/16 x C3 的特征图。

  • Stage 3: 输入 H/16 x W/16 x C3 的特征图,经过 L3 个 Transformer 模块处理,再通过 Overlap Patch Merging 下采样。输出一个尺寸为 H/32 x W/32 x C4 的特征图。

  • Stage 4: 输入 H/32 x W/32 x C4 的特征图,经过 L4 个 Transformer 模块处理。注意:在最后一个阶段,为了保留最深层的语义信息,不再进行下采样,因此没有 Overlap Patch Merging 模块。输出的特征图尺寸仍然是 H/32 x W/32 x C4

最终,编码器完成了它的使命,输出了四个不同尺度的特征图,共同构成了送往解码器的特征金字塔:

  • F 1 F_1 F1: H/4 x W/4 x C1 (高分辨率,包含丰富的细节和边缘信息)
  • F 2 F_2 F2: H/8 x W/8 x C2
  • F 3 F_3 F3: H/16 x W/16 x C3
  • F 4 F_4 F4: H/32 x W/32 x C4 (低分辨率,包含丰富的全局语义和上下文信息)

轻量级的全 MLP 解码器 (Lightweight All-MLP Decoder) 详解

SegFormer 的解码器是其 “简单高效” 设计哲学的集中体现。它的核心任务是接收来自编码器的多尺度特征金字塔,高效地将这些特征融合,并最终生成像素级的分割预测。

核心设计理念:为何选择极简解码器?

许多经典的分割模型(如 FCN, U-Net, DeepLabv3+)都设计了结构相对复杂的解码器,包含多个卷积层、上采样层甚至注意力模块。SegFormer 则反其道而行之,采用了极其简洁的全 MLP (Multi-Layer Perceptron) 设计。

这一决策背后的核心洞见在于:SegFormer 强大的分层 Transformer 编码器已经承担了绝大部分的特征提取和上下文聚合工作。由于其独特的结构,编码器具有非常大的有效感受野 (Effective Receptive Field, ERF),这意味着即使是模型浅层(高分辨率)的特征图也已经隐式地包含了丰富的全局上下文信息。因此,一个复杂的、为扩大感受野而设计的、计算昂贵的解码器变得不再必要。

基于此,SegFormer 的解码器不包含任何复杂的标准卷积或注意力模块,完全由线性层(在实现上通常是 1x1 卷积)构成,从而使得解码器部分的参数量极少,计算开销几乎可以忽略不计。

解码器数据流:分步详解

下面我们以一个具体的例子(例如 MiT-B2 模型,输入图像 512x512,解码器嵌入维度 C_embed=256,类别数 N_cls=19)来详细拆解解码器的四个主要步骤。

第一步:统一通道维度 (Unify Channel Dimension)
  • 目的: 将来自编码器的四个不同尺度、不同通道数 ( F 1 , F 2 , F 3 , F 4 F_1, F_2, F_3, F_4 F1,F2,F3,F4) 的特征图,映射到一个统一的、固定的通道维度 C_embed,为后续的融合做准备。

  • 实现: 对每个特征图使用一个独立的 1x1 卷积(即 MLP 层)。

  • 公式:
    F i ^ = Linear i ( F i ) ∀ i ∈ { 1 , 2 , 3 , 4 } \hat{F_i} = \text{Linear}_i(F_i) \quad \forall i \in \{1, 2, 3, 4\} Fi^=Lineari(Fi)i{1,2,3,4}

  • 数据流示例:

    • 输入:
      • F 1 F_1 F1: 128x128x64
      • F 2 F_2 F2: 64x64x128
      • F 3 F_3 F3: 32x32x320
      • F 4 F_4 F4: 16x16x512
    • 过程:
      • F 1 F_1 F1 -> 1x1 Conv -> F 1 ^ \hat{F_1} F1^ (128x128x256)
      • F 2 F_2 F2 -> 1x1 Conv -> F 2 ^ \hat{F_2} F2^ (64x64x256)
      • F 3 F_3 F3 -> 1x1 Conv -> F 3 ^ \hat{F_3} F3^ (32x32x256)
      • F 4 F_4 F4 -> 1x1 Conv -> F 4 ^ \hat{F_4} F4^ (16x16x256)
    • 输出: 四个通道数均为 C_embed=256 的特征图 F i ^ \hat{F_i} Fi^
第二步:上采样至统一尺寸 (Upsample to Uniform Size)
  • 目的: 将所有经过通道统一的特征图,在空间分辨率上对齐到最大的特征图尺寸(即 H/4 x W/4)。
  • 实现: 使用计算量极小且高效的双线性插值 (Bilinear Interpolation)
  • 公式:
    F i ^ ′ = Upsample × 2 i − 1 ( F i ^ ) \hat{F_i}' = \text{Upsample}_{\times 2^{i-1}}(\hat{F_i}) Fi^=Upsample×2i1(Fi^)
  • 数据流示例:
    • 输入: 四个通道数为 256 的特征图 F i ^ \hat{F_i} Fi^
    • 过程:
      • F 1 ^ \hat{F_1} F1^ (128x128x256) -> 保持不变
      • F 2 ^ \hat{F_2} F2^ (64x64x256) -> 上采样 2 倍 -> F 2 ^ ′ \hat{F_2}' F2^ (128x128x256)
      • F 3 ^ \hat{F_3} F3^ (32x32x256) -> 上采样 4 倍 -> F 3 ^ ′ \hat{F_3}' F3^ (128x128x256)
      • F 4 ^ \hat{F_4} F4^ (16x16x256) -> 上采样 8 倍 -> F 4 ^ ′ \hat{F_4}' F4^ (128x128x256)
    • 输出: 四个尺寸均为 128x128x256 的特征图。
第三步:特征拼接 (Concatenate)
  • 目的: 这是特征融合的关键步骤。将所有处理后的、尺寸和通道都对齐的特征图聚合在一起,形成一个包含了所有尺度信息的强大特征表示。
  • 实现: 在通道维度上进行拼接。
  • 公式:
    F f u s e = Concat ( [ F 1 ^ , F 2 ^ ′ , F 3 ^ ′ , F 4 ^ ′ ] ) F_{fuse} = \text{Concat}([\hat{F_1}, \hat{F_2}', \hat{F_3}', \hat{F_4}']) Ffuse=Concat([F1^,F2^,F3^,F4^])
  • 数据流示例:
    • 输入: 四个 128x128x256 的特征图。
    • 过程: 将它们在通道维度上拼接。
    • 输出: 一个融合后的特征图 F f u s e F_{fuse} Ffuse,尺寸为 128 x 128 x (4 * 256) = 128 x 128 x 1024
第四步:最终融合与分类预测 (Final Fusion & Prediction)
  • 目的: 对融合后的高维特征进行最后一次信息提炼,并输出最终的像素级别分类结果。
  • 实现: 包含两个串联的 1x1 卷积层。
  • 公式:
    F o u t = Linear f u s e ( F f u s e ) M 1 / 4 = Linear p r e d ( F o u t ) F_{out} = \text{Linear}_{fuse}(F_{fuse}) \\ M_{1/4} = \text{Linear}_{pred}(F_{out}) Fout=Linearfuse(Ffuse)M1/4=Linearpred(Fout)
  • 数据流示例:
    1. 融合: 输入的 F f u s e F_{fuse} Ffuse (128x128x1024) 经过一个 1x1 卷积,将其通道数从 4 * C_embed 降维回 C_embed。输出 F o u t F_{out} Fout,尺寸为 128x128x256
    2. 预测: F o u t F_{out} Fout (128x128x256) 再经过一个 1x1 卷积(作为分类头),将其通道数从 C_embed 变为 N_cls。输出 1/4 尺寸的分割掩码 M 1 / 4 M_{1/4} M1/4,其尺寸为 128x128x19。这个张量的每个空间位置 (h, w) 上的 19 维向量代表了该像素属于各个类别的分数(logits)。
最后一步:恢复原始分辨率
  • 将解码器输出的 1/4 尺寸分割掩码 M 1 / 4 M_{1/4} M1/4 (128x128x19),通过一次简单的双线性插值上采样 4 倍,得到最终与原图尺寸一致的分割结果,尺寸为 512 x 512 x 19

总结

SegFormer的设计哲学是化繁为简,追求极致的效率与性能平衡

  • 在编码器端,它通过分层架构重叠块嵌入生成了强大的多尺度特征;通过高效自注意力机制大幅降低了计算复杂度;通过混合前馈网络 (Mix-FFN) 巧妙地替代了复杂的位置编码,增强了模型的鲁棒性。
  • 在解码器端,它充分利用了编码器强大的特征表达能力,设计了一个极度轻量级的全MLP解码器,避免了复杂的模块,在显著减少参数量和计算量的同时,保证了出色的分割性能。
    正是这种简洁而深刻的设计,使得SegFormer在多个主流语义分割基准测试(如ADE20K, Cityscapes)中,实现了在同等甚至更少计算量下的SOTA(State-of-the-Art)性能,成为该领域一个里程碑式的工作。
Logo

欢迎加入我们的广州开发者社区,与优秀的开发者共同成长!

更多推荐