Open-Source AI Cookbook documentation

用自定义生物医学数据集微调视觉 Transformer 模型

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Open In Colab

用自定义生物医学数据集微调视觉 Transformer 模型

作者: Emre Albayrak

本指南概述了在自定义生物医学数据集上微调视觉 transformer(ViT)模型的过程。它包括加载数据集和准备数据集的步骤,为不同的数据拆分设置图像转换,配置和初始化 ViT 模型,以及定义具有评估和可视化工具的训练过程。

数据集信息

自定义数据集是手工制作的,包含 780 张图片,分为 3 类(良性,恶性,正常)。

attachment:datasetinfo.png

模型信息

我们所要微调的模型是 Google 的 "vit-large-patch16-224" 模型,该模型可以在 Hugging Face 的模型库中找到。这个模型是在 ImageNet-21k 数据集上进行预训练的,该数据集包含 1400 万张图片和 21,843 个类别。之后,它在 ImageNet 2012 数据集上进行了微调,该数据集有 100 万张图片和 1000 个类别,图像分辨率统一为 224x224。Google 还提供了其他几种 ViT(Vision Transformer)模型,它们具有不同的图像尺寸和分割块大小。

现在,让我们开始吧。

开始

首先,让我们安装库。

!pip install datasets transformers accelerate torch scikit-learn matplotlib wandb

(可选) 我们会把我们的模型推送到 hugging face hub 上,所以我们必须登录。

# from huggingface_hub import notebook_login
# notebook_login()

数据集准备

数据集库会自动从数据集中拉取图片和类别。若需更详细的信息,你可以访问这个链接

from datasets import load_dataset

dataset = load_dataset("emre570/breastcancer-ultrasound-images")
dataset

我们已获得数据集。但我们没有验证集。为了创建验证集,我们将根据测试集的大小计算验证集作为训练集的一部分的大小。然后我们将训练数据集拆分为新的训练和验证子集。

# Get the numbers of each set
test_num = len(dataset["test"])
train_num = len(dataset["train"])

val_size = test_num / train_num

train_val_split = dataset["train"].train_test_split(test_size=val_size)
train_val_split

我们已获得分离的训练集。现在让我们将其与测试集合并。

from datasets import DatasetDict

dataset = DatasetDict(
    {"train": train_val_split["train"], "validation": train_val_split["test"], "test": dataset["test"]}
)
dataset

完美!我们的数据集已经准备好了。现在我们将子集分配给不同的变量。我们将在后面使用它们以便于引用。

train_ds = dataset["train"]
val_ds = dataset["validation"]
test_ds = dataset["test"]

我们可以看到这张图片是一个与标签关联的PIL.Image对象。

train_ds[0]

我们还可以查看训练集的特征。

train_ds.features

让我们从数据集中显示每个类的一个图像。

>>> import matplotlib.pyplot as plt

>>> # Initialize a set to keep track of shown labels
>>> shown_labels = set()

>>> # Initialize the figure for plotting
>>> plt.figure(figsize=(10, 10))

>>> # Loop through the dataset and plot the first image of each label
>>> for i, sample in enumerate(train_ds):
...     label = train_ds.features["label"].names[sample["label"]]
...     if label not in shown_labels:
...         plt.subplot(1, len(train_ds.features["label"].names), len(shown_labels) + 1)
...         plt.imshow(sample["image"])
...         plt.title(label)
...         plt.axis("off")
...         shown_labels.add(label)
...         if len(shown_labels) == len(train_ds.features["label"].names):
...             break

>>> plt.show()

数据处理

数据集已经准备好了。但我们还没有为微调做好准备。我们将按照以下步骤进行:

  • 标签映射: 我们将标签 ID 与其对应的名字之间进行转换,这对于模型训练和评估非常有用。

  • 图像处理: 然后,我们使用 ViTImageProcessor 来标准化输入图像的大小,并应用特定于预训练模型的归一化。同时,我们还将为训练、验证和测试定义不同的转换,以使用 torchvision 提高模型的泛化能力。

  • 转换函数: 实现函数以将转换应用于数据集,将图像转换为 ViT 模型所需格式和尺寸。

  • 数据加载: 设置一个自定义的整理函数以正确地批量处理图像和标签,并创建一个 DataLoader 以在模型训练期间高效地加载数据和批量处理。

  • 批量准备: 检索并显示样本批量中数据的形状,以验证处理是否正确并为模型输入做好准备。

