Amazon представила компилятор NNVM для фреймворков машинного обучения
Amazon представила компилятор NNVM. Он разработан специально для оптимальной работы фреймворков машинного обучения на различных платформах.
На данный момент существует широкий выбор фреймворков для разработки алгоритмов машинного обучения. Также есть возможность запускать код на огромном количестве устройств: от мобильных телефонов до облачных дата-центров. Это разнообразие является проблемой для разработчиков ИИ.
Во-первых, необходимо поддерживать на должном уровне работу многих фреймворков. Во-вторых, нужно гарантировать высокую производителность на разном железе, используя портируемый код, который подойдёт как для запуска в веб-браузере, так и на видеокартах дата-центров. В-третьих, необходимо убедиться, что работа фреймворков будет обеспечена на ещё не вышедшем железе.
Компилятор всему голова
Amazon считает, что ключом к решению подобных проблем является используемый компилятор. Группа исследователей Вашингтонского университета и AWS представили свой подход в виде компилятора NNVM, основанный на применении стека TVM. Его целью является предоставление многократно используемой цепочки инструментов для компиляции описания высокоуровневых нейронных сетей из фреймворков глубинного обучения в низкоуровневые коды, используемые в бэкенде.
Компилятор NNVM
Целью компилятора является представление входных данных из разных фреймворков в качестве стандартизированных вычислительных графов с последующим их переводом в исполняемые графы.
Компилятор поддерживает модели в форматах OpenML (фреймворки Keras и Caffe), Apache MXNet и используемого Facebook и Microsoft открытого формата ONNX (Open Neural Network Exchange), при помощи которого передаются модели для обучения в фреймворках Caffe2, PyTorch и CNTK (Cognitive Toolkit). Результатом компиляции является код для различных бекэндов, включая вычислительные ядра CUDA, OpenCL и Metal. Также возможна генерация кода LLVM, на основе которого формируются машинные инструкции для архитектур x86 и ARM или представление WebAssembly.
Процесс компиляции кода
- Формирование промежуточного графа вычислений на основе входных данных, полученных из фронтенд-интерфейсов фреймворков;
- Оптимизация графа и выделение в нём операторов с подпрограммами обработки данных;
- Компиляция операторов в исполняемые модули и развёртывание для различных бэкендов с минимальными зависимостями.
Полученные после компиляции модули могут быть сгенерированы в необходимый код на разных языках программирования: С++, Python, JavaScript, Java, Objective-C — для дальнейшего использования на мобильных платформах, видеокартах, серверах и в веб-браузерах.
Производительность компилятора
Сотрудники Amazon сравнили в производительности фреймворк MXNet и новый компилятор NNVM на двух разных аппаратных кофигурациях: ARM-процессор на Raspberry Pi и видеопроцессор Nvidia в облачных сервисах Amazon.
В случае видеокарт Nvidia для фреймворка MXNet в качестве бекэнда использовалась библиотека cuDNN на графическом ускорителе Nvidia K80. Компилятор NNVM показал ускорение в 1,2 раза для наборов данных ResNet18 и MobileNet.
На Raspberry Pi демонстрируемые результаты ещё лучше. Скорость обработки ResNet18 у NNVM в 2,2 раза выше, а MobileNet — в 11,5 раз. Это связано главным образом с тем, что в MXNet глубина свёртки не оптимизирована (из-за отсутствия подобных операторов в библиотеке dnn), тогда как NNVM использует прямую генерацию эффективного кода.
Код компилятора и инструкция по его использованию доступны на GitHub.