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

作者: Emre Albayrak

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

##
数据集信息

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

##
模型信息

我们所要微调的模型是 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)

    &#123;'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