图像分类

图像分类为图像分配一个标签或类别。与文本或音频分类不同，其输入是构成图像的像素值。图像分类有许多应用场景，例如在自然灾害后检测损坏情况、监测作物健康状况，或协助筛查医学影像中的疾病迹象。

本指南将介绍如何：

在 Food-101 数据集上微调 ViT，以分类图像中的食物。
使用您微调后的模型进行推理。

要查看与此任务兼容的所有架构和检查点，我们建议查阅任务页面

在开始之前，请确保已安装所有必要的库：

pip install transformers datasets evaluate accelerate pillow torchvision scikit-learn trackio

我们建议您登录 Hugging Face 账户，以将您的模型上传并与社区共享。当系统提示时，请输入您的 token 以完成登录：

>>> from huggingface_hub import notebook_login

>>> notebook_login()

加载 Food-101 数据集

首先，从 🤗 Datasets 库中加载 Food-101 数据集的一个较小子集。这样你可以先进行实验，确保一切正常运行，然后再花更多时间在完整数据集上进行训练。

>>> from datasets import load_dataset

>>> food = load_dataset("ethz/food101", split="train[:5000]")

使用 train_test_split 方法将数据集的 train 子集拆分为训练集和测试集：

>>> food = food.train_test_split(test_size=0.2)

然后来看一个示例：

>>> food["train"][0]
{'image': <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=512x512 at 0x7F52AFC8AC50>,
 'label': 79}

数据集中的每个示例包含两个字段：

image：该食物的 PIL 图像
label：该食品项的标签类别

为了让模型更容易根据标签 ID 获取标签名称，请创建一个字典，将标签名称映射为整数，反之亦然：

>>> labels = food["train"].features["label"].names
>>> label2id, id2label = dict(), dict()
>>> for i, label in enumerate(labels):
...     label2id[label] = str(i)
...     id2label[str(i)] = label

现在，您可以将标签 ID 转换为标签名称：

>>> id2label[str(79)]
'prime_rib'

预处理

下一步是加载一个 ViT 图像处理器，将图像处理为张量：

>>> from transformers import AutoImageProcessor

>>> checkpoint = "google/vit-base-patch16-224-in21k"
>>> image_processor = AutoImageProcessor.from_pretrained(checkpoint)

对图像应用一些图像变换，以提高模型对过拟合的鲁棒性。这里将使用 torchvision 的 transforms 模块，但您也可以使用任何您喜欢的图像库。

随机裁剪图像的一部分，调整其大小，并使用图像的均值和标准差对其进行归一化：

>>> from torchvision.transforms import RandomResizedCrop, Compose, Normalize, ToTensor

>>> normalize = Normalize(mean=image_processor.image_mean, std=image_processor.image_std)
>>> size = (
...     image_processor.size["shortest_edge"]
...     if "shortest_edge" in image_processor.size
...     else (image_processor.size["height"], image_processor.size["width"])
... )
>>> _transforms = Compose([RandomResizedCrop(size), ToTensor(), normalize])

然后创建一个预处理函数，应用这些变换并返回图像的 pixel_values——即模型的输入：

>>> def transforms(examples):
...     examples["pixel_values"] = [_transforms(img.convert("RGB")) for img in examples["image"]]
...     del examples["image"]
...     return examples

要对整个数据集应用预处理函数，请使用 🤗 Datasets 的 with_transform 方法。在加载数据集元素时，转换会即时应用：

>>> food = food.with_transform(transforms)

现在使用 DefaultDataCollator 创建一个示例批次。与 🤗 Transformers 中的其他数据整理器不同，DefaultDataCollator 不会应用填充等额外预处理。

>>> from transformers import DefaultDataCollator

>>> data_collator = DefaultDataCollator()

评估

在训练过程中包含一个指标通常有助于评估模型的性能。您可以使用 🤗 Evaluate 库快速加载评估方法。对于此任务，请加载 accuracy 指标（请参阅 🤗 Evaluate 的快速入门以了解如何加载和计算指标）：

>>> import evaluate

>>> accuracy = evaluate.load("accuracy")

然后创建一个函数，将你的预测值和标签传递给 compute 以计算准确率：

>>> import numpy as np


>>> def compute_metrics(eval_pred):
...     predictions, labels = eval_pred
...     predictions = np.argmax(predictions, axis=1)
...     return accuracy.compute(predictions=predictions, references=labels)

您的 compute_metrics 函数已准备就绪，在配置训练时您将再次用到它。

训练

如果您不熟悉使用 Trainer 对模型进行微调，请查看此处的基础教程！

现在您可以开始训练模型了！使用 AutoModelForImageClassification 加载 ViT。指定标签数量以及预期标签数，并设置标签映射：

>>> from transformers import AutoModelForImageClassification, TrainingArguments, Trainer

>>> model = AutoModelForImageClassification.from_pretrained(
...     checkpoint,
...     num_labels=len(labels),
...     id2label=id2label,
...     label2id=label2id,
... )

此时，只剩下三个步骤：

在 TrainingArguments 中定义训练超参数。请注意不要删除未使用的列，因为这会丢弃 image 列。如果没有 image 列，就无法创建 pixel_values。请设置 remove_unused_columns=False 以防止此行为！另一个必需的参数是 output_dir，用于指定模型的保存位置。通过设置 push_to_hub=True，您可以将该模型推送到 Hub（上传模型需要登录 Hugging Face）。在每个 epoch 结束时，Trainer 会评估准确率并保存训练检查点。
将训练参数与模型、数据集、分词器、数据整理器以及 compute_metrics 函数一起传递给 Trainer。
调用 train() 来微调你的模型。

>>> training_args = TrainingArguments(
...     output_dir="my_awesome_food_model",
...     remove_unused_columns=False,
...     eval_strategy="epoch",
...     save_strategy="epoch",
...     learning_rate=5e-5,
...     per_device_train_batch_size=16,
...     gradient_accumulation_steps=4,
...     per_device_eval_batch_size=16,
...     num_train_epochs=3,
...     warmup_steps=0.1,
...     logging_steps=10,
...     report_to="trackio",
...     run_name="food101",
...     load_best_model_at_end=True,
...     metric_for_best_model="accuracy",
...     push_to_hub=True,
... )

>>> trainer = Trainer(
...     model=model,
...     args=training_args,
...     data_collator=data_collator,
...     train_dataset=food["train"],
...     eval_dataset=food["test"],
...     processing_class=image_processor,
...     compute_metrics=compute_metrics,
... )

>>> trainer.train()

训练完成后，使用 push_to_hub() 方法将您的模型共享到 Hub，以便所有人都能使用您的模型：

>>> trainer.push_to_hub()

如需了解如何针对图像分类任务微调模型的更详细示例，请参阅对应的 PyTorch notebook。

推理

很好，既然你已经完成了模型微调，现在就可以用它进行推理了！

加载您想要进行推理的图像：

>>> ds = load_dataset("ethz/food101", split="validation[:10]")
>>> image = ds["image"][0]

尝试使用微调后的模型进行推理的最简单方法是将其用于 pipeline()。使用你的模型实例化一个用于图像分类的 pipeline，并将你的图像传入其中：

>>> from transformers import pipeline

>>> classifier = pipeline("image-classification", model="my_awesome_food_model")
>>> classifier(image)
[{'score': 0.31856709718704224, 'label': 'beignets'},
 {'score': 0.015232225880026817, 'label': 'bruschetta'},
 {'score': 0.01519392803311348, 'label': 'chicken_wings'},
 {'score': 0.013022331520915031, 'label': 'pork_chop'},
 {'score': 0.012728818692266941, 'label': 'prime_rib'}]

您也可以手动复现 pipeline 的结果：

加载图像处理器以预处理图像，并将 input 作为 PyTorch 张量返回：

>>> from transformers import AutoImageProcessor
>>> import torch

>>> image_processor = AutoImageProcessor.from_pretrained("my_awesome_food_model")
>>> inputs = image_processor(image, return_tensors="pt")

将输入传递给模型并返回 logits：

>>> from transformers import AutoModelForImageClassification

>>> model = AutoModelForImageClassification.from_pretrained("my_awesome_food_model")
>>> with torch.no_grad():
...     logits = model(**inputs).logits

获取概率最高的预测标签，并使用模型的 id2label 映射将其转换为标签：

>>> predicted_label = logits.argmax(-1).item()
>>> model.config.id2label[predicted_label]
'beignets'

Update on GitHub