Новый делегат Tensorflow Lite XNNPACK обеспечивает лучшую в своем классе производительность на процессорах x86 и ARM.

TL; DR: новый делегат Tensorflow Lite XNNPACK обеспечивает лучшую в своем классе производительность на процессорах x86 и ARM - в некоторых случаях более чем в 10 раз быстрее, чем бэкэнд Tensorflow Lite по умолчанию.

Tensorflow Lite - один из моих любимых программных пакетов. Он обеспечивает простое и быстрое развертывание на различном оборудовании и теперь поставляется с широким кругом делегатов для ускорения вывода - GPU, Core ML и Hexagon, и это лишь некоторые из них. Однако одним из недостатков Tensorflow Lite является то, что он разработан с учетом мобильных приложений и поэтому не оптимизирован для процессоров Intel и AMD x86. Лучшая поддержка x86 находится в дорожной карте разработки Tensorflow Lite, но на данный момент Tensorflow Lite в основном полагается на преобразование инструкций ARM Neon в SSE через мост Neon_2_SSE.

Однако есть новый делегат Tensorflow Lite для вычислений с плавающей запятой на базе ЦП, XNNPACK, который поддерживает оптимизацию x86 AVX и AVX-512. В этом посте я расскажу вам об использовании XNNPACK и покажу несколько тестов.

Установка и использование XNNPACK

Инструкцию по использованию XNNPACK можно найти здесь. В частности, теперь есть флаг сборки, который по умолчанию включает делегат XNNPACK. Это удобно, поскольку до сих пор нельзя было загружать делегаты Tensorflow Lite в Python. Команда для сборки Tensorflow из исходников будет выглядеть так:

bazel build --define tflite_with_xnnpack=true \ 
  //tensorflow/tools/pip_package:build_pip_package

Инструмент тестирования Tensorflow Lite теперь также имеет флаг для включения делегата XNNPACK. Например, для профилирования на машине x86 сначала создайте инструмент профилирования:

bazel build -c opt --verbose_failures \    
  tensorflow/lite/tools/benchmark:benchmark_model

Затем запустите профилировщик с помощью следующей команды:

bazel-bin/tensorflow/lite/tools/benchmark/benchmark_model \
  --graph=<model path> --warmup_runs=50 --num_runs=1000 \
  --enable_op_profiling=true --use_xnnpack=true

Оптимизация для XNNPACK

Важно убедиться, что модели подходят для XNNPACK, поскольку он поддерживает только подмножество всех операторов Tensorflow Lite. Например, стандартные реализации Keras часто используют явные уровни заполнения и реализуют верхний уровень глобального пула с помощью оператора среднего. При использовании обычного бэкэнда TFLite время выполнения увеличивается только на несколько процентных пунктов, но эти операции не поддерживаются XNNPACK, что приводит к значительным накладным расходам - ​​30% в случае MobileNet V2 с 8 потоками (см. Ниже)!

Заполнение легко исправить, заменив явные слои заполнения на заполнение, встроенное в операции свертки:

# Before
x = layers.ZeroPadding2D(padding=((3, 3), (3, 3)), name='conv1_pad')(img_input)
x = layers.Conv2D(64, 7, strides=2, use_bias=use_bias, name='conv1_conv', padding='valid')(x)
# After
x = layers.Conv2D(64, 7, strides=2, use_bias=use_bias, name='conv1_conv', padding='same')(img_input)

Уровень глобального среднего пула можно заменить средним уровнем пула с большим ядром:

# Before
x = layers.GlobalAveragePooling2D()(x)
x = layers.Dense(classes, activation='softmax')(x)
# After
# Use XNNPACK compatible average pooling
x = layers.AveragePooling2D(pool_size=(7, 7))(x)
    
# Implement the top dense layer as a convolution, so we don't need to remove spatial dims
x = layers.Conv2D(classes, kernel_size=1)(x)
x = layers.Softmax()(x)

Учтите, что вам придется переучивать модель. Вы можете найти исправленные версии этих моделей в репо к статье здесь.

Контрольные точки

РУКА

Итак, тесты! Сначала я решил протестировать MobileNet V3 на своем Galaxy S8:

Я тестировал с использованием 1 потока более 1000 итераций с 50 итерациями разогрева.

Как видите, XNNPACK предлагает отличную производительность поверх стандартной серверной части ЦП Tensorflow Lite. Стоит отметить, что XNNPACK поддерживает инструкции ARM Float 16, включенные в новые процессоры ARMv8.2-A (например, A55), но, к сожалению, у меня их нет под рукой. Бэкэнд графического процессора по-прежнему работает быстрее, особенно для более крупных моделей. Однако для этого требуется OpenGL ES 3.1 или выше, который доступен только на ~ 2/3 всех устройств Android (см. Долю рынка здесь.

x86

Теперь о x86. Я решил сравнить с пакетом Intel OpenVino, используя MobileNet V2 и ResNet50. Для тестирования я использовал экземпляр Google Cloud N2 Cascade Lake с 8 виртуальными ЦП. С 1 потоком:

И 8 потоков:

Как видите, Tensorflow Lite, использующий делегат XNNPACK, работает превосходно, в некоторых случаях более чем в 10 раз быстрее, чем бэкэнд Tensorflow Lite по умолчанию. Производительность приближается к OpenVino для MobileNet V2, но отстает от ResNet 50. Я не считаю это большой проблемой, поскольку архитектуры на основе глубинной свертки, такие как MobileNet V2, гораздо лучше подходят для развертывания ЦП. XNNPACK также имеет лучшее масштабирование для нескольких ядер ЦП, чем стандартный бэкэнд. Обратите внимание, что тестовый инструмент TFLite соответствует режиму задержки OpenVINO, поэтому было бы интересно посмотреть, что XNNPACK может предоставить, если он настроен на пропускную способность.

Резюме

Tensorflow Lite теперь может предложить отличную производительность x86 с помощью нового делегата XNNPACK, в некоторых случаях превосходя пакет Intel OpenVino. Основным недостатком XNNPACK является то, что он предназначен только для вычислений с плавающей запятой. Квантование 8-битной модели может легко привести к ›двукратному увеличению производительности, с еще большим увеличением при развертывании на новых процессорах Intel Cascade Lake, которые поддерживают инструкции AVX-512 VNNI. Поддержка 8-битного квантования на x86 входит в план развития Tensorflow Lite, возможно, даже в следующем выпуске.

Точно так же для мобильных развертываний XNNPACK превосходит бэкэнд по умолчанию Tensorflow Lite. Хотя делегат GPU по-прежнему работает быстрее, чем XNNPACK, XNNPACK полезен на устройствах, которые не поддерживают вычисления на GPU.

Мне особенно нравится XNNPACK, потому что он упрощает процесс развертывания, позволяя оставаться в экосистеме Tensorflow. Теперь можно преобразовать модель один раз с помощью Tensorflow Lite и развернуть ее на нескольких платформах, уменьшив количество требуемых различных пакетов программного обеспечения. Также стоит отметить, что процессоры AMD становятся все менее редкостью, а OpenVino - это продукт Intel. Я попытался протестировать на экземпляре Google Cloud N2D EYPC, но, к сожалению, мне не удалось увеличить квоту.

Надеюсь, вы нашли эту статью полезной. Код для воспроизведения этих тестов находится здесь и здесь.