Развертывание моделей машинного обучения на устройствах IoT с помощью ONNX

Одна среда выполнения, чтобы управлять ими всеми

В этой статье мы узнаем, как развертывать модели машинного обучения (ML) с помощью ONNX, экосистемы, которая позволяет нам разделить среду обучения и вывода для моделей ML. Сначала мы увидим, почему и когда это полезно, затем обзор того, как работает ONNX, и, наконец, мы рассмотрим пример с использованием компьютерного зрения. Начиная с обученной модели в PyTorch, мы экспортируем ее в ONNX и запускаем в Google Colab обе версии модели, чтобы проверить результат и найти различия в производительности, что является ключевым фактором для небольших устройств IoT, таких как Raspberry Pi 4. » в котором мы будем запускать версию модели ONNX.

Ожидается, что у вас есть базовые знания Python в применении к машинному обучению, доступ к Google Colab или эквивалентному рабочему пространству и Raspberry Pi.

Когда нам нужно изменить рабочий процесс развертывания?

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

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

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

Вот где инструмент, предназначенный для умозаключений, пригодится. Это может помочь нам получить максимальную отдачу от наших моделей, оптимизировать процесс и уменьшить количество зависимостей, необходимых для их выполнения. Имея это в виду, Open Neural Network Exchange (ONNX) является одним из таких инструментов. Среди его особенностей мы получим:

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

Как работает экосистема ONNX?

Первой частью экосистемы ONNX является спецификация формата, которая включает в себя:

Определения вычислительных графов для представления моделей
Поддерживаемые типы данных
Встроенные операторы.

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

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

Пример моделей классификации изображений

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

Мы перенесем этот конвейер с Pytorch на ONNX, для этого нам придется изменить модель и, в большинстве случаев, этап предварительной обработки. Это связано с тем, что мы обычно выполняем часть манипуляций с изображениями в той же среде, на которой обучается модель, поэтому, чтобы удалить зависимость от среды, нам придется изменить все части, в которых она используется в конвейере. В этой статье мы сосредоточимся на ключевых сегментах кода, но можем посмотреть полные файлы для этого примера в этом репозитории.

Начнем с клонирования репо:

!git clone https://github.com/Luis-ramirez-r/cv-onnx.git

Теперь, когда у нас есть все файлы, необходимые для выполнения задачи, давайте откроем блокнот Main.ipynb и запустим первую строку. Это сделает несколько вещей: во-первых, загрузите нашу начальную модель в PyTorch, выполните конвейер вывода и, наконец, сохраните модель. Если все работает правильно, в результате вы получите породу собаки «Самоед» и путь модели в следующей строке:

>>> !python '1_pytorch_model_test.py'
Samoyed
The model has been saved at models/googlenet.pt

Теперь, когда мы знаем, что модель работает, мы хотим преобразовать ее из PyTorch в ONNX, для этого мы можем обратиться к документации PyTorch.

Оттуда мы видим, что есть функция для экспорта моделей в формат ONNX с использованием torch.onnx.export . Чтобы использовать его, нам нужно предоставить в качестве входных данных: объект модели, то есть модель, которую мы загрузили в скрипте 1, образец входных данных, для этого мы будем читать наше тестовое изображение Dog.jpeg и путь к выходному файлу. . Для остальных параметров мы можем оставить значение по умолчанию. Существует также функция под названием check_model, которую мы будем использовать для подтверждения того, что модель имеет действительную схему. Если в файле есть ошибка, функция выдаст ошибку.

Код для этого шага находится в файле: 2_export_model.py.

>>> !python '2_export_model.py'
The model has been saved at: models/googlenet.onnx

Теперь мы готовы снова запустить конвейер, но с ключевым изменением. На этот раз с использованием ONNX-Runtime вместо PyTorch. Чтобы загрузить модель ONNX, нам нужно создать сеанс вывода, для которого требуются два параметра: путь к модели и поставщик выполнения (EP). EP позволяет среде выполнения оптимизировать вывод модели на платформу. На данный момент мы будем использовать «CPUExecutionProvider».

Код для этого шага находится в файле 3_onnx_model_base_pipeline.py, и мы можем запустить его так же, как и с предыдущими файлами, и результат должен быть таким же.

>>> !python '3_onnx_model_og_pipeline.py'
Samoyed

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

Чтобы устранить зависимость от PyTorch, мы должны переписать функцию предварительной обработки, для этого мы будем использовать NumPy:

И теперь мы закончили, на этом этапе мы, наконец, выполняем нашу модель без зависимости от PyTorch.

Полный код этого этапа находится в файле 4_onnx_new_pipeline.py, и, как видите, импортировать PyTorch не требуется.

