Как ускорить глубокую нейросеть: обзор тензорных компиляторов

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

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

Основное отличие заключается в области применения. Компиляторы общего назначения преобразуют код с языков высокого уровня в машинный код для всех типов программных систем. Тензорные компиляторы — это узкоспециализированные инструменты для машинного обучения, которые решают задачи оптимизации операций над многомерными матрицами (тензорами) и их преобразования в низкоуровневое представление, которое собирается в исполняемый модуль.

Сравним различные характеристики и особенности компиляторов общего назначения и тензорных.

• Компиляторы общего назначения работают с исходным кодом на языках общего назначения (например, C, C++, Java).
• Тензорные компиляторы — с описаниями тензорных операций или ML-моделями (например, в формате PyTorch, TensorFlow).

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

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

Отметим, что после работы тензорного компилятора может быть использован компилятор общего назначения.

Компиляторы общего назначения в основном работают с CPU, но некоторые поддерживают компиляцию для GPU и других архитектур. Спектр аппаратных платформ для тензорных компиляторов шире: CPU, GPU, NPU, TPU, FPGA и edge-устройства. Рассмотрим их подробнее:

• CPU (архитектуры x86, ARM) — это универсальные процессоры для самых разных задач, включая операции с тензорами.
• GPU (NVIDIA CUDA, AMD ROCm) — для распараллеливаемых, вычислительно интенсивных задач.
• NPU (нейронные процессоры, например, Samsung Exynos NPU) — для оптимизации тензорных вычислений генерируется машинный код, специально адаптированный под архитектуру конкретного NPU.
• TPU (тензорные процессоры, например Google TPU) — для ускорения тензорных вычислений.
• FPGA (программируемые логические интегральные схемы) — для пользовательских вычислений с низкой задержкой в специализированных приложениях.
• Edge-устройства (NVIDIA Jetson, Apple Neural Engine) — для энергоэффективного и быстрого выполнения моделей машинного обучения.

Компиляторы общего назначения использует широкий круг инженеров и разработчиков. Они нужны для преобразования исходного кода, который написан на высокоуровневых языках программирования, в исполняемый машинный код. Такие компиляторы используются в игровой и веб-разработке, создании системного ПО, программировании встроенных систем и других областях.

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

Основные области применения тензорных компиляторов:

• В машинном обучении для ускорения выполнения нейронных сетей путем оптимизации тензорных операций, таких как матричное умножение, свертка и редукция.
• При анализе данных для ускорения конвейеров обработки, включающих масштабные тензорные операции.
• Во встроенных системах для обеспечения эффективных тензорных вычислений в условиях ограниченных ресурсов, например, на edge-устройствах: смарт-камерах, устройствах интернета вещей, носимых смарт-устройствах и так далее.

Благодаря узкой специализации тензорные компиляторы достигают более высокой эффективности для своих конкретных задач по сравнению с компиляторами общего назначения.

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

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

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

Тензорные компиляторы состоят из двух основных компонент: фронтенд и бэкенд. На разных этапах тензорные компиляторы работают с некоторым промежуточным представлением сети (intermediate representation, IR). Схема работы в общем виде показана на рисунке:

Компилятор принимает на вход модель в формате какого-либо фреймворка глубокого обучения, например PyTorch, TensorFlow или ONNX. На выходе он генерирует код, который оптимизирован для исполнения на конкретной аппаратной платформе.

Фронтенд обеспечивает чтение и преобразование модели в высокоуровневый IR — граф вычислений. Затем оптимизирует его с помощью аппаратно-независимых техник.

Бэкенд принимает оптимизированный фронтендом граф вычислений и преобразует его в низкоуровневый IR — программный код на некотором языке. Потом выполняет аппаратно-зависимые оптимизации и генерирует код под заданную платформу.

Фреймворки для оптимизации инференса можно разделить на фронтенд- и бэкенд-ориентированные. В числе первых — NNFusion, XLA, Apache TVM, OpenVINO (nGraph) и некоторые другие. Такие фреймворки ориентированы на разработчиков нейронных сетей. К ним относятся тензорные компиляторы и инструменты, которые предоставляют высокоуровневый API. Вторая группа, например Glow и MLIR, предназначена в основном для инженеров по оптимизации ПО.

В таблице приведена общая техническая информация по некоторым фреймворкам.

Нас в основном интересуют фронтенд-ориентированные решения — такие высокоуровневые инструментарии не требуют глубокого погружения в область оптимизации программного кода. Мы рассмотрим ряд «живых» проектов: Glow, XLA, OpenVINO и Apache TVM.

Если вы хотите больше узнать про построение и использование тензорных компиляторов для ускорения вывода глубоких нейронных сетей, то рекомендуем для самостоятельного изучения бесплатный курс от сотрудников института ИТММ ННГУ им. Н. И. Лобачевского. Вы пройдете путь от истории глубокого обучения до практических советов по оптимизации нейронной сети с помощью Apache TVM. Преподаватели могут использовать курс для формирования собственных образовательных программ и даже рассчитывать на небольшие методические консультации коллег из ННГУ. Курс распространяется по лицензии Apache 2.0.

