Привет, я Денис Соколов, руковожу R&D в Zenia — это платформа для йоги и фитнеса, которая использует ИИ для трекинга поз человека (подробнее об этом — в другой моей статье). Наша система распознавания работает на трёх платформах — iOS, Android, Web. В этой статье поговорим о ключевых отличиях между ними. Расскажу, как устроена подготовка моделей компьютерного зрения к использованию, какими фреймворками пользуемся для запуска на устройствах клиентов, какие сложности решали и чем остались довольны. Если вы занимаетесь запуском нейронных сетей на мобильных устройствах или вебе, статья для вас.

Обучение модели

Цикл жизни модели состоит из двух этапов — обучения и ее исполнения на устройстве. Для Zenia мы обучаем модели на серверах, а запускаем на устройствах пользователей. Это очень разные условия как с точки зрения окружения — операционных систем и железа, так и с точки зрения требований: необходимой функциональности и размера библиотек. Различия настолько большие, что фреймворки разделяются по назначению даже в пределах одной экосистемы — PyTorch и PyTorch Mobile, TensorFlow и TensorFlow Lite. Существуют отдельные решения, максимально эффективно использующие особенности платформ: CoreML и SNPE для мобильных устройств, TensorRT и OpenVino для серверов и десктопов. 

Для обучения моделей мы используем фреймворк PyTorch. В основе PyTorch лежит динамический граф вычислений. Это значит, что работу модели мы описываем императивно: явно задаем последовательность операций и даже можем менять операции «на лету». У этого подхода есть преимущества — простота и скорость разработки: последовательность операций задаётся явно, можно легко и быстро получить доступ к промежуточным выходам модели.

Но есть и недостатки такой динамичности: результатом обучения, который можно «опубликовать» для исполнения, является комбинация из кода, описывающего последовательность операций, и обученных весов. Он сильно привязан к конкретной среде исполнения — интерпретатору Python и полноценной сборке PyTorch. Такой формат плохо переносим (для исполнения требуется портирование и сборка большого объёма кода) и его сложно совмещать с другими компонентами системы (по сравнению, например, с библиотеками на C).

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

Далее коротко обсудим, как перевести модель из одного формата в другой, зачем это нужно и как устроен этот процесс.

Конвертация модели

Чтобы выжать максимальную пользу от оптимизаторов, динамичность из модели нужно убрать, то есть скомпилировать её. В PyTorch эту роль выполняет модуль torch.jit. Он отдает низкоуровневое представление модели в виде статической последовательности операций, которая лучше подходит для оптимизации (TorchScript).

К сожалению, этот формат на сегодняшний день подходит только для использования внутри экосистемы PyTorch - PyTorch Mobile, который пока находится в бета версии, и  PyTorch C++ API (а также CoreML - об этом далее). Поэтому продолжаем конвертировать.

Универсальным решением конвертационных вопросов должен был стать ONNX — промежуточный формат, работу над которым начали Microsoft и Facebook в 2017 году. К сожалению, даже в 2021 году далеко не все фреймворки умеют работать с ним напрямую. А это значит, что приходится переживать конвертацию дважды. 

Любая конвертация модели потенциально несёт проблемы. Не поддерживается какой-нибудь слой, не совмещаются версии исходного и целевого фреймворка с конкретной версией ONNX. Даже если конвертация технически удалась, это ещё не гарантирует успех: из одного слоя может получиться пять. Или реализация какой-нибудь операции будет не такой, как вы хотели. И даже если подобрать правильную комбинацию параметров и версий, нет гарантии, что при обновлениях фреймворков что-нибудь не сломается.

А если вы используете кастомные слои, то возможностей для приключений становится ещё больше.

Что можно сделать, чтобы облегчить себе жизнь:

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

2. Тестировать модель до и после конвертации, в том числе на скорость работы

3. С помощью netron отсматривать получившийся после конвертации граф на предмет аномалий

4. Использовать только базовые слои — никаких channel shuffle, если вы можете себе это позволить

5. Приготовиться страдать :)

Альтернатива — использовать TensorFlow и страдать во время разработки. С точки зрения удобства исполнения на конечных устройствах TensorFlow, безусловно, выигрывает: большинство inference-фреймворков понимают его формат.

Отличия между платформами, на которых исполняется модель

iOS