Запустим скрипт:

>>> !python '4_onnx_pipeline.py'
Samoyed

Как и ожидалось, мы получили одинаковый результат для всех трех случаев, PyTorch, ONNX + PyTorch и только ONNX, теперь давайте посмотрим, в чем отличия.

Как вывод ONNX сравнивается с Pytorch?

Для всех случаев мы получили один и тот же прогноз от модели: Самоед, какая порода соответствует входному изображению, но есть ли различия между двумя моделями?

С точки зрения производительности мы получаем более интересный результат, где время вывода значительно меньше, когда мы используем ONNX, 129 мс против 41 мс, когда вывод делается в ЦП. Чтобы воспроизвести вывод с помощью графического процессора, сначала нам нужно изменить среду выполнения Colab на графический процессор:

Для кода в PyTorch никаких изменений не требуется, а для кода ONNX нужно только добавить CUDA EP:

После выполнения кода для обеих платформ мы получаем следующие значения времени вывода: 10 мс для PyTorch и 6 мс для ONNX. Хотя в этом случае разница меньше, мы по-прежнему получаем почти вдвое большую скорость в ONNX по сравнению с PyTorch.

А как насчет Интернета вещей?

До сих пор мы использовали ONNX в Google Colab для подготовки и тестирования конвейера, и теперь мы готовы к последнему шагу. Каких отличий можно ожидать при развертывании на устройстве IoT?

Устройства IoT бывают разных видов, но единственное, что мы можем ожидать, — это меньший размер и меньшее потребление, что подразумевает меньшую вычислительную мощность и меньшую стоимость. Это достигается за счет перехода на архитектуру ARM, ту же, что используется в наших телефонах. Примером таких устройств является Raspberry Pi, именно его мы и собираемся использовать. Изменение архитектуры может быть проблемой во многих случаях, но, поскольку ONNX является кросс-платформенной средой выполнения, на этот раз нам не о чем беспокоиться.

Я использую версию оперативной памяти Raspberry Pi4 4 ГБ. Если у вас нет доступа к нему, альтернативой является использование виртуальной машины с архитектурой ARM, например экземпляры AWS Graviton.

Настройка Raspberry Pi 4

В отличие от установки Colab, где большинство библиотек были доступны с самого начала, в этом случае нам нужно будет установить все зависимости в наш Raspberry Pi. Надеюсь, поскольку это среда вывода, только их не так много.

Чтобы все было чисто, вы можете клонировать репозиторий на свое устройство и искать файл requirements.txt, в котором есть все библиотеки, необходимые для выполнения модели, как видите, их всего четыре.

Чтобы установить библиотеки на Raspberry Pi, просто запустите код:

pip install -r requirements.txt

И теперь мы готовы к работе, вы можете запустить тот же файл, который мы использовали в Colab 4_onnx_new_pipeline.py, и вы получите тот же результат, собаки размножаются.

>>> python '4_onnx_pipeline.py'
Samoyed

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

Что дальше?

Пока мы рассматриваем один вариант использования экосистемы ONNX, преобразование модели PyTorch в ONNX и ее развертывание на Raspberry Pi, но это только верхушка айсберга. Что произойдет, если мы начнем с другого фреймворка или захотим использовать другое целевое устройство?

Для PyTorch мы можем напрямую экспортировать модели в формате .onnx, но для многих фреймворков нам приходится использовать внешнюю библиотеку, такую как onnxmltools, которая среди прочих поддерживает scikit-learn, Keras, H2O.

Если вы планируете развернуть свои модели на другом устройстве, таком как Jetson nano или нейронный вычислительный накопитель, обязательно установите версию среды выполнения и EP, которые лучше оптимизированы для вашей платформы, здесь вы можете посмотреть инструкции по установке среды выполнения ONNX на всех поддерживаемых платформах, а здесь — список всех доступных EP.

Заключение

В этой статье мы описываем проблемы развертывания моделей машинного обучения, когда используется несколько платформ или когда у нас есть ограничения производительности, и представляем экосистему ONNX в качестве альтернативы для решения этих проблем. Затем мы применили практический подход к компьютерному зрению, в котором мы:

Использовали Google Colab в качестве платформы для разработки
Экспортировал модель из PyTorch в ONNX и соответствующим образом изменил этап предварительной обработки.
Подтвердите, что мы получили тот же результат
Сравните производительность между ONNX и Pytorch
Развернул модель на Raspberry Pi

И, наконец, просмотрите некоторые рекомендации для дальнейшей работы, например. для развертывания моделей на других устройствах или использования разных EP для оптимизации времени вывода.