使用自定義資料集微調語義分割模型

本指南展示瞭如何微調 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）。