Этот компилятор машинного обучения используется в качестве бэкенда для высокоуровневых фреймворков глубокого обучения, например PyTorch. Интеграция Glow в библиотеку PyTorch пока разрабатывается. Чтобы воспользоваться интеграцией в PyTorch, компилятор нужно собрать из исходников.

Компилятор поддерживает загрузку моделей в форматах библиотек Caffe2, ONNX, PyTorch и TensorFlow Lite. Есть встроенная поддержка квантизации моделей в формат весов INT8 и INT16. Запуск моделей с динамической формой входных данных отсутствует.

Минус Glow — скудность документации, из-за чего трудно разобраться, как использовать этот инструмент. Впрочем, это характерно для многих open source-проектов.

Что почитать и где скачать:

• описание принципов работы,
• исходный код на GitHub,
• лицензия: Apache 2.0.

Компилятор с открытым исходным кодом, который оптимизирует модели глубокого обучения в формате популярных фреймворков для эффективного выполнения на CPU, GPU и ускорителях машинного обучения.

XLA можно интегрировать в библиотеки TensorFlow, Keras и PyTorch (для поддержки Google TPU, NVIDIA GPU). Компилятор работает с моделями в форматах TensorFlow и Keras. Также есть поддержка квантизации моделей в формат весов INT8 и INT16. Исполнение квантованных операций пока находится в экспериментальном режиме. Ускорение вычислений обеспечивается за счет интеграции высокопроизводительных библиотек Eigen, MKL, cuDNN, TensorRT.

Что почитать и где скачать:

• документация,
• исходный код на GitHub,
• интеграция c:
  • TensorFlow,
  • PyTorch,
  • JAX,
• лицензия: Apache 2.0.

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

Ключевая особенность OpenVINO — поддержка всего «зоопарка» устройств компании Intel, включая CPU, GPU, NPU, FPGA. Он позволяет работать с моделями в форматах PyTorch, TensorFlow, TensorFlow Lite, PaddlePaddle, ONNX и Keras.

Пожалуй, OpenVINO — один из немногих инструментов, которые предоставляют такой большой спектр форматов. При этом обеспечивается поддержка понижения точности весов моделей до INT8 и FP16 и последующего запуска этих моделей. Понижение точности осуществляется с помощью Neural Network Compression Framework (NNCF), который раньше был частью проекта OpenVINO. Высокие показатели производительности инференса достигаются, в частности, за счет использования вычислительных примитивов библиотеки oneDNN.

Отметим возможность запускать модели с динамической формой входных данных. После загрузки модели перед запуском инференса можно явно установить размеры входного тензора. OpenVINO может быть интегрирован в качестве бэкенда в ONNX Runtime и PyTorch.

OpenVINO — наиболее зрелый и высокоуровневый инструментарий с точки зрения фронтенд-разработчика. Разнообразие поддерживаемых аппаратных решений позволяет написать единый код, который будет эффективно исполняться на множестве устройств. Сейчас в разработке находится поддержка RISC-V CPU. Вспомогательные инструменты, которые обеспечивают решение инфраструктурных задач, — приятный бонус.

OpenVINO предлагает подробную документацию, но без должного опыта в ней иногда сложно найти нужный фрагмент.

Что почитать и где скачать:

• документация,
• исходный код на GitHub,
• опыт разработки тензорного компилятора под RISC-V CPU с помощью OpenVINO и MLIR в YADRO,
• интеграция как бэкенда с:
  • ONNX Runtime,
  • PyTorch,
• лицензия: Apache 2.0.

Тензорный компилятор, который специализируется на выполнении нейронных сетей на CPU, GPU и ускорителях машинного обучения. TVM активно развивается, а обнаруженные баги регулярно закрываются. У вас есть выбор: использовать последний из доступных релизов или брать основную ветку репозитория и собирать компилятор из исходников.

TVM можно рассматривать как самостоятельный тензорный компилятор, в котором есть несколько высокоуровневых API для инференса глубоких моделей. Он способен загружать, конвертировать и компилировать модели из всех наиболее популярных фреймворков: PyTorch, ONNX, TensorFlow Lite, Caffe, MXNet, TensorFlow (в ONNX через пакет tf2onnx) и Keras. Есть встроенная поддержка квантизации в формат весов INT8 — по крайней мере, если верить документации и результатам экспериментов для нескольких классов моделей. Также есть возможность запуска моделей, квантованных внешними инструментами.

TVM — активно развивающийся тензорный компилятор, вокруг которого сложилось большое сообщество разработчиков и пользователей. В репозитории на GitHub более 3400 форков и 900 контрибьюторов — значительно больше, чем у остальных рассмотренных инструментов, что говорит о его перспективности.

Что почитать и где скачать:

• документация,
• исходный код на GitHub,
• что «под капотом»,
• лицензия: Apache 2.0.

Резюмируем особенности описанных тензорных компиляторов:

OpenVINO и Apache TVM — самые продвинутые и функциональные фронтенд-ориентированные инструменты. Оба проекта активно развиваются — разработчики добавляют поддержку аппаратных платформ, а также улучшают функционал.