Transformers 中的 SegFormer：MiT 编码器与全 MLP 解码器的语义分割实践指南

本文围绕 Transformers 官方模型文档 SegFormer 展开，系统讲解 SegFormer 在库中的架构实现（MiT 层级 Transformer 编码器 + 轻量全 MLP 解码器）、基于 pipeline / AutoModel 的语义分割推理用法，以及 SegformerImageProcessor 中 do_reduce_labels 等关键预处理机制。读完本文，你可以直接复现 ADE20K 预训练模型的推理流程，理解各配置项的取值来源，并结合仓库自带脚本完成语义分割微调。

SegFormer 模型概述

SegFormer（出自论文 SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers）是一个语义分割模型，它将层级式 Transformer 编码器（Mix Transformer，简称 MiT）与一个轻量的"全 MLP"解码器（all-MLP decoder）组合在一起。与传统语义分割方案相比，它有兩個显著的设计取向：

• 不使用显式位置编码：从源码实现看（SegformerMixMLP），Mix-FFN 中夹了一层 3x3 深度卷积（SegformerDepthWiseConv），注释中明确说明该卷积"隐式编码位置信息"，替代了标准 ViT/BeiT MLP 中的显式位置嵌入；
• 解码器极其简单：解码头 SegformerDecodeHead 只由线性投影、双线性上采样与 1x1 卷积融合构成，避免了复杂的上采样结构，使整个模型更简单、高效、可扩展，官方文档也指出该设计在 ADE20K 与 Cityscapes 等基准上达到了当时的先进水平。

所有原始 SegFormer 检查点托管在 Hugging Face 的 NVIDIA 组织页（例如 nvidia/segformer-b0-finetuned-ade-512-512），仓库中的示例代码与文档示例统一使用该检查点。

快速上手：Pipeline 与 AutoModel 两种方式

官方文档给出了两种最典型的推理路径，以下完整保留并补充了输出形状的说明。

方式一：使用 image-segmentation Pipeline

from transformers import pipeline

pipeline = pipeline(task="image-segmentation", model="nvidia/segformer-b0-finetuned-ade-512-512")
pipeline("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg")

Pipeline 会自动完成图像加载、预处理、模型推理与解码后处理，适合快速验证与演示。

方式二：使用 AutoModelForSemanticSegmentation + AutoProcessor

import requests
from PIL import Image

from transformers import AutoModelForSemanticSegmentation, AutoProcessor

url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg"
image = Image.open(requests.get(url, stream=True).raw)

processor = AutoProcessor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
model = AutoModelForSemanticSegmentation.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512", device_map="auto")

inputs = processor(images=image, return_tensors="pt").to(model.device)
outputs = model(**inputs)
logits = outputs.logits  # shape [batch, num_labels, height, width]

几点说明：

• AutoProcessor 会按后端自动解析出对应的图像处理器。从 auto_mappings.py 的映射可以看到，segformer 对应 torchvision 后端的 SegformerImageProcessor 与 pil 后端的 SegformerImageProcessorPil；
• AutoModelForSemanticSegmentation 则解析为 SegformerForSemanticSegmentation（见 modeling_auto.py）；
• 输出的 logits 形状为 [batch, num_labels, height/4, width/4]——解码头在 1/4 分辨率处产生预测，需要还原到原图尺寸时，可用 processor.post_process_semantic_segmentation(outputs, target_sizes=[[h, w]]) 内部的双线性插值完成放大（实现见 image_processing_segformer.py），不指定 target_sizes 时直接对 logits 取 argmax。

实践要点：输入尺寸、裁剪填充与 do_reduce_labels

官方文档 Notes 部分列出的三条实践要点是微调 SegFormer 时最容易踩坑的地方，这里逐条展开并结合源码印证。

1. 支持任意输入尺寸，内部按 patch 尺寸对齐

文档明确指出：SegFormer 可以接受任意输入尺寸，预处理器会把输入填充到可被 config.patch_sizes 整除的尺寸。这与模型实现一致——SegformerOverlapPatchEmbeddings 采用带对称填充（padding=patch_size // 2）的步幅卷积做重叠分块（modeling_segformer.py），编码器内部没有任何固定尺寸断言，各 stage 只依赖当前 feature map 的高宽 height/width 做形状变换，因此推理时对输入尺寸没有硬性约束。

2. 最关键的预处理：随机裁剪/填充到统一尺寸再归一化

文档强调，最重要的预处理步骤是将所有图像随机裁剪并填充到相同尺寸（如 512x512 或 640x640），然后做归一化。这与 SegformerImageProcessor 的默认行为一致：

• 默认 size = {"height": 512, "width": 512}，default_to_square = True，即默认把所有输入裁剪/填充为 512x512 正方形（image_processing_segformer.py）；
• 图像本身走 rescale(1/255) + ImageNet 均值/方差归一化（IMAGENET_DEFAULT_MEAN/STD），插值方式为双线性（BILINEAR）；
• 分割标注图走不同路径：_preprocess_image_like_inputs 中，segmentation_maps 会强制 do_normalize=False、do_rescale=False，并改用 NEAREST_EXACT 最近邻插值（防止标签被插值"抹花"），最终转为 int64 并压缩掉通道维后放入 data["labels"]（image_processing_segformer.py）。

3. do_reduce_labels：处理 ADE20K 这类"0 号背景不算类"的数据集

ADE20K 等数据集在标注中用 0 表示背景，但类别表里并不包含"背景"这一类。为此 SegformerImageProcessor 提供 do_reduce_labels 参数，将所有标签值减 1，并把原来为 0 的背景位置替换为 255——255 正是损失函数里的 ignore_index（对应配置项 semantic_loss_ignore_index=255），从而保证背景像素不参与损失计算。若你的数据集本身就带有背景类，则应保持 do_reduce_labels=False（这也是默认值）。

源码中的处理顺序值得注意（reduce_label）：先把 0 整体替换为 255，再对所有值 -1，最后把减出来的 254 还原回 255——即先占位、再平移、再复位，确保背景最终稳定落在 255。

from transformers import SegformerImageProcessor

processor = SegformerImageProcessor(do_reduce_labels=True)

损失侧与之配套：SegformerForSemanticSegmentation.forward 在多类场景使用 CrossEntropyLoss(ignore_index=config.semantic_loss_ignore_index)，且计算损失前会把 logits 双线性上采样到标注图尺寸（modeling_segformer.py）；单通道（num_labels == 1）时退化为带有效像素掩码的 BCEWithLogitsLoss。

SegformerConfig：编码器结构的完整参数表

SegformerConfig 定义了 MiT 编码器的全部结构超参数。默认值即为论文中 SegFormer-B0 的规格（注意：这是最小的 B0 变体，更大模型如 b3/b5 只是各 stage 的 hidden_sizes 等维度更大）：

参数	默认值	含义
num_channels	3	输入图像通道数
num_encoder_blocks	4	编码器 stage（块）数
depths	(2, 2, 2, 2)	每个 stage 内的 Transformer 层数
hidden_sizes	(32, 64, 160, 256)	每个 stage 的隐藏维度
patch_sizes	(7, 3, 3, 3)	各 stage 入口的重叠分块卷积核大小
strides	(4, 2, 2, 2)	各 stage 入口卷积的步幅（逐级下采样，总计 1/16）
num_attention_heads	(1, 2, 5, 8)	各 stage 的注意力头数
sr_ratios	(8, 4, 2, 1)	各 stage 的序列缩减比率（作用于 K/V）
mlp_ratios	(4, 4, 4, 4)	Mix-FFN 隐层相对输入的比例
hidden_act	"gelu"	Mix-FFN 激活函数
hidden_dropout_prob	0.0	隐藏层 dropout
attention_probs_dropout_prob	0.0	注意力概率 dropout
classifier_dropout_prob	0.1	解码头分类前的 dropout
drop_path_rate	0.1	随机深度（DropPath）最大丢弃率，按层线性递增分配
layer_norm_eps	1e-6	LayerNorm epsilon
decoder_hidden_size	256	解码头统一通道数
semantic_loss_ignore_index	255	语义分割损失中忽略的标签索引
initializer_range	0.02	权重初始化标准差
reshape_last_stage	True	最后一个 stage 输出是否 reshape 回 (B, C, H, W) 张量形式

配置使用示例（来自 configuration_segformer.py 的官方 docstring）：

from transformers import SegformerModel, SegformerConfig

# 初始化一个 nvidia/segformer-b0-finetuned-ade-512-512 风格的配置
configuration = SegformerConfig()

# 基于该配置初始化模型
model = SegformerModel(configuration)

# 访问模型配置
configuration = model.config

从 modeling_segformer.py 可以看到 drop_path_rate 的实际用法：SegformerModel 在构造时为全部 sum(depths) 个块生成线性递减的 DropPath 概率序列，越深的块丢弃概率越高，这是标准的随机深度（stochastic depth）退火策略。

核心架构源码解析

以下按数据流顺序拆解 MiT 编码器与解码头，对应实现均在 modeling_segformer.py。

重叠分块嵌入（Overlapping Patch Embeddings）

SegformerOverlapPatchEmbeddings 用一次带对称填充的 nn.Conv2d（核大小 patch_size、步幅 stride、padding=patch_size//2）完成像素到 token 序列的投影，再接 LayerNorm，返回 (embeddings, height, width)。重叠分块让相邻 token 共享像素信息，是 MiT 区别于普通 ViT 分块的关键（L65-L84）。

序列缩减注意力（Sequence Reduction Attention）

SegformerAttention 借鉴了 PvT 的设计：query 使用完整序列，而 key/value 先经过 SegformerSequenceReduction 用 kernel=stride=sr_ratio 的步幅卷积做空间缩减，把序列从 (B, H*W, C) 压到 (B, (H/sr)*(W/sr), C)，从而把注意力开销从 O(N²) 降下来。首个 stage 的 sr_ratio=8 缩减最激进，最后一层 sr_ratio=1 则退化为标准全序列注意力（L139-L203）。实现中通过 ALL_ATTENTION_FUNCTIONS 分发，SegformerPreTrainedModel 声明支持 SDPA、Flash Attention、Flex Attention 等加速后端（L354-L372），并支持梯度检查点与 torch.compile 全图编译。

Mix-FFN 与 Transformer 块

SegformerMixMLP 的计算路径为 fc1 → DWConv(3x3) → 激活 → fc2，两次 dropout 分别使用 hidden_dropout_prob；3x3 深度卷积承担位置信息注入。SegformerLayer 采用 pre-norm 结构，注意力分支与 FFN 分支各接一次 DropPath 残差（L271-L306）。

编码器主干 SegformerModel

SegformerModel.forward 的逻辑非常简洁：将 pixel_values 依次穿过 num_encoder_blocks 个 SegformerStage，返回 BaseModelOutput(last_hidden_state=...)；前 3 个 stage 输出 reshape 为 (B, C, H, W) 的多尺度特征图，最后一个 stage 的输出形态受 reshape_last_stage 控制（L375-L401）。

全 MLP 解码头 SegformerDecodeHead

解码头的每一步都很轻（L478-L532）：

1. 每个 stage 的特征图经一个 nn.Linear（SegformerMLP）投影到统一的 decoder_hidden_size；
2. 统一通道后双线性上采样到首个 stage 的分辨率（即输入的 1/4 分辨率）；
3. 按 stage 倒序 torch.cat 拼接（浅层特征排在前面），再过 1x1 卷积融合（linear_fuse）+ BatchNorm2d + ReLU + dropout，对应原实现的 ConvModule；
4. 最后用 1x1 卷积 classifier 输出 (B, num_labels, H/4, W/4) 的 logits。

任务头

• SegformerForSemanticSegmentation：主干 + 解码头，训练时把 logits 上采样到标注尺寸后计算交叉熵（或单通道的 BCE），返回 SemanticSegmenterOutput；
• SegformerForImageClassification：复用同一 MiT 主干，将最后 stage 的特征展平后做全局平均池化，再过一个 nn.Linear(hidden_sizes[-1], num_labels)，即用于 ImageNet 类任务的分类头（L410-L462）。

微调与测试资源

• 微调脚本：examples/pytorch/semantic-segmentation/run_semantic_segmentation.py 与 run_semantic_segmentation_no_trainer.py 演示了如何用 Trainer（--model_name_or_path 指向 nvidia/segformer-b0-finetuned-ade-512-512，配合 --dataset_name 自定义数据集）或纯 PyTorch 循环微调 AutoModelForSemanticSegmentation 支持的任何模型（SegFormer、BEiT、DPT 等），详细说明见 examples/pytorch/semantic-segmentation/README.md；其中 Trainer 版本示例会在 sidewalk-semantic 等数据集上输出 segformer_outputs/ 并推送到 Hub，no_trainer 版本可用 accelerate launch 启动；
• 建模测试：tests/models/segformer/test_modeling_segformer.py 中 SegformerModelTester 使用极小配置（depths=[1,1,1,1]、hidden_sizes=[8,8,16,16]、decoder_hidden_size=32、64x64 输入）覆盖 SegformerModel、SegformerForSemanticSegmentation、SegformerForImageClassification 的前向与损失路径，SegformerConfigTester 则断言 hidden_sizes、num_attention_heads、num_encoder_blocks 三个属性存在；
• 图像处理测试：tests/models/segformer/test_image_processing_segformer.py 覆盖 preprocess 与 post_process_semantic_segmentation 的输入输出行为。

小结

SegFormer 在 Transformers 中的实现完整保留了原论文"简单而高效"的设计内核：重叠分块 + 序列缩减注意力构成 MiT 编码器，DWConv-MixFFN 免去显式位置编码，全 MLP 解码头以 1x1 卷积完成多尺度融合，logits 以 1/4 分辨率输出后按需插值还原。配合 SegformerImageProcessor 对 512x512 统一尺寸、最近邻标注插值与 do_reduce_labels 背景掩码的默认处理，推理（pipeline 或 AutoModel 两行代码即可）与微调（仓库自带脚本）都有现成路径。如果你需要自定义数据集上微调，从 examples/pytorch/semantic-segmentation 的脚本与 docs/source/en/model_doc/segformer.md 中的预处理要点入手即可。