Как ускорить модель PyTorch на Android с помощью графического вывода и библиотеки MACE
В последние годы наблюдается тенденция к использованию логического вывода на мобильных телефонах. В Tensorflow поддержка GPU на мобильных устройствах встроена в стандартную библиотеку, но в случае с PyTorch она еще не реализована, поэтому нам нужно использовать сторонние библиотеки. В этой статье мы рассмотрим связку PyTorch➔ONNX➔libMACE.
Установка libMACE
Далее мы рассмотрим шаги по установке libMACE.
Давайте создадим отдельную виртуальную среду, которая будет содержать библиотеку и все, что ей нужно для работы.
Затем нам нужно установить систему сборки bazel, согласно официальному руководству по установке.
Для создания библиотеки MACE требуется Android NDK. Нам нужна версия r15c, которую мы загрузим в каталог ~/VENV/opt/android-ndk-r15c/.
Далее нам нужно установить дополнительные библиотеки.
На последнем этапе создадим скрипт для настройки переменных окружения android_env.sh (показан ниже).
Преобразование модели
MACE использует собственный формат представления нейронных сетей, поэтому нам необходимо преобразовать исходную модель. Процесс преобразования состоит из нескольких этапов. Мы рассмотрим это на примере ResNet 50 из библиотеки torchvision.
На первом этапе мы конвертируем модель PyTorch в формат ONNX.
После преобразования содержимое папки должно выглядеть так.
На втором этапе нам нужно сохранить модель в собственном формате libMACE. Создадим конфигурационный файл по гайду.
В файле должен быть указан абсолютный путь к файлу ONNX (model_file_path), контрольная сумма SHA256 (model_sha256_checksum), геометрия и имена входных и выходных тензоров (input_tensors , input_shapes, output_tensors и output_shapes), а также формат данных, в нашем случае NCHW .
Контрольную сумму можно рассчитать с помощью утилиты sha256sum.
После создания файла конфигурации мы можем запустить скрипт преобразования модели.
Если преобразование завершилось без ошибок, отобразится следующий текст.
Результатом преобразования будут файлы reset_model.data и reset_model.pb в ~/VENV/opt/mace/build/resnet_model/model/.
Настройка проекта Android Studio
Для нашего приложения в Android Studio нам нужно указать тип Собственное приложение C++.
Далее нам понадобится бинарная сборка библиотеки MACE из репозитория.
Давайте создадим каталог /app/libmace с папками arm64-v8 и armeabi-v7a, куда мы скопируем версии libmace.so для архитектуры cpu_gpu из libmace-v0.13.0/lib/.
В CMakeLists.txt нам нужно добавить MACE include dir.
Далее в файле CMakeLists.txt нам нужно создать библиотеку lib_mace и добавить ее в список target_link_libraries. Нам также нужно добавить -ljnigraphics в этот список для поддержки JNI Bitmap.
В файле app/build.gradle нам нужно добавить подразделы abiFilters и externalNativeBuild в defaultConfig.
В разделе android нам нужно добавить запись sourceSets (только для старой версии Android Studio).
В AndroidManifest.xml мы добавляем разрешение на чтение файловой системы.
Далее создадим папку assets (Создать новую➔Folder➔Assets Folder) и скопируем resnet_model.pb и resnet_model.data к нему.
Загрузка модели
Во-первых, мы добавим функцию загрузки модели в MainActivity.java.
Мы также добавим его реализацию в native_lib.cpp.
Для библиотеки MACE требуется специальная конфигурация запуска, которая показана ниже.
Давайте посмотрим на код более подробно. Во-первых, нам нужно выбрать устройство для вычислений (device_type). В случае с графическим процессором MACE подготавливает двоичные файлы OpenCL, поэтому ему требуется каталог storage_path, где библиотека может их сохранить. Далее указываем приоритет задачи(GPUperfHint и GPUriorityHint). Если мы выберем высокий приоритет, интерфейс может зависнуть.
После создания начальной конфигурации мы загружаем нейронную сеть в память и создаем MaceEngine.
Нейронная сеть возвращается как shared_ptr, которую мы не можем напрямую передать в MainActivity, поэтому мы введем промежуточный класс ModelData (показан ниже).
Результатом loadModel является указатель на объект этого типа.
Возвращает указатель как длинный.
Вывод модели
Для вывода модели мы объявляем функцию классификации в MainActivity.
Также мы добавляем его определение в native_lib.cpp.
Далее в этом разделе вы найдете пошаговое описание этой функции.
Во-первых, мы восстанавливаем указатель моделей из long.
Далее нам нужно подготовить данные в формате NCHW (пример кода приведен ниже).
После загрузки данных нам нужно сформировать параметры нейронной сети в виде словаря тензоров, как этот код. Формы тензоров должны соответствовать формам, указанным в файле resnet_model.yml.
Теперь данные готовы. Запустим вывод моделей.
Прогнозируемый класс будет номером вывода с максимальным значением.
Итак, основная часть приложения создана. Также нам нужна некоторая логика Java, связанная с взаимодействием с пользователем. В статье он не описан, но его можно найти в репозитории проекта. После его добавления мы можем собрать программу и запустить ее на мобильном телефоне.
Полный исходный код можно скачать здесь.
