使用自定义数据集微调语义分割模型

本指南展示了如何微调 Segformer，一个最先进的语义分割模型。我们的目标是为披萨配送机器人构建一个模型，使其能够识别行驶路径和障碍物 🍕🤖。我们首先将在 Segments.ai 上标注一组人行道图像。然后，我们将使用 🤗 transformers（一个提供易于使用的最先进模型实现的开源库）来微调预训练的 SegFormer 模型。在此过程中，您将学习如何使用 Hugging Face Hub，它是最大的开源模型和数据集目录。

语义分割的任务是对图像中的每个像素进行分类。您可以将其视为一种更精确的图像分类方式。它在医学图像和自动驾驶等领域有广泛的应用。例如，对于我们的披萨配送机器人，重要的是要知道人行道在图像中的确切位置，而不仅仅是是否有个人行道。

由于语义分割是一种分类，因此用于图像分类和语义分割的网络架构非常相似。2014 年，Long 等人一篇开创性的论文使用卷积神经网络进行语义分割。最近，Transformer 已被用于图像分类（例如 ViT），现在它们也被用于语义分割，进一步推动了最先进技术的发展。

SegFormer 是 Xie 等人在 2021 年推出的一种语义分割模型。它具有一个不使用位置编码（与 ViT 不同）的分层 Transformer 编码器和一个简单的多层感知器解码器。SegFormer 在多个常见数据集上实现了最先进的性能。让我们看看我们的披萨配送机器人在人行道图像上的表现如何。

让我们开始安装必要的依赖项。因为我们要将数据集和模型推送到 Hugging Face Hub，所以我们需要安装 Git LFS 并登录到 Hugging Face。

git-lfs 的安装在您的系统上可能有所不同。请注意，Google Colab 预装了 Git LFS。

pip install -q transformers datasets evaluate segments-ai
apt-get install git-lfs
git lfs install
huggingface-cli login

1. 创建/选择数据集

任何机器学习项目的第一步都是收集一个好的数据集。为了训练语义分割模型，我们需要一个带有语义分割标签的数据集。我们可以使用 Hugging Face Hub 中的现有数据集，例如 ADE20k，或者创建我们自己的数据集。

对于我们的披萨配送机器人，我们可以使用现有的自动驾驶数据集，例如 CityScapes 或 BDD100K。然而，这些数据集是由在道路上行驶的汽车捕获的。由于我们的配送机器人将在人行道上行驶，这些数据集中的图像与我们的机器人在现实世界中看到的数据之间将存在不匹配。

我们不希望我们的配送机器人感到困惑，所以我们将使用在人行道上拍摄的图像创建我们自己的语义分割数据集。我们将在接下来的步骤中展示如何标注我们拍摄的图像。如果您只想使用我们已完成的标注数据集，您可以跳过 “创建您自己的数据集” 部分，并从 “使用 Hugging Face Hub 中的数据集” 继续。

创建您自己的数据集

要创建您的语义分割数据集，您需要两样东西：

1. 覆盖模型在现实世界中将遇到的各种情况的图像
2. 分割标签，即每个像素代表一个类别/分类的图像。

我们已经在比利时拍摄了一千张人行道图像。收集和标注这样的数据集可能需要很长时间，因此您可以从小数据集开始，如果模型表现不佳，再进行扩展。

为了获得分割标签，我们需要指示这些图像中所有区域/对象的类别。这可能是一项耗时的工作，但使用正确的工具可以显著加快任务速度。对于标注，我们将使用 Segments.ai，因为它具有用于图像分割的智能标注工具和易于使用的 Python SDK。

在 Segments.ai 上设置标注任务

首先，在 https://segments.ai/join 创建一个账户。接下来，创建一个新的数据集并上传您的图像。您可以通过网页界面或 Python SDK（参见Notebook）进行此操作。

标注图像

现在原始数据已加载，请访问 segments.ai/home 并打开新创建的数据集。点击“开始标注”并创建分割掩码。您可以使用 ML 驱动的超像素和自动分割工具来加快标注速度。

将结果推送到 Hugging Face Hub

标注完成后，创建一个包含标注数据的新数据集版本。您可以在 Segments.ai 的版本选项卡上执行此操作，也可以通过 SDK 以编程方式执行，如笔记本中所示。

