Классификация изображений

Классификация изображений присваивает изображению метку или класс. В отличие от классификации текста или аудио, входными данными здесь являются значения пикселей, из которых состоит изображение. У классификации изображений существует множество применений, например, обнаружение повреждений после стихийных бедствий, мониторинг состояния посевов или помощь в скрининге медицинских изображений на наличие признаков заболеваний.

В этом руководстве показано, как:

1. Настройте ViT на датасете Food-101 для классификации блюда на изображении.
2. Используйте дообученную модель для инференса.

Чтобы увидеть все архитектуры и контрольные точки, совместимые с этой задачей, рекомендуем обратиться к странице задачи

Прежде чем начать, убедитесь, что у вас установлены все необходимые библиотеки:

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

Мы рекомендуем вам войти в свой аккаунт Hugging Face, чтобы загрузить модель и поделиться ею с сообществом. При появлении запроса введите токен для входа:

>>> from huggingface_hub import notebook_login

>>> notebook_login()

Загрузка набора данных Food-101

Начните с загрузки небольшого подмножества набора данных Food-101 из библиотеки 🤗 Datasets. Это даст вам возможность поэкспериментировать и убедиться, что всё работает, прежде чем тратить больше времени на обучение на полном наборе данных.

>>> from datasets import load_dataset

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

Разбейте train-подмножество датасета на обучающую и тестовую выборки с помощью метода train_test_split:

>>> 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: класс метки для предмета еды

Чтобы упростить модели получение имени метки по её идентификатору, создайте словарь, который сопоставляет имя метки с целым числом и наоборот:

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

Теперь вы можете преобразовать идентификатор метки в её название:

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

Предобработка

Следующий шаг — загрузить процессор изображений ViT, чтобы преобразовать изображение в тензор:

>>> from transformers import AutoImageProcessor

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

Примените к изображениям некоторые преобразования, чтобы сделать модель более устойчивой к переобучению. Здесь вы используете модуль transforms из torchvision, но можете использовать любую другую библиотеку для работы с изображениями.

Вырежьте случайную часть изображения, измените её размер и нормализуйте с помощью среднего значения и стандартного отклонения изображения:

>>> 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

Чтобы применить функцию предобработки ко всему набору данных, используйте метод with_transform в 🤗 Datasets. Преобразования применяются на лету при загрузке элемента набора данных:

>>> food = food.with_transform(transforms)

Теперь создайте пакет примеров с помощью DefaultDataCollator. В отличие от других collator'ов данных в 🤗 Transformers, DefaultDataCollator не выполняет дополнительное предобработка, например, выравнивание (padding).

>>> 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, загляните в базовое руководство здесь!

Теперь вы готовы приступить к обучению модели! Загрузите ViT с помощью AutoModelForImageClassification. Укажите количество меток, а также ожидаемое число меток и соответствия меток:

>>> 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, который указывает, куда сохранять модель. Вы сможете загрузить эту модель в Hub, установив push_to_hub=True (для загрузки модели необходимо войти в аккаунт Hugging Face). В конце каждой эпохи Trainer будет оценивать точность и сохранять контрольную точку обучения.
2. Передайте аргументы обучения в Trainer вместе с моделью, набором данных, токенизатором, data collator и функцией 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_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.

Вывод

Отлично, теперь, когда вы дообучили модель, вы можете использовать её для инференса!

Загрузите изображение, для которого вы хотите выполнить инференс:

>>> 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")

Передайте входные данные в модель и получите логиты:

>>> 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'