标签映射

id2label = {id: label for id, label in enumerate(train_ds.features["label"].names)}
label2id = {label: id for id, label in id2label.items()}
id2label, id2label[train_ds[0]["label"]]

图像处理

from transformers import ViTImageProcessor

model_name = "google/vit-large-patch16-224"
processor = ViTImageProcessor.from_pretrained(model_name)
from torchvision.transforms import (
    CenterCrop,
    Compose,
    Normalize,
    RandomHorizontalFlip,
    RandomResizedCrop,
    ToTensor,
    Resize,
)

image_mean, image_std = processor.image_mean, processor.image_std
size = processor.size["height"]

normalize = Normalize(mean=image_mean, std=image_std)

train_transforms = Compose(
    [
        RandomResizedCrop(size),
        RandomHorizontalFlip(),
        ToTensor(),
        normalize,
    ]
)
val_transforms = Compose(
    [
        Resize(size),
        CenterCrop(size),
        ToTensor(),
        normalize,
    ]
)
test_transforms = Compose(
    [
        Resize(size),
        CenterCrop(size),
        ToTensor(),
        normalize,
    ]
)

创建转换函数

def apply_train_transforms(examples):
    examples["pixel_values"] = [train_transforms(image.convert("RGB")) for image in examples["image"]]
    return examples


def apply_val_transforms(examples):
    examples["pixel_values"] = [val_transforms(image.convert("RGB")) for image in examples["image"]]
    return examples


def apply_test_transforms(examples):
    examples["pixel_values"] = [val_transforms(image.convert("RGB")) for image in examples["image"]]
    return examples

将转换函数应用于每个集合

train_ds.set_transform(apply_train_transforms)
val_ds.set_transform(apply_val_transforms)
test_ds.set_transform(apply_test_transforms)
train_ds.features
train_ds[0]

看起来我们将像素值转换成了张量。

数据加载

import torch
from torch.utils.data import DataLoader


def collate_fn(examples):
    pixel_values = torch.stack([example["pixel_values"] for example in examples])
    labels = torch.tensor([example["label"] for example in examples])
    return {"pixel_values": pixel_values, "labels": labels}


train_dl = DataLoader(train_ds, collate_fn=collate_fn, batch_size=4)

批量准备

>>> batch = next(iter(train_dl))
>>> for k, v in batch.items():
...     if isinstance(v, torch.Tensor):
...         print(k, v.shape)
pixel_values torch.Size([4, 3, 224, 224])
labels torch.Size([4])

完美!现在我们为微调过程做好了准备。

微调模型

现在我们将配置和微调模型。我们首先使用特定的标签映射和预训练设置初始化模型,调整大小不匹配的问题。训练参数被设置用来定义模型的学习过程,包括保存策略、批量大小和训练轮次,结果将通过 Weights & Biases 进行记录。Hugging Face Trainer 然后将实例化以管理训练和评估,利用自定义数据整理器和模型的内置处理器。最后,在训练之后,模型的性能将在测试数据集上进行评估,并打印指标以评估其准确性。

首先,我们调用我们的模型。

from transformers import ViTForImageClassification

model = ViTForImageClassification.from_pretrained(
    model_name, id2label=id2label, label2id=label2id, ignore_mismatched_sizes=True
)

这里有一个微妙的细节。ignore_mismatched_sizes 参数。

当你在一个新数据集上微调预训练模型时,有时你的图像的输入大小或模型架构的特定细节(比如分类层中的标签数量)可能与模型最初训练时的大小不完全匹配。这可能会由于各种原因而发生,例如当你在完全不同类型的图像数据(如医学图像或专业相机图像)上使用在一种类型的图像数据(如ImageNet中的自然图像)上训练的模型时。 将 ignore_mismatched_sizes 设置为 True 允许模型调整其层以适应大小差异,而不会抛出错误。

