TDLR; В этой статье представлены новые улучшения среды выполнения ONNX для ускоренного обучения и описаны 4 ключевых шага для ускорения обучения существующей модели PyTorch с помощью среды выполнения ONNX (ORT).
Что такое ONNX Runtime (ORT)?
ONNX Runtime - это ориентированный на производительность механизм вывода для моделей ONNX (Open Neural Network Exchange). ONNX Runtime был разработан с упором на производительность и масштабируемость для поддержки тяжелых рабочих нагрузок в крупномасштабных производственных сценариях. Он также имеет возможности расширения для совместимости с новейшими разработками оборудования.
Недавно на // Build 2020 Microsoft анонсировала новые возможности для выполнения оптимизированного обучения с ONNX Runtime в дополнение к логическому выводу.
Эти оптимизации привели к скорости 45% собственных внутренних моделей Transformer NLP от Microsoft.
В рамках нашей работы по возвращению сообществу Microsoft разработала пример репозитория, который демонстрирует, как интегрировать обучение ORT в официальную реализацию большой модели BERT Nvidia с параметрами более 8,3 млрд.
Эту реализацию можно найти по ссылке выше, и ее можно даже обучить на своих данных с помощью Azure ML с примером записной книжки.
Однако, хотя это замечательно для начала, было бы неплохо понять, что происходит под капотом, что на первый взгляд может показаться немного подавляющим.
В оставшейся части этой статьи я рассмотрю 4 основных модификации, которые необходимо внести в модели pyTorch для использования преимуществ ORT, и укажу вам, где в репозитории примеров кода вы можете глубоко погрузиться, чтобы узнать больше.
Я не буду вдаваться в подробности модификации, но я объясню основные концепции, чтобы вы могли начать свое собственное путешествие по ОРТ.
Шаг 1. Настройка распределенной учебной среды ORT
Для обучения больших нейронных сетей часто требуются распределенные вычислительные кластеры. Таким образом, мы можем запустить версию нашего скрипта для каждого графического процессора на каждой виртуальной машине в нашем кластере. Чтобы правильно гарантировать, что наши данные и градиенты нашей модели обновляются, нам нужно назначить каждой версии нашего скрипта поезда:
- Мировой рейтинг - рейтинг процесса во всех экземплярах виртуальных машин.
- Локальный рейтинг - рейтинг для процесса сценария на данном экземпляре виртуальной машины.
Ort_supplement предоставляет функцию настройки, которая настраивает среду выполнения ONNX для распределенного обучения для Open MPI и службы машинного обучения Azure.
device = ort_supplement.setup_onnxruntime_with_mpi(args)
Шаг 2: Создайте модель тренера ОРТ
Когда у нас есть распределенная обучающая среда, следующим шагом будет загрузка модели pyTorch в ORT для обучения. Для этого мы используем обучающий метод create_ort из скрипта ort_supplement.
model = ort_supplement.create_ort_trainer(args, device, model)
Модель тренера ORT требует для реализации нескольких важных аргументов:
- Модель pyTorch в комплекте с функцией потерь
- Функция оптимизатора по умолчанию мы используем Lamb Optimizer
- описание модели с использованием объекта IODescription для объяснения входных и выходных тензорных измерений модели. Код для примера описания NVIDIA BERT выглядит следующим образом
4. Шаги накопления градиента - количество шагов, которые нужно выполнить в экземпляре скрипта перед синхронизацией градиента.
5. Версия Opset - версия набора операций для среды выполнения ONNX. Последняя версия - 12
6. Функция постобработки (необязательно), которая запускается после того, как среда выполнения ONNX преобразует модель pyTorch в ONNX, которую в дальнейшем можно использовать для поддержки неподдерживаемых операций и оптимизации графика модели ONNX. дальше
7. Атрибуты оптимизатора карты - сопоставляет имена весов с набором параметров оптимизации.
Для получения дополнительной информации ознакомьтесь со следующим ресурсом репозитория ОРТ по обучению.
Шаг 3. Назовите ОРТ "Учебные шаги по обучению модели"
После инициализации нашей модели нам нужно вызвать run_ort_training_steps, чтобы продвинуться вперед с нашей моделью, вычислить ее локальные потери и распространить ее агрегированный градиент.
loss, global_step = ort_supplement.run_ort_training_step(args, global_step, training_steps, model, batch) # Runs the actual training steps
Шаг 4. Экспорт обученной модели ONNX
Наконец, как только мы завершим все итерации распределенного обучения нашей модели. Хотя экспорт в ONNX не является обязательным для оценки, он позволяет нам воспользоваться преимуществами ускоренного логического вывода ORT. Мы можем экспортировать нашу модель ORT в формат ONNX для оценки, вызвав функцию model.save_as_onnx и предоставив ей место назначения вывода. Эту функцию также можно использовать для проверки нашей модели в каждую эпоху.
model.save_as_onnx(out_path)
Заключение
С помощью четырех вышеперечисленных функций у вас есть ключевые инструменты, необходимые для понимания полного учебного примера BERT LARGE ONNX. Чтобы увидеть их в действии, посмотрите сценарий run_pretraining_ort ниже.
Следующие шаги
Теперь, когда вы больше знакомы с тем, как использовать ORT SDK, взгляните на несколько других действительно интересных сообщений и примеров в блогах ONNX.
По мере развития среды выполнения мы всегда ищем новых участников, ознакомьтесь с нашими рекомендациями по вкладу здесь. Надеюсь, это поможет вам на пути к более эффективному глубокому обучению.
об авторе
Аарон (Ари) Борнштейн - исследователь искусственного интеллекта, увлеченный историей, занимающийся новыми технологиями и вычислительной медициной. В качестве инженера с открытым исходным кодом в команде Microsoft Cloud Developer Advocacy он сотрудничает с израильским сообществом Hi-Tech Community для решения реальных проблем с помощью технологий, меняющих правила игры, которые затем документируются, публикуются в открытом доступе и передаются остальному миру.