Apple сделал свою платформу очень дружелюбной к машинному обучению. Есть 2 стандартных API — высокоуровневый Core ML и низкоуровневый MPS (Metal Performance Shaders). Они доступны на всех актуальных версиях iOS, поэтому о поддержке можно не беспокоиться. 

Core ML крайне удобен и позволяет, не сильно заморачиваясь, получить отличную производительность, в том числе за счёт автоматического использования Neural Engine - сопроцессора, ускоряющего операции, типичные для нейронных сетей. Если вы планируете работать с нейронными сетями на iOS, крайне рекомендую прочитать книгу Core ML Survival Guide. В ней много практических советов о том, как подружиться с Core ML и организовать пайплайн данных. К сожалению, автор прекратил её поддержку, но многие приёмы всё ещё актуальны.

Core ML очень удобен, когда нужно сделать так, чтобы модель «просто работала». Но если у вас не одна модель, а несколько в сложном пайплайне, или вы активно используете GPU для рендеринга и хотели бы избежать лишнего перекидывания буферов с CPU на GPU, вам понадобится работать с MPS - низкоуровневым фреймворком для работы с шейдерами. 

Если вас заинтересовала эта тема, рекомендую посмотреть прекрасную серию докладов Андрея Володина из Prisma AI: Machine Learning + Mobile: настоящее и будущее, Байтик к байтику, или Как выжать из телефона всё и не расплавить его Introducing Smelter: движок для инференса  

Android

Запуск нейронных сетей на Android — это дикий запад железа, прошивок, драйверов и китайских устройств за $50. Думаю, это одна из главных причин, по которой запуск нейронных сетей на GPU Android-устройств долгое время был крайне ненадежным вариантом ускорить производительность.

Есть много фреймворков для запуска нейронных сетей на Android-устройствах. От производителей железа и телефонов: MACE (Xiaomi), HiAI (Huawei), SNPE (Qualcomm), ncnn (Tencent), MNN (Alibaba), а также версии популярных фреймворков, адаптированные для мобильных устройств — PyTorch Mobile, TensorFlow Lite.

Google занимается стандартизацией API для нейронных сетей — NNAPI. Он призван облегчить жизнь разработчикам фреймворков и библиотек в поддержке GPU и NPU (сопроцессоров для ускорения работы нейронных сетей) и должен работать, начиная с Android 8.1. К сожалению, эта прекрасная идея пока не очень хорошо работает на практике и фактически поддерживается не на всех устройствах. Попытка запустить модель с использованием NNAPI запросто может привести к тому, что на самом деле работать она будет на CPU.

Библиотека для запуска модели — это лишь часть пайплайна, который проходит изображение от камеры до получения результата, видимого пользователю. Если вы хотите оптимизировать процесс, нужно глубоко разбираться с Camera API (которых на Android уже 3) и OpenGL. Или воспользоваться готовым решением — MediaPipe. 

Этот фреймворк предоставляет как готовые решения (распознавание ключевых точек на лице, трекинг рук и позы человека), так и примитивы для построения собственного графа вычислений. У него достаточно высокий порог входа, но решиться на его преодоление может помочь тот факт, что этот фреймворк сейчас является стандартом для построения подобных пайплайнов в Google. Если вы планируете его использовать, обязательно прочитайте whitepaper и на самом базовом уровне разберитесь в работе OpenGL - рекомендую курс The Cherno.

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

Вместе со сложностями поддержка Android несёт и важное бизнес-преимущество. Android — это не только мобильные телефоны, но и телевизоры, и другие умные устройства. Поддержка Android позволила нам быстро портировать приложение на SberPortal и SberBoxTop.

Здесь можно посмотреть крутой обзор мобильных фреймворков с уклоном в Android от автора AI Benchmark - бенчмарка железа мобильных устройств.

Web

В последнее время всё переезжает в браузеры или браузерные технологии. Машинное обучение — не исключение. Ситуация с GPU и браузерами сложная: браузеры сейчас — это почти операционная система, которая добавляет ещё один уровень абстракции над железом. 

С одной стороны, это круто: код для WebGL может работать как на встроенных, так и на дискретных GPU. С другой стороны, он будет работать не так эффективно из-за дополнительного посредника. Стандарты WebAssembly и WebGPU, которые сейчас разрабатывают, должны побороть этот недостаток и дать более низкоуровневый доступ к возможностям железа. Есть и «стандартный» API для нейронных сетей в браузере — WebNN.

