Тонкая настройка Vision Transformer (ViT) для классификации персонажей фильмов Marvel.

Если вы просто хотите сначала опробовать модель, вы можете сделать это здесь.

В области компьютерного зрения долгое время доминировали сверточные нейронные сети (CNN). С другой стороны, Трансформеры стали архитектурой глубокого обучения для любой задачи НЛП. Но в 2021 году, с появлением Vision Transformers paper, трансформеры тоже начали участвовать в гонке задач Computer Vision.

В этой статье мы настраиваем модель Vision Transformer для обнаружения персонажей фильмов Marvel.

Немного о Vision Transformer

Vision Transformer (ViT) — это, по сути, обычная сеть преобразователей, применяемая к изображениям вместо текста. Если вы хотите понять, как работают трансформеры, это отличное видео.

Как показано на рис.1, входное изображение разбивается на отдельные патчи. Эти отдельные патчи подобны токенам в задаче НЛП, где обычно используются преобразователи. Эти отдельные части изображения сглаживаются в линейные векторы. Затем добавляются позиционные векторы для встраивания позиционной информации этих патчей. Затем они передаются в блок преобразователя энкодера. В блок преобразователя также передается специальный токен CLS, который мы можем использовать для предсказания класса изображения. Выходные данные внедрения CLS из блока преобразователя подаются в сеть прямой связи с классификатором softmax для прогнозирования класса.

Точная настройка преобразователя зрения

Подстроить ViT к нашим изображениям очень просто. Спасибо потрясающим библиотекам от Huggingface🤗.

Давайте создадим классификатор, который идентифицирует персонажей фильмов Marvel. Сначала нам нужны изображения разных персонажей Marvel.

Подготовить набор данных из изображений Google

git clone https://github.com/Joeclinton1/google-images-download.git && cd google-images-download && pip install -e . && cd ..

Приведенная выше команда устанавливает инструмент, который вы можете использовать для загрузки изображений из Google Images. После установки выполните приведенную ниже команду, чтобы начать загрузку изображений. Параметр keywords должен быть строкой различных условий поиска, каждый из которых представляет класс. limit — это количество изображений, которые вы хотите использовать для каждого класса.

googleimagesdownload  --keywords "Iron Man,Captain America,Thor,Spider Man,Docter Strage,Black Panther,Ant Man,Captain Marvel,Hulk,Black Widow,Hawkeye Avengers,Scarlet Witch,Vision Avengers,Bucky Barnes,Falcon Avengers,Loki" --limit 100

Загруженные файлы будут похожи на этот формат.

downloads/Iron Man/ironman.jpg
downloads/Iron Man/iron_man2.jpg
.
.
downloads/Loki/loki.jpg

Мы хотим, чтобы файлы были разделены на обучающие и проверочные наборы для нашей модели. Мы можем сделать это с помощью другого инструмента split-folders. Установите его:

pip install split-folders

Теперь, чтобы разделить папки на наборы train и val с разделением 80–20 процентов, выполните. Укажите имя папки после --

splitfolders --ratio .8 .2  -- downloads

Теперь вы должны увидеть папки, похожие на эту структуру.

output/train/Iron Man/ironman.jpg
output/train/Iron Man/iron_man2.jpg
.
.
output/train/Loki/loki.jpg
output/val/Iron Man/ironman3.jpg
output/val/Iron Man/iron_man22.jpg
.
output/val/Loki/loki_.jpg

Точная настройка кода

Установите необходимые библиотеки

pip install -U transformers datasets

После импорта необходимых пакетов указываем расположение нашего набора данных в строке номер 6,7. Измените его, чтобы он соответствовал местоположению в вашей системе.

Строка 12–20 прикрепляет функцию преобразования к набору данных, эта функция преобразования применяется для преобразования необработанных изображений PIL, считанных библиотекой набора данных, в тензорный формат, который подходит для подачи в модель. Мы также добавляем информацию о этикетке в пакет в функции преобразования.

Строки 21–30 определяют метрику для расчета точности. Строка 34–41 определяет модель и средство извлечения признаков из контрольной точки, заданной model_name_or_path . Работа экстрактора признаков заключается в применении любых операций нормализации, стандартизации и изменения размера к изображению, чтобы сделать его готовым для загрузки в модель.

Строка 23 определяет функцию сопоставления данных, которая преобразует пакеты, поступающие в виде списков dicts, в словарь пакетных тензоров.

Строка 44 и далее определяет объект Trainer с параметрами для обучения. Обратитесь к документации, чтобы получить более глубокое представление о каждом параметре. trainer.train() должен начать обучение. После этого сохраняем модель и логируем метрики. Будет создана папка, которая будет содержать вес модели и конфигурацию. Путь к этой папке указан output_dir в классе TrainingArguments.

Со 100 изображениями на класс и текущими настройками я получил точность теста 0,8683. Похоже, есть возможности для улучшения за счет увеличения размера набора данных или настройки некоторых параметров в TrainingArguments.

Тестирование нашего классификатора Мстителей

Приведенный выше код использует обученную модель для предсказания персонажа Мстителя. Строка 6–8 загружает модель с пути обученной модели. Строка 11–13 загружает изображение в объект PIL Image. В строках 16–17 выполняется извлечение и прогнозирование признаков. Строка 19–21 извлекает окончательные прогнозы уровня, также известные как логиты, и находит индекс класса, соответствующий максимальному значению логитов.

Последний шаг: наслаждайтесь игрой с моделью.

Входное изображение:

Predicted class: Iron Man

Входное изображение:

Predicted class: Black Widow

Дополнительный шаг: нажмите на Hub, чтобы похвастаться

Если вы хотите подтолкнуть модель к обнимающемуся хабу модели лица, это очень просто. Просто установите push_to_hub в значение True в TrainingArguments, и ваша модель будет перенесена в вашу учетную запись с обнимающим лицом. Вам также нужно будет войти в систему, чтобы обнять лицо, выполнив следующую команду.

pip install huggingface_hub
huggingface-cli login

Ссылки и ссылки