例如,这个模型训练的类数是1000,即 torch.Size([1000]),它期望一个具有 torch.Size([1000]) 类的输入。我们的数据集有3类,即 torch.Size([3]) 类。如果我们直接给它,它会抛出错误,因为类别数量不匹配。

然后,为这个模型定义来自谷歌的训练参数。

(可选) 注意,由于我们将 report_to 参数设置为 wandb,指标将被保存在 Weights & Biases 中。W&B 将要求你提供一个 API 密钥,因此你应该创建一个账户和一个API密钥。如果你不希望这样做,你可以删除 report_to 参数。

from transformers import TrainingArguments, Trainer
import numpy as np

train_args = TrainingArguments(
    output_dir="output-models",
    save_total_limit=2,
    report_to="wandb",
    save_strategy="epoch",
    evaluation_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=10,
    per_device_eval_batch_size=4,
    num_train_epochs=40,
    weight_decay=0.01,
    load_best_model_at_end=True,
    logging_dir="logs",
    remove_unused_columns=False,
)

我们现在可以使用 Trainer 开始微调过程。

trainer = Trainer(
    model,
    train_args,
    train_dataset=train_ds,
    eval_dataset=val_ds,
    data_collator=collate_fn,
    tokenizer=processor,
)
trainer.train()
Epoch 训练损失 验证损失 准确率
40 0.174700 0.596288 0.903846

微调过程已完成。接下来,我们继续使用测试集评估模型。

>>> outputs = trainer.predict(test_ds)
>>> print(outputs.metrics)
{'test_loss': 0.40843912959098816, 'test_runtime': 4.9934, 'test_samples_per_second': 31.242, 'test_steps_per_second': 7.81}

{'test_loss': 0.3219967782497406, 'test_accuracy': 0.9102564102564102, 'test_runtime': 4.0543, 'test_samples_per_second': 38.478, 'test_steps_per_second': 9.619}

(可选) 将模型推送到 Hub

我们可以使用 push_to_hub 将我们的模型推送到 Hugging Face Hub。

model.push_to_hub("your_model_name")

太棒了!让我们可视化结果。

结果

我们已经完成了微调。让我们看看我们的模型是如何预测类别的,使用 scikit-learn 的混淆矩阵显示,并展示召回率。

什么是混淆矩阵?

混淆矩阵是一种特定的表格布局,它允许可视化算法的性能,通常是监督学习模型,在一组已知真实值的测试数据上。它特别有用,因为可以检查分类模型的性能,因为它显示了真实标签与预测标签的频率。

让我们绘制我们模型的混淆矩阵。

>>> from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay

>>> y_true = outputs.label_ids
>>> y_pred = outputs.predictions.argmax(1)

>>> labels = train_ds.features["label"].names
>>> cm = confusion_matrix(y_true, y_pred)
>>> disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=labels)
>>> disp.plot(xticks_rotation=45)

什么是召回率?

召回率是分类任务中使用的性能指标,用于衡量模型正确识别数据集中所有相关实例的能力。具体来说,召回率评估了模型正确预测为阳性的实际阳性比例。

让我们使用 scikit-learn 打印召回率。

>>> from sklearn.metrics import recall_score

>>> # Calculate the recall scores
>>> # 'None' calculates recall for each class separately
>>> recall = recall_score(y_true, y_pred, average=None)

>>> # Print the recall for each class
>>> for label, score in zip(labels, recall):
...     print(f"Recall for {label}: {score:.2f}")
Recall for benign: 0.90
Recall for malignant: 0.86
Recall for normal: 0.78

良性召回率为0.90, 恶性召回率为0.86, 正常召回率为0.78

结论

在这个指南中,我们介绍了如何使用医学数据集训练一个 ViT 模型。它涵盖了关键的步骤,如数据集准备、图像预处理、模型配置、训练、评估和结果可视化。通过利用 Hugging Face 的 Transformers 库、scikit-learn 和 PyTorch Torchvision,它促进了高效的模型训练和评估,提供了关于模型性能及其准确分类生物医学图像能力的宝贵见解。

< > Update on GitHub