Фреймворк для браузера выбрать проще, чем для Android. Есть явный лидер - tfjs, есть альтернативные варианты, например, OpenCV.js и ONNX Runtime Web. 

Использование нейронных сетей в браузере осложняется тем, что пока что нет стандартных способов построения эффективного пайплайна. Приходится изобретать свои решения — например, с использованием web workers, чтобы снять нагрузку с основного потока. Браузер часто ведёт себя непредсказуемо, перераспределяя лимиты по использованию железа на своё усмотрение. Это сильно осложняет профилирование и поиск проблем с производительностью.

Для поддержки максимального количество устройств с наилучшим пользовательским опытом применяют подстройку системы распознавания под конкретное устройство. С помощью бенчмарка определяется «мощность» платформы. В зависимости от этого регулируются целевой fps и разрешение входного изображения или выбирается более лёгкая версия модели.

Обвязка модели

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

Можно написать её несколько раз на каждой платформе, но тогда с каждой новой платформой скорость разработки замедляется и повышается вероятность ошибок. Также появляются сложности со стороны организации разработки: на каждой из платформ свой язык программирования, свои инструменты и подходы к разработке. Приходится выбирать: развивать экспертизу по каждой из платформ в отделе R&D или иметь один источник правды (эталонную систему) и передавать работу по интеграции в соответствующие команды. 

Мы пошли по первому пути. Во многом из-за того, что наша компания начиналась с R&D, и другого выбора у нас не было. До того, как в компании стали появляться отделы, отвечающие за конкретную платформу, наш R&D-отдел отвечал за полный цикл доставки, от набора данных и обучения моделей до интеграции системы в конкретную платформу.

Когда количество платформ растёт, а код пре- и пост-процессинга становится всё более сложным, хочется иметь возможность писать его только один раз. До недавнего времени единственным выбором для создания кросс-платформенных модулей был C/C++. Ситуацию поменял не так давно появившийся Kotlin Multiplatform.

Мы начали Android проект на Kotlin с расчётом на то, что Kotlin Multiplatform станет шагом в сторону переиспользования кода между платформами. Нам повезло, и мы попали в момент, когда KMM (Kotlin Multiplatform Mobile) уже был достаточно зрелым для реальных проектов. 

Впервые мы попробовали возможности Kotlin Native, когда писали собственные модули для MediaPipe. Несмотря на то, что интерфейс библиотеки, генерируемый Kotlin Native не совсем удобный, многословный и требует некоторых ритуалов для преобразования данных в нужный формат, цель была достигнута. Мы начали использовать единую кодовую базу. 

Когда мы решили портировать Zenia в Web, то уже неплохо знали особенности работы с KMM и решили использовать Kotlin/JS, чтобы переиспользовать бОльшую часть логики. И это сработало, причём без больших проблем.

Скорее всего, полученный результат не так эффективен и производителен, как если бы мы написали кросс-платформенную библиотеку на C++. Но работа с Kotlin многократно проще и приятнее. А главное — позволяет быстрее и с меньшими усилиями получить желаемый результат, что очень важно для стартапа.

Резюме и советы практикам

В этой статье мы совсем не касались способов оптимизации запуска моделей на серверах, десктопах и специализированных устройствах, таких, как NVIDIA Jetson и Intel Movidius.

Если планируете запускать модель на устройствах пользователей, с самого начала разработки стоит запланировать, какими фреймворками будете пользоваться, и накидать базовый пайплайн. Если для вашей задачи уже есть готовые решения, ознакомьтесь сначала с ними, позапускайте на ваших устройствах. Так вы лучше поймёте возможности платформ. Сейчас я бы выбрал CoreML на iOS, Camera X ImageAnalysis на Android c переходом на MediaPipe для более сложных пайплайнов, tf.js для браузеров. Но если вы попробуете PyTorch Mobile и расскажете об этом — буду благодарен :)

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

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

Один вовремя написанный юнит или интеграционный тест поможет вам сэкономить кучу времени на отладке сломанного пайплайна. Работа с мобилками не такая привычная, но вооружившись Streamlit, Flask/FastAPI и вебсокетами, можно за пару дней соорудить из телефона сервер и гонять кусочки пайплайна, перетаскивая картинки в браузере.