Transformers 文档

图像分类

Hugging Face's logo
加入 Hugging Face 社区

并获得增强型文档体验

开始使用

图像分类

图像分类为图像分配标签或类别。与文本或音频分类不同,输入是构成图像的像素值。图像分类有许多应用,例如检测自然灾害后的损害、监测作物健康状况或帮助筛选医学图像以寻找疾病迹象。

本指南说明了如何

  1. 微调 ViTFood-101 数据集上,对图像中的食物项目进行分类。
  2. 使用微调后的模型进行推理。

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

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

pip install transformers datasets evaluate accelerate pillow torchvision scikit-learn

我们鼓励您登录您的 Hugging Face 帐户,以上传和与社区共享您的模型。在系统提示时,输入您的令牌以登录

>>> from huggingface_hub import notebook_login

>>> notebook_login()

加载 Food-101 数据集

首先从 🤗 Datasets 库加载 Food-101 数据集的较小子集。这将让您有机会尝试并确保一切正常,然后再花更多时间在完整数据集上进行训练。

>>> from datasets import load_dataset

>>> food = load_dataset("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)
Pytorch
隐藏 Pytorch 内容

对图像应用一些图像变换,以使模型对过度拟合更加稳健。在这里,您将使用 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()
TensorFlow
隐藏 TensorFlow 内容

为了避免过度拟合,并使模型更加稳健,请在数据集的训练部分添加一些数据增强。在这里,我们使用 Keras 预处理层来定义训练数据的变换(包括数据增强)和验证数据的变换(仅中心裁剪、调整大小和归一化)。您可以使用 tf.image 或您喜欢的任何其他库。

>>> from tensorflow import keras
>>> from tensorflow.keras import layers

>>> size = (image_processor.size["height"], image_processor.size["width"])

>>> train_data_augmentation = keras.Sequential(
...     [
...         layers.RandomCrop(size[0], size[1]),
...         layers.Rescaling(scale=1.0 / 127.5, offset=-1),
...         layers.RandomFlip("horizontal"),
...         layers.RandomRotation(factor=0.02),
...         layers.RandomZoom(height_factor=0.2, width_factor=0.2),
...     ],
...     name="train_data_augmentation",
... )

>>> val_data_augmentation = keras.Sequential(
...     [
...         layers.CenterCrop(size[0], size[1]),
...         layers.Rescaling(scale=1.0 / 127.5, offset=-1),
...     ],
...     name="val_data_augmentation",
... )

接下来,创建函数以将适当的变换应用于一批图像,而不是一次应用于一张图像。

>>> import numpy as np
>>> import tensorflow as tf
>>> from PIL import Image


>>> def convert_to_tf_tensor(image: Image):
...     np_image = np.array(image)
...     tf_image = tf.convert_to_tensor(np_image)
...     # `expand_dims()` is used to add a batch dimension since
...     # the TF augmentation layers operates on batched inputs.
...     return tf.expand_dims(tf_image, 0)


>>> def preprocess_train(example_batch):
...     """Apply train_transforms across a batch."""
...     images = [
...         train_data_augmentation(convert_to_tf_tensor(image.convert("RGB"))) for image in example_batch["image"]
...     ]
...     example_batch["pixel_values"] = [tf.transpose(tf.squeeze(image)) for image in images]
...     return example_batch


... def preprocess_val(example_batch):
...     """Apply val_transforms across a batch."""
...     images = [
...         val_data_augmentation(convert_to_tf_tensor(image.convert("RGB"))) for image in example_batch["image"]
...     ]
...     example_batch["pixel_values"] = [tf.transpose(tf.squeeze(image)) for image in images]
...     return example_batch

使用 🤗 Datasets set_transform 实时应用变换

food["train"].set_transform(preprocess_train)
food["test"].set_transform(preprocess_val)

作为最后的预处理步骤,使用 DefaultDataCollator 创建一批示例。与 🤗 Transformers 中的其他数据收集器不同,DefaultDataCollator 不会应用额外的预处理,例如填充。

>>> from transformers import DefaultDataCollator

>>> data_collator = DefaultDataCollator(return_tensors="tf")

评估

在训练期间包含指标通常有助于评估模型的性能。您可以使用 🤗 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 函数现在可以使用了,您将在设置训练时返回它。

训练

Pytorch
隐藏 Pytorch 内容

如果您不熟悉使用 Trainer 微调模型,请查看基本教程 这里

您现在可以开始训练您的模型了!使用 AutoModelForImageClassification 加载 ViT。指定标签数量以及预期标签数量和标签映射

>>> from transformers import AutoModelForImageClassification, TrainingArguments, Trainer

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

此时,只剩下三个步骤了

  1. TrainingArguments 中定义您的训练超参数。重要的是,不要删除未使用的列,因为这会删除 image 列。没有 image 列,您无法创建 pixel_values。将 remove_unused_columns=False 设置为防止此行为!唯一其他必需参数是 output_dir,它指定模型的保存位置。您将通过设置 push_to_hub=True 将此模型推送到 Hub(您需要登录 Hugging Face 才能上传您的模型)。在每个纪元结束时,Trainer 将评估准确率并保存训练检查点。
  2. 将训练参数传递给 Trainer,以及模型、数据集、令牌化器、数据收集器和 compute_metrics 函数。
  3. 调用 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_ratio=0.1,
...     logging_steps=10,
...     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"],
...     tokenizer=image_processor,
...     compute_metrics=compute_metrics,
... )

