Как быстро настроить обучение с использованием нескольких графических процессоров для оптимизации гиперпараметров с помощью PyTorch Lightning

Вероятно, большинство людей, читающих этот пост, хотя бы однажды обучили модель машинного обучения, что заняло значительное количество времени. Так было со мной, когда я настраивал RoBERTa для высокоточной классификации текста. Использование графического процессора для обучения может на порядок сократить время обучения, в зависимости от данных и модели. Это часто имеет место для сверточных нейронных сетей (CNN) и данных изображения или видеопотока, если передача пакета на GPU занимает меньше времени, чем передача его через модель.

Одна голова - хорошо, две - лучше. Это тот случай, когда доступно более одного графического процессора. Для меня одна из самых привлекательных особенностей PyTorch Lightning - это возможность бесшовного обучения с использованием нескольких графических процессоров, которая требует минимальной модификации кода. PyTorch Lightning - это оболочка поверх PyTorch, которая направлена ​​на стандартизацию рутинных разделов реализации модели машинного обучения. Дополнительный прирост скорости за счет дополнительных графических процессоров особенно удобен для трудоемких задач, таких как настройка гиперпараметров. В этой статье я хотел бы поделиться своим опытом использования PyTorch Lightining и Optuna, библиотеки Python для автоматической настройки гиперпараметров.

Данные и предварительная обработка

Для этого примера я выбрал набор данных Intel Image Classification, доступный на Kaggle. Полный код можно найти в этой записной книжке. Набор данных разделен на обучающий и тестовый подмножества, по 14034 изображения в каждом. В наборе данных шесть классов: горы, ледник, море, улица, здания и лес. Вот несколько примеров изображений:

Большинство изображений имеют размер 150 на 150 пикселей с небольшим количеством выбросов по ширине. Pytorch Lightning предоставляет базовый класс LightningDataModule для загрузки, подготовки и обслуживания выбранного набора данных для модели. Он определяет методы, которые вызываются во время обучения, проверки и тестирования. Мой дочерний класс выглядит так:

где transforms импортированы из torchvision. Случайный переворот по горизонтали используется только для увеличения обучающих данных; ToTensor() в преобразованиях предварительной обработки масштабирует все каналы до интервала [0, 1] (путем деления более 255). Этот класс также разделяет данные поезда на наборы для обучения и проверки, причем 20% изображений попадают в последний. Еще одна замечательная особенность PyTorch Lightning заключается в том, что нет необходимости указывать устройство (а) при инициализации IntelDataModule, поскольку позже все делается автоматически обучающим.

Контрольная модель

Для начала давайте соберем базовый CNN. Он будет служить эталоном для сравнения оптимизированной модели. Чтобы уменьшить количество повторений кода, давайте представим базовый строительный блок для нашей модели.

Это просто сверточный слой с необязательной нормой партии и функцией активации. Тестовая модель будет выглядеть так

Модель не имеет Linear слоев и является полностью сверточной. Хороший трюк, который добавляет немного скорости на GPU из книги Орелиена Джерона Практическое машинное обучение с scikit-learn, Keras и Tensorflow ». Методы классов ..._step и ..._epoch_end базового класса LightningModule определяют действия, которые необходимо предпринять в конце этапов обучения / проверки / тестирования и эпох соответственно. Это место для передачи выходных данных в метрические, журнальные количества, представляющие интерес. Я решил не регистрироваться после каждого шага, а скорее в конце эпохи. Это также место, где рассчитывались метрики, что может сэкономить немного времени на вычисления. Обучение модели в PyTorch Lightning выполняется с помощью класса тренера и очень просто.

Большинство аргументов говорят сами за себя с deterministic параметром, контролирующим воспроизводимость. Для демонстрационных целей я выбрал фиксированную скорость обучения 0,0001 для 20 эпох. Эта эталонная модель достигает точности 0,8213 на тестовом наборе данных, вот матрица неточностей.

Модель часто путает ледники с горами. Поскольку первые часто идут с горами на заднем плане, в этом трудно винить. В целом, это достойный, но, наверное, не лучший результат. Один из способов его улучшения - настройка гиперпараметров.

Обучение с использованием нескольких графических процессоров и оптимизация гиперпараметров

Автоматический выбор гиперпараметров, включая типы слоев, можно упростить с помощью nn.ModuleDict модулей от Pytorch:

где Residual - полностью сверточный модуль из статьи Глубокое остаточное обучение для распознавания изображений. Настраиваемая модель для оптимизации выглядит так

В приведенном выше коде ..._epoch_end методы идентичны модели Benchmark и опущены для краткости. Однако необходимо определить некоторые дополнительные методы ..._step_end. Эти методы вызываются после того, как все части пакета были оценены на всех устройствах, и результаты агрегированы здесь. backbone - это компонент модели, который будет оптимизирован. Он создается с помощью функции, которая выбирает гиперпараметры из заранее определенных наборов и передает их conv_block

Распределение вероятностей для каждого параметра контролируется Optuna, которая является библиотекой швейцарского армейского ножа для точной настройки моделей Pytorch, Tensorflow, Scikit-learn и других. Его просто настроить и легко использовать. Для этого требуется определяемая пользователем целевая функция, которая принимает trial объект, отвечающий за выбор гиперпараметров. Он создает модель из выбранных гиперпараметров и использует ее для получения оценки для оптимизации. После того, как вышеуказанная функция create_backbone создает настраиваемый компонент для экземпляра модели, он обучается, и точность проверки используется в качестве объективной оценки:

Каждая модель в приведенной выше функции обучается на двух графических процессорах, что контролируется параметрами gpus=2 и accelerator='dp'. Значение 'dp' указывает, что все графические процессоры расположены на одном компьютере. После определения objective запустить процедуру настройки несложно.

Объект выборки в приведенном выше уравнении определяет стратегию выборки гиперпараметров. TPESampler выбирает больше точек из региона с гиперпараметрами, получившими более высокие баллы, и обновляет эти регионы после каждого trial. Подобно GridSearchCV лучшим параметрам и оценке Sklearn, которые можно извлечь с помощью study.best_score и study.best_params. При оптимизированных параметрах точность улучшилась на 1,27% до 0,8340:

Резюме

Мне определенно нравится PyTorch Lightning за благородную цель упрощения и стандартизации создания моделей машинного обучения. Особенно полезна бесшовная масштабируемость распределенного обучения, которую можно получить почти бесплатно. Очевидным местом для применения увеличения скорости обучения является оптимизация гиперпараметров, которую помогает реализовать Optuna. Несмотря на то, что Optuna была одной из многих таких библиотек, ее легко настроить для моделей практически из любого фреймворка под небом. Надеюсь, этот небольшой пример поможет вам подняться по горам, которые стоят между вашими хорошими и классными моделями!