请注意，创建版本可能需要几秒钟。您可以在 Segments.ai 上的版本选项卡中查看您的版本是否仍在创建中。

现在，我们将通过 Segments.ai Python SDK 将该版本转换为 Hugging Face 数据集。如果您尚未设置 Segments Python 客户端，请按照notebook中“在 Segments.ai 上设置标注任务”部分的说明进行操作。

请注意，转换可能需要一段时间，具体取决于您的数据集大小。

from segments.huggingface import release2dataset

release = segments_client.get_release(dataset_identifier, release_name)
hf_dataset = release2dataset(release)

如果我们检查新数据集的特征，我们可以看到图像列和相应的标签。标签由两部分组成：注释列表和分割位图。注释对应于图像中的不同对象。对于每个对象，注释包含一个 `id` 和一个 `category_id`。分割位图是一个图像，其中每个像素都包含该像素处对象的 `id`。更多信息可以在相关文档中找到。

对于语义分割，我们需要一个包含每个像素的 `category_id` 的语义位图。我们将使用 Segments.ai SDK 中的 `get_semantic_bitmap` 函数将位图转换为语义位图。为了将此函数应用于数据集中的所有行，我们将使用 `dataset.map`。

from segments.utils import get_semantic_bitmap

def convert_segmentation_bitmap(example):
    return {
        "label.segmentation_bitmap":
            get_semantic_bitmap(
                example["label.segmentation_bitmap"],
                example["label.annotations"],
                id_increment=0,
            )
    }


semantic_dataset = hf_dataset.map(
    convert_segmentation_bitmap,
)

您还可以重写 `convert_segmentation_bitmap` 函数以使用批处理，并将 `batched=True` 传递给 `dataset.map`。这将显著加快映射速度，但您可能需要调整 `batch_size` 以确保进程不会耗尽内存。

我们将要微调的 SegFormer 模型需要特定名称的特征。为方便起见，我们现在就将其匹配到这个格式。因此，我们将 `image` 特征重命名为 `pixel_values`，将 `label.segmentation_bitmap` 重命名为 `label`，并丢弃其他特征。

semantic_dataset = semantic_dataset.rename_column('image', 'pixel_values')
semantic_dataset = semantic_dataset.rename_column('label.segmentation_bitmap', 'label')
semantic_dataset = semantic_dataset.remove_columns(['name', 'uuid', 'status', 'label.annotations'])

我们现在可以将转换后的数据集推送到 Hugging Face Hub。这样，您的团队和 Hugging Face 社区就可以使用它。在下一节中，我们将看到如何从 Hub 加载数据集。

hf_dataset_identifier = f"{hf_username}/{dataset_name}"

semantic_dataset.push_to_hub(hf_dataset_identifier)

使用 Hugging Face Hub 中的数据集

如果您不想创建自己的数据集，而是在 Hugging Face Hub 上找到了适合您用例的数据集，您可以在此处定义标识符。

例如，您可以使用完整的人行道标注数据集。请注意，您可以直接在浏览器中查看示例。

hf_dataset_identifier = "segments/sidewalk-semantic"

2. 加载并准备用于训练的 Hugging Face 数据集

现在我们已经创建了一个新数据集并将其推送到 Hugging Face Hub，我们可以在一行代码中加载数据集。

from datasets import load_dataset

ds = load_dataset(hf_dataset_identifier)

让我们打乱数据集并将其分成训练集和测试集。

ds = ds.shuffle(seed=1)
ds = ds["train"].train_test_split(test_size=0.2)
train_ds = ds["train"]
test_ds = ds["test"]

我们将提取标签数量和人类可读的 ID，以便稍后正确配置分割模型。

import json
from huggingface_hub import hf_hub_download

repo_id = f"datasets/{hf_dataset_identifier}"
filename = "id2label.json"
id2label = json.load(open(hf_hub_download(repo_id=hf_dataset_identifier, filename=filename, repo_type="dataset"), "r"))
id2label = {int(k): v for k, v in id2label.items()}
label2id = {v: k for k, v in id2label.items()}

num_labels = len(id2label)

图像处理器和数据增强

SegFormer 模型期望输入的形状是特定的。为了将我们的训练数据转换为预期的形状，我们可以使用 `SegFormerImageProcessor`。我们可以使用 `ds.map` 函数预先将图像处理器应用于整个训练数据集，但这会占用大量磁盘空间。相反，我们将使用一个*转换*，它只会在数据实际使用时（即时）准备一批数据。这样，我们就可以开始训练，而无需等待进一步的数据预处理。

在我们的转换中，我们还将定义一些数据增强，使我们的模型对不同的光照条件更具鲁棒性。我们将使用 `ColorJitter` 函数从 `torchvision` 中随机改变批处理中图像的亮度、对比度、饱和度和色调。

from torchvision.transforms import ColorJitter
from transformers import SegformerImageProcessor

processor = SegformerImageProcessor()
jitter = ColorJitter(brightness=0.25, contrast=0.25, saturation=0.25, hue=0.1) 

def train_transforms(example_batch):
    images = [jitter(x) for x in example_batch['pixel_values']]
    labels = [x for x in example_batch['label']]
    inputs = processor(images, labels)
    return inputs


def val_transforms(example_batch):
    images = [x for x in example_batch['pixel_values']]
    labels = [x for x in example_batch['label']]
    inputs = processor(images, labels)
    return inputs


# Set transforms
train_ds.set_transform(train_transforms)
test_ds.set_transform(val_transforms)

3. 微调 SegFormer 模型

加载模型进行微调

SegFormer 的作者定义了 5 种大小递增的模型：B0 到 B5。下表（摘自原始论文）显示了这些不同模型在 ADE20K 数据集上的性能，并与其他模型进行了比较。

在这里，我们将加载最小的 SegFormer 模型 (B0)，它在 ImageNet-1k 上预训练。它的大小只有大约 14MB！使用一个小型模型将确保我们的模型可以在我们的披萨配送机器人上流畅运行。

from transformers import SegformerForSemanticSegmentation

pretrained_model_name = "nvidia/mit-b0" 
model = SegformerForSemanticSegmentation.from_pretrained(
    pretrained_model_name,
    id2label=id2label,
    label2id=label2id
)

设置 Trainer

为了对我们的数据进行模型微调，我们将使用 Hugging Face 的 Trainer API。我们需要设置训练配置和评估指标才能使用 Trainer。

首先，我们将设置 `TrainingArguments`。这定义了所有训练超参数，例如学习率、 epoch 数量、模型保存频率等。我们还指定在训练后将模型推送到 Hub (`push_to_hub=True`) 并指定模型名称 (`hub_model_id`)。

from transformers import TrainingArguments

epochs = 50
lr = 0.00006
batch_size = 2

hub_model_id = "segformer-b0-finetuned-segments-sidewalk-2"

training_args = TrainingArguments(
    "segformer-b0-finetuned-segments-sidewalk-outputs",
    learning_rate=lr,
    num_train_epochs=epochs,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    save_total_limit=3,
    evaluation_strategy="steps",
    save_strategy="steps",
    save_steps=20,
    eval_steps=20,
    logging_steps=1,
    eval_accumulation_steps=5,
    load_best_model_at_end=True,
    push_to_hub=True,
    hub_model_id=hub_model_id,
    hub_strategy="end",
)

接下来，我们将定义一个函数来计算我们想要使用的评估指标。因为我们正在进行语义分割，所以我们将使用 平均交并比 (mIoU)，可以直接在 `evaluate` 库中访问。IoU 表示分割掩码的重叠程度。平均 IoU 是所有语义类别 IoU 的平均值。请参阅这篇博客文章以获取图像分割评估指标的概述。

因为我们的模型输出的 logits 维度是 height/4 和 width/4，所以我们必须在计算 mIoU 之前对其进行上采样。

import torch
from torch import nn
import evaluate

metric = evaluate.load("mean_iou")

def compute_metrics(eval_pred):
  with torch.no_grad():
    logits, labels = eval_pred
    logits_tensor = torch.from_numpy(logits)
    # scale the logits to the size of the label
    logits_tensor = nn.functional.interpolate(
        logits_tensor,
        size=labels.shape[-2:],
        mode="bilinear",
        align_corners=False,
    ).argmax(dim=1)

    pred_labels = logits_tensor.detach().cpu().numpy()
    metrics = metric.compute(
        predictions=pred_labels,
        references=labels,
        num_labels=len(id2label),
        ignore_index=0,
        reduce_labels=processor.do_reduce_labels,
    )
    
    # add per category metrics as individual key-value pairs
    per_category_accuracy = metrics.pop("per_category_accuracy").tolist()
    per_category_iou = metrics.pop("per_category_iou").tolist()

    metrics.update({f"accuracy_{id2label[i]}": v for i, v in enumerate(per_category_accuracy)})
    metrics.update({f"iou_{id2label[i]}": v for i, v in enumerate(per_category_iou)})
    
    return metrics

最后，我们可以实例化一个 `Trainer` 对象。

from transformers import Trainer

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_ds,
    eval_dataset=test_ds,
    compute_metrics=compute_metrics,
)

现在我们的训练器已设置完毕，训练就像调用 `train` 函数一样简单。我们无需担心 GPU 的管理，训练器会处理好一切。

trainer.train()

训练完成后，我们可以将微调后的模型和图像处理器推送到 Hub。

这也会自动创建一个包含我们模型结果的模型卡。我们将在 `kwargs` 中提供一些额外信息，使模型卡更完整。

kwargs = {
    "tags": ["vision", "image-segmentation"],
    "finetuned_from": pretrained_model_name,
    "dataset": hf_dataset_identifier,
}

processor.push_to_hub(hub_model_id)
trainer.push_to_hub(**kwargs)

4. 推理

现在到了激动人心的部分，使用我们微调过的模型！在本节中，我们将展示如何从 Hub 加载模型并将其用于推理。

但是，您也可以直接在 Hugging Face Hub 上试用您的模型，这得益于由托管推理 API 提供支持的酷炫小部件。如果您在之前的步骤中将模型推送到了 Hub，您应该会在模型页面上看到一个推理小部件。您可以通过在模型卡中定义示例图像 URL，为该小部件添加默认示例。请参阅此模型卡作为示例。

使用 Hugging Face Hub 中的模型

我们首先使用 `SegformerForSemanticSegmentation.from_pretrained()` 从 Hub 加载模型。

from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation

processor = SegformerImageProcessor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
model = SegformerForSemanticSegmentation.from_pretrained(f"{hf_username}/{hub_model_id}")

接下来，我们将从我们的测试数据集中加载一张图片。

image = test_ds[0]['pixel_values']
gt_seg = test_ds[0]['label']
image

为了分割这张测试图像，我们首先需要使用图像处理器准备图像。然后我们通过模型对其进行前向传播。

我们还需要记住将输出的 logits 缩放到原始图像大小。为了获得实际的类别预测，我们只需要对 logits 应用 `argmax`。

from torch import nn

inputs = processor(images=image, return_tensors="pt")
outputs = model(**inputs)
logits = outputs.logits  # shape (batch_size, num_labels, height/4, width/4)

# First, rescale logits to original image size
upsampled_logits = nn.functional.interpolate(
    logits,
    size=image.size[::-1], # (height, width)
    mode='bilinear',
    align_corners=False
)

# Second, apply argmax on the class dimension
pred_seg = upsampled_logits.argmax(dim=1)[0]

现在是时候展示结果了。我们将在旁边展示结果以及地面实况掩码。

您觉得呢？您会用这些分割信息将我们的披萨配送机器人送上路吗？

结果可能还不完美，但我们总是可以扩展我们的数据集以使模型更健壮。我们现在还可以训练一个更大的 SegFormer 模型，看看它的表现如何。

5. 结论

就是这样！您现在知道如何创建自己的图像分割数据集以及如何使用它来微调语义分割模型。

我们在此过程中向您介绍了一些有用的工具，例如：

• Segments.ai 用于标注您的数据
• 🤗 datasets 用于创建和共享数据集
• 🤗 transformers 用于轻松微调最先进的分割模型
• Hugging Face Hub 用于共享我们的数据集和模型，以及为我们的模型创建推理小部件

我们希望您喜欢这篇帖子并有所收获。请随时在 Twitter 上与我们分享您的模型（@TobiasCornille、@NielsRogge 和 @huggingface）。