Предобученные модели в torchvision: ResNet, EfficientNet и ViT

Обучение глубокой нейронной сети на ImageNet с нуля требует миллионов размеченных изображений и недель вычислений на GPU. Предварительно обученные модели уже выучили общие визуальные признаки — границы, текстуры, формы и части объектов высокого уровня — на основе этого огромного набора данных.

Повторно используя эти веса, Вы получаете пользу от обучения на 1,2 миллиона изображений, не оплачивая вычислительную стоимость обучения. В этом заключается основная идея переноса обучения: признаки, выученные на одной крупной задаче, хорошо переносятся на связанные задачи меньшего масштаба. torchvision.models предоставляет десятки предварительно обученных архитектур, готовых к загрузке и использованию.

ResNet-50 (остаточная сеть с 50 слоями) ввела пропускающие соединения, которые напрямую добавляют вход блока к его выходу: output = F(x) + x. Благодаря этому градиенты могут проходить непосредственно через операцию сложения, что позволяет обучать очень глубокие сети без затухания градиентов.

ResNet-50 содержит примерно 25 миллионов параметров и состоит из одного начального свёрточного слоя 7×7, максимального объединения, четырёх остаточных блоков (layer1–layer4) и слоя глобального усредняющего объединения, за которым следует полносвязная голова на 1000 классов для классификации изображений ImageNet. Именно последний слой fc заменяется для пользовательских задач.

Перед выполнением вывода входные данные должны соответствовать предварительной обработке, использовавшейся при обучении на ImageNet: их нужно изменить до размера не менее 224×224 и нормализовать со средним и стандартным отклонением ImageNet. API torchvision.transforms выполняет эту обработку. Перед выводом всегда вызывайте model.eval(), чтобы отключить прореживание и режим обучения пакетной нормализации.

Модель выдаёт 1000 логитов — по одному на каждый класс ImageNet. Мы применяем softmax, чтобы получить вероятности, и выбираем классы с наибольшими значениями top-k. Теперь torchvision.models включает названия категорий в метаданные весов, поэтому отдельный файл с метками больше не нужен.

EfficientNet (2019) представила составное масштабирование: систематическое одновременное масштабирование ширины, глубины и разрешения входных данных с использованием одного составного коэффициента. Вместо произвольного увеличения ширины или глубины сети EfficientNet уравновешивает все три измерения, обеспечивая оптимальный компромисс между точностью и эффективностью.

Семейство EfficientNet простирается от efficientnet_b0 (5 млн параметров) до efficientnet_b7 (66 млн параметров). EfficientNet-B0 достигает сопоставимой с ResNet-50 точности, используя в 8 раз меньше параметров и в 6 раз меньше FLOP, поэтому хорошо подходит для развёртывания на мобильных и периферийных устройствах. PyTorch предоставляет все восемь вариантов.

Vision Transformers (ViT) (2020) применяют архитектуру трансформера непосредственно к изображениям, без каких-либо свёрток. Изображение разбивается на сетку фрагментов фиксированного размера (например, 16×16 пикселей), каждый фрагмент разворачивается и проецируется в вектор встраивания. Эти векторы фрагментов обрабатываются подобно токенам слов в NLP.

В начало последовательности фрагментов добавляется токен CLS (классификации). После прохождения через несколько блоков кодировщика трансформера с механизмом самовнимания выход токена CLS используется для классификации. Для превосходства над CNN ViT требует больших обучающих наборов данных, но предварительно обученные модели ViT из torchvision позволяют сразу использовать эту мощь в Ваших задачах.

Выбор между этими архитектурами зависит от Ваших ограничений и задачи. ResNet-50 — надёжный вариант по умолчанию: хорошо изученная модель, сильная базовая линия, множество руководств и реализаций. EfficientNet-B0/B2 выигрывает, когда важны скорость вывода и размер модели — например, в мобильных приложениях, системах реального времени или периферийном оборудовании.

ViT особенно хорошо показывает себя в крупномасштабных задачах и получает преимущества от самоконтролируемого предварительного обучения (DINO, CLIP). Ей требуется больше вычислительных ресурсов и памяти, чем CNN с сопоставимой точностью. Для большинства пользовательских задач классификации изображений с наборами данных среднего размера начните с EfficientNet-B2 или ResNet-50, а затем попробуйте ViT, если у Вас есть ресурсы для её дообучения.

Перед изменением предварительно обученной модели для своей задачи изучите её архитектуру, чтобы понять, какие слои нужно заменить. Используйте print(model), чтобы увидеть дерево слоёв, а для программного перебора — named_modules() или named_children().

Главная идея заключается в следующем: каждая модель torchvision заканчивается классификационной головой, рассчитанной на 1000 классов ImageNet. Чтобы адаптировать модель к задаче с num_classes, отличным от 1000, замените этот последний слой. Извлекатель признаков (всё до головы) сохраняет признаки, выученные на ImageNet.

При замене классификационной головы необходимо знать размерность признаков, выдаваемых базовой сетью (всем, кроме последнего слоя). Эта размерность является размером входа новой классификационной головы.

Распространённые размерности выхода базовой сети: ResNet-50 выдаёт 2048, EfficientNet-B0 — 1280, а ViT-B/16 — 768. Эти векторы признаков вычисляются с помощью глобального усредняющего объединения пространственных карт признаков, в результате чего для каждого изображения получается один вектор. Ваша заменяемая голова получает этот вектор на вход.

Можно быстро проверить корректность весов предварительно обученной модели, запустив её на хорошо известном тестовом изображении и проверив, совпадает ли наиболее вероятное предсказание с ожидаемой меткой. Эта проверка работоспособности подтверждает, что веса загружены правильно, а конвейер предварительной обработки настроен верно.

Помимо этой проверки, рекомендуется всегда сверяться с ожидаемым форматом входных данных в метаданных весов: размером входа (224×224 для большинства моделей), порядком каналов (RGB, а не BGR) и константами нормализации (средним и стандартным отклонением ImageNet). Использование неправильной нормализации — распространённая ошибка, приводящая к плохим результатам переноса обучения.

Типичный рабочий процесс использования предварительно обученных моделей для пользовательских задач таков: (1) загрузить предварительно обученную модель с помощью weights=...IMAGENET1K..., (2) заменить последнюю классификационную голову новой nn.Linear, размер которой соответствует количеству Ваших классов, (3) при необходимости сначала заморозить веса базовой сети и (4) обучать модель с меньшей скоростью обучения, чем при обучении с нуля.

Новая классификационная голова начинает со случайными весами и должна обучиться на Ваших данных. Базовая сеть начинает с превосходными признаками и требует лишь небольших корректировок. Поэтому дифференциальные скорости обучения — очень низкая для базовой сети и более высокая для головы — часто улучшают скорость сходимости и итоговую точность.

В этом уроке Вы узнали, что ResNet-50 использует пропускающие соединения для обучения очень глубоких сетей и выдаёт признаки размерности 2048; EfficientNet обеспечивает лучший компромисс между точностью и эффективностью благодаря составному масштабированию ширины, глубины и разрешения; а ViT применяет самовнимание трансформера к фрагментам изображений без свёрток. Далее мы изучим извлечение признаков — заморозим предварительно обученную базовую сеть и обучим только новую классификационную голову на пользовательском наборе данных.