Amazon открыл код компилятора NNVM для систем машинного обучения
Компания Amazon представила компилятор NNVM, предназначенный для компиляции предоставляемых системами машинного обучения высокоуровневых графов вычислений в оптимизированный набор машинных кодов. Код проекта распространяется под лицензией Apache 2.0.
Система поддерживает компиляцию моделей в форматах OpenML (поддерживаются фреймворки Keras и Caffe), Apache MXNet и недавно представленного FaceBook и Microsoft открытого формата ONNX (Open Neural Network Exchange), при помощи которого могут передавать модели обучения из фреймворков Caffe2, PyTorch и CNTK (Cognitive Toolkit). На выходе может генерироваться код для различных бэкендов, включая ядра для вычислений на стороне GPU при помощи CUDA, OpenCL и Metal, а также промежуточный код LLVM, на основе которого могут формироваться машинные инструкции для архитектур x86 и ARM.
Компиляция включает в себя несколько стадий:
- Формирование промежуточного представления графа вычислений на основе данных, полученных от различных фреймворков машинного обучения.
- Оптимизация полученного графа вычислений и выделение операторов с подпрограммами оптимизации графа.
- Компиляция в исполняемые модули и развертывание для различных аппаратных бэкендов с минимальными зависимостями. Для построения кода на основе промежуточного представления используется стек TVM, предоставляющий операторы, реализованные и оптимизированные для целевых аппаратных платформ.
Архитектура NNVM позволяет легко добавлять новые фронтэнды, операторы и оптимизации графа, без необходимости изменения базовых интерфейсов. Полученные после компиляции модули могут быть сгенерированы в виде кода на c++, python, javascript, java, objective-c для запуска на серверах, мобильных устройствах, встраиваемых системах и в web-браузере.
Проведение тестирования генерации кода для CPU ARM и GPU NVIDIA показало, что формируемый компилятором итоговый код превосходит по производительности фреймворк MXNet. Размер формируемого модуля зависит в основном от размера runtime TVM, который при сборки для Raspberry Pi и мобильных устройств занимает около 300 Кб.