>>> trainer.train()

训练完成后,使用 push_to_hub() 方法将您的模型共享到 Hub,以便每个人都可以使用您的模型

>>> trainer.push_to_hub()
TensorFlow
隐藏 TensorFlow 内容

如果您不熟悉使用 Keras 微调模型,请先查看 基本教程

要使用 TensorFlow 微调模型,请按照以下步骤操作

  1. 定义训练超参数,并设置优化器和学习率计划。
  2. 实例化预训练模型。
  3. 将 🤗 Dataset 转换为 tf.data.Dataset
  4. 编译您的模型。
  5. 添加回调,并使用 fit() 方法运行训练。
  6. 将您的模型上传到 🤗 Hub 以与社区共享。

首先定义超参数、优化器和学习率计划

>>> from transformers import create_optimizer

>>> batch_size = 16
>>> num_epochs = 5
>>> num_train_steps = len(food["train"]) * num_epochs
>>> learning_rate = 3e-5
>>> weight_decay_rate = 0.01

>>> optimizer, lr_schedule = create_optimizer(
...     init_lr=learning_rate,
...     num_train_steps=num_train_steps,
...     weight_decay_rate=weight_decay_rate,
...     num_warmup_steps=0,
... )

然后,使用 TFAutoModelForImageClassification 加载 ViT 以及标签映射

>>> from transformers import TFAutoModelForImageClassification

>>> model = TFAutoModelForImageClassification.from_pretrained(
...     checkpoint,
...     id2label=id2label,
...     label2id=label2id,
... )

使用 to_tf_dataset 和您的 data_collator 将您的数据集转换为 tf.data.Dataset 格式

>>> # converting our train dataset to tf.data.Dataset
>>> tf_train_dataset = food["train"].to_tf_dataset(
...     columns="pixel_values", label_cols="label", shuffle=True, batch_size=batch_size, collate_fn=data_collator
... )

>>> # converting our test dataset to tf.data.Dataset
>>> tf_eval_dataset = food["test"].to_tf_dataset(
...     columns="pixel_values", label_cols="label", shuffle=True, batch_size=batch_size, collate_fn=data_collator
... )

使用 compile() 为训练配置模型

>>> from tensorflow.keras.losses import SparseCategoricalCrossentropy

>>> loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
>>> model.compile(optimizer=optimizer, loss=loss)

要从预测中计算准确率并将模型推送到 🤗 Hub,请使用 Keras 回调。将您的 compute_metrics 函数传递给 KerasMetricCallback,并使用 PushToHubCallback 上传模型

>>> from transformers.keras_callbacks import KerasMetricCallback, PushToHubCallback

>>> metric_callback = KerasMetricCallback(metric_fn=compute_metrics, eval_dataset=tf_eval_dataset)
>>> push_to_hub_callback = PushToHubCallback(
...     output_dir="food_classifier",
...     tokenizer=image_processor,
...     save_strategy="no",
... )
>>> callbacks = [metric_callback, push_to_hub_callback]

最后,您可以开始训练模型了!使用训练集和验证集、纪元数以及您的回调调用 fit() 以微调模型

>>> model.fit(tf_train_dataset, validation_data=tf_eval_dataset, epochs=num_epochs, callbacks=callbacks)
Epoch 1/5
250/250 [==============================] - 313s 1s/step - loss: 2.5623 - val_loss: 1.4161 - accuracy: 0.9290
Epoch 2/5
250/250 [==============================] - 265s 1s/step - loss: 0.9181 - val_loss: 0.6808 - accuracy: 0.9690
Epoch 3/5
250/250 [==============================] - 252s 1s/step - loss: 0.3910 - val_loss: 0.4303 - accuracy: 0.9820
Epoch 4/5
250/250 [==============================] - 251s 1s/step - loss: 0.2028 - val_loss: 0.3191 - accuracy: 0.9900
Epoch 5/5
250/250 [==============================] - 238s 949ms/step - loss: 0.1232 - val_loss: 0.3259 - accuracy: 0.9890

恭喜!您已微调模型并将其共享到 🤗 Hub。您现在可以使用它进行推理了!

要更深入地了解如何微调图像分类模型,请查看相应的 PyTorch 笔记本

推理

太好了,现在你已经微调了模型,你可以用它来进行推理!

加载你想进行推理的图像

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

尝试使用微调后的模型进行推理的最简单方法是在 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 的结果

Pytorch
隐藏 Pytorch 内容

加载一个图像处理器来预处理图像,并将 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'
TensorFlow
隐藏 TensorFlow 内容

加载一个图像处理器来预处理图像,并将 input 返回为 TensorFlow 张量

>>> from transformers import AutoImageProcessor

>>> image_processor = AutoImageProcessor.from_pretrained("MariaK/food_classifier")
>>> inputs = image_processor(image, return_tensors="tf")

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

>>> from transformers import TFAutoModelForImageClassification

>>> model = TFAutoModelForImageClassification.from_pretrained("MariaK/food_classifier")
>>> logits = model(**inputs).logits

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

>>> predicted_class_id = int(tf.math.argmax(logits, axis=-1)[0])
>>> model.config.id2label[predicted_class_id]
'beignets'
< > 在 GitHub 上更新