使用自定义数据集微调 SegFormer 进行语义分割

Hugging Face 详细介绍了一种微调 SegFormer 的工作流程，SegFormer 是一种最先进的语义分割模型，用于处理自定义图像数据集。通过利用预训练权重和 Hugging Face 生态系统，此过程可以创建专门的模型——例如为披萨外卖机器人设计的、用于识别人行道和障碍物的模型。

理解 SegFormer 和语义分割

语义分割是为图像中的每个单独像素进行分类的过程，它提供了比标准图像分类更细粒度的细节。这种能力对于医学成像和自动驾驶等应用至关重要，在这些应用中需要精确的边界检测（例如，识别人行道的准确边缘）。

SegFormer 由 Xie 等人在 2021 年推出，改进了之前的卷积神经网络 (CNN) 和 Vision Transformer (ViT) 方法。其架构由以下部分组成：

分层 Transformer 编码器：与 ViT 不同，SegFormer 不使用位置编码。
一个简单的 MLP 解码器：一个处理编码器输出的多层感知器解码器。

数据准备和增强

有效的语义分割需要具有精确分割图的数据集。虽然存在 ADE20k、CityScapes 或 BDD100K 等通用数据集，但通常需要特定领域的数据以避免分布不匹配。例如，在人行道上运行的机器人需要从人行道视角而非汽车视角捕获的数据。

数据集加载和处理

使用 datasets 库，可以加载自定义数据集（例如 segments/sidewalk-semantic）并将其拆分为训练集和测试集。为了确保模型接收到正确格式的数据，使用了 SegFormerImageProcessor。

即时转换 (On-the-Fly Transforms)

为了优化磁盘空间和训练速度，Hugging Face 建议使用通过 set_transform 进行的转换。这会即时准备数据批次，而不是预先对整个数据集进行预处理。为了增加模型对不同光照条件的适应性，torchvision.transforms.ColorJitter 被集成到训练流水线中，以随机调整亮度、对比度、饱和度和色调。

微调工作流程

模型选择

SegFormer 提供五种模型大小（B0 到 B5）。对于边缘部署——例如在送餐机器人上——由于其占用空间小（约 14MB）且效率高，建议使用 B0 模型。微调过程通常从在 ImageNet-1k 上预训练的模型 (nvidia/mit-b0) 开始。

训练配置

微调通过 Hugging Face Trainer API 进行管理。关键配置参数包括：

超参数：学习率（例如 0.00006）和 epoch 数量（例如 50）。
评估指标：使用平均交并比 (mIoU) 来衡量预测的分割掩码与地面真值 (ground truth) 之间的重叠程度。
Logit 上采样：由于 SegFormer 输出的 logits 在原始图像分辨率的 1/4（高度/4，宽度/4），因此在计算 mIoU 之前，必须使用双线性插值将它们上采样以匹配标签大小。

推理和部署

一旦完成微调，模型及其图像处理器可以推送到 Hugging Face Hub。这可以实现轻松的共享，并可以通过托管的推理 API 创建推理组件 (inference widget) 进行实时测试。

执行推理

要对新图像进行推理，需要执行以下步骤：

预处理：使用 SegformerImageProcessor 处理图像。
前向传播：将处理后的图像通过模型以获得 logits。
重缩放：使用 nn.functional.interpolate 将 logits 上采样到原始图像尺寸。
预测：在类别维度上应用 argmax 操作，以确定最终的像素级类别预测。