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

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

本指南展示了如何对Segformer进行微调,这是一种最先进的语义分割模型。我们的目标是为披萨送餐机器人构建一个模型,以便它可以看到驶向何处并识别障碍物🍕🤖。我们首先将在Segments.ai上标记一组人行道图像。然后,我们将使用🤗 transformers库对预训练的SegFormer模型进行微调,该库是一个开源库,提供了易于使用的最先进模型的实现。在此过程中,您将学习如何使用Hugging Face Hub,这是最大的开源模型和数据集目录。

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

由于语义分割是一种分类类型,用于图像分类和语义分割的网络架构非常相似。2014年,Long等人的一篇开创性论文使用卷积神经网络进行了语义分割。最近,Transformers已被用于图像分类(例如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

任何ML项目的第一步都是组装一个好的数据集。为了训练语义分割模型,我们需要一个带有语义分割标签的数据集。我们可以使用Hugging Face Hub上的现有数据集,例如ADE20k,或创建自己的数据集。

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

我们不希望我们的送餐机器人感到困惑,因此我们将使用我们自己在人行道上拍摄的图像创建我们自己的语义分割数据集。我们将在下一步中展示如何标记我们拍摄的图像。如果您只想使用我们已完成的标记数据集,可以跳过“创建自己的数据集”部分,从“使用Hub上的数据集”继续。

创建自己的数据集

为了创建您的语义分割数据集,您需要两样东西:

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

我们事先拍摄了比利时人行道的一千张图像。收集和标记这样的数据集可能需要很长时间,因此您可以从一个较小的数据集开始,并在模型表现不够好时扩展它。

以下是人行道数据集中原始图像的一些示例。

为了获得分割标签,我们需要指定这些图像中所有区域/对象的类别。这可能是一个耗时的任务,但使用合适的工具可以显著加快任务的完成速度。为了标记,我们将使用Segments.ai,因为它提供了用于图像分割的智能标记工具和易于使用的Python SDK。

在Segments.ai上设置标记任务

首先,在https://segments.ai/join上创建一个帐户。接下来,创建一个新的数据集并上传您的图像。您可以使用Web界面或通过Python SDK(参见笔记本)完成此操作。

标记图像

现在原始数据已加载完成,请前往segments.ai/home并打开新创建的数据集。点击“开始标记”并创建分割掩码。您可以使用基于机器学习的超像素和自动分割工具来加速标记。

提示:使用超像素工具时,可以滚动改变超像素大小,并点击并拖动以选择分割区域。

将结果推送到Hugging Face Hub

完成标记后,创建一个包含标记数据的新数据集发布。您可以在Segments.ai上的发布选项卡上完成此操作,或者通过SDK以示例笔记本所示进行编程。

请注意,创建发布可能需要几秒钟的时间。您可以在Segments.ai上的发布选项卡上查看发布是否还在创建中。

现在,我们将通过Segments.ai Python SDK将发布转换为Hugging Face数据集。如果您尚未设置Segments Python客户端,请按照示例笔记本中的“在Segments.ai上设置标注任务”部分的说明进行设置。

请注意,转换的时间可能会很长,具体取决于数据集的大小。

from segments.huggingface import release2dataset

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

如果我们检查新数据集的特征,我们可以看到图像列和相应的标签。标签由两部分组成:注释列表和分割位图。注释对应图像中的不同对象。对于每个对象,注释包含一个idcategory_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)

使用来自Hub的数据集

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

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

hf_dataset_identifier = "segments/sidewalk-semantic"

既然我们已经创建了一个新的数据集并将其推送到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模型期望输入具有特定的形状。为了将我们的训练数据转换为匹配预期形状的数据,我们可以使用 SegFormerFeatureExtractor 。我们可以使用 ds.map 函数提前将特征提取器应用于整个训练数据集,但这可能会占用大量磁盘空间。相反,我们将使用 transform ,当需要使用数据时才会准备一批数据(即时处理)。这样,我们可以开始训练而不必等待进一步的数据预处理。

在我们的 transform 中,我们还将定义一些数据增强操作,以使我们的模型对不同的光照条件更具鲁棒性。我们将使用来自 torchvisionColorJitter 函数来随机改变批次中图像的亮度、对比度、饱和度和色调。

from torchvision.transforms import ColorJitter
from transformers import SegformerFeatureExtractor

feature_extractor = SegformerFeatureExtractor()
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 = feature_extractor(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 = feature_extractor(images, labels)
    return inputs


# 设置 transforms
train_ds.set_transform(train_transforms)
test_ds.set_transform(val_transforms)

加载模型进行微调

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
)

设置训练器

为了在我们的数据上对模型进行微调,我们将使用Hugging Face的Trainer API。我们需要设置训练配置和一个评估指标来使用Trainer。

首先,我们将设置 TrainingArguments 。这定义了所有的训练超参数,如学习率、迭代次数、保存模型的频率等。我们还指定在训练后将模型推送到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的平均值。请参阅这篇博文,了解有关图像分割评估指标的概述。

由于我们的模型输出的是尺寸为高度/4和宽度/4的logits,所以在计算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)
    # 将logits缩放到与标签相同的大小
    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()
    # 目前使用_compute而不是compute
    # 请参阅此问题以了解更多信息:https://github.com/huggingface/evaluate/pull/328#issuecomment-1286866576
    metrics = metric._compute(
            predictions=pred_labels,
            references=labels,
            num_labels=len(id2label),
            ignore_index=0,
            reduce_labels=feature_extractor.do_reduce_labels,
        )
    
    # 将每个类别的指标作为独立的键值对添加
    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,
}

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

现在是令人兴奋的部分,使用我们微调的模型!在本节中,我们将展示如何从Hub加载模型并用于推理。

但是,您也可以直接在Hugging Face Hub上尝试您的模型,感谢托管推理API提供的酷炫小部件。如果您在上一步中将模型推送到Hub,您应该在模型页面上看到一个推理小部件。您可以通过在模型卡片中定义示例图像URL来向小部件添加默认示例。请参阅此模型卡片作为示例。

从Hub使用模型

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

from transformers import SegformerFeatureExtractor, SegformerForSemanticSegmentation

feature_extractor = SegformerFeatureExtractor.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 = feature_extractor(images=image, return_tensors="pt")
outputs = model(**inputs)
logits = outputs.logits  # 形状为 (batch_size, num_labels, height/4, width/4)

# 首先,将logits放缩到原始图片的大小
upsampled_logits = nn.functional.interpolate(
    logits,
    size=image.size[::-1], # (height, width)
    mode='bilinear',
    align_corners=False
)

# 其次,对类别维度应用argmax
pred_seg = upsampled_logits.argmax(dim=1)[0]

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

你觉得怎么样?基于这个分割信息,你会让我们的披萨送餐机器人上路吗?

结果可能还不完美,但我们可以通过扩展数据集来使模型更加稳健。我们现在还可以训练一个更大的SegFormer模型,并看看它的表现如何。

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

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

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

我们希望您喜欢这篇文章并从中学到了一些东西。随时在Twitter上与我们分享您自己的模型(@TobiasCornille,@NielsRogge和@huggingface)。