TFLite C++ API for Android ビルド手順メモ

Android C / C++ Docker Edge Computing TensorFlow TensorFlow Lite
このエントリーをはてなブックマークに追加 follow us in feedly

R&D チームの奥村（@izariuo440）です。TensorFlow 2.3 以降、TFLite が XNNPACK で高速になりそうな機運を感じたので Android での性能を少し確認してみました。XNNPACK を有効化してスレッド数を適切に設定すれば、CPU でも従来の 1.2 〜 2.0 倍の FPS が得られそうな感触をつかめました。

背景と動機

2020/02/05、TensorFlow Lite に XNNPACK1 が取り込まれました。2020 年 3 月末の TensorFlow Dev Summit では、さまざまなプラットフォームでの高速推論の一環として「Better CPU performance」で XNNPACK が取り上げられており、TensorFlow 2.3 で XNNPACK 対応がリリースされるとの記載があり、CPU での推論処理が 20% 〜 50% 高速化されるようです。TensorFlow 2.3.0 rc0 のリリースノートでは Android/iOS に関する記載も出ています。

私は DNN の高速推論に興味があり、これまで NVIDIA GPU、Edge TPU、WebAssembly SIMD といった環境に向けて取り組んできました。Edge TPU と WebAssembly SIMD は TFLite C++ API によるものということもありますし、何より TensorFlow 2.3.0 のリリースを待たずに性能を試したいので Android 向けに TFLite C++ API を手っ取り早くビルドする手順をメモしておきます。

ビルド手順

Android で TFLite C++ API を利用するには libtensorflowlite.so をビルドする必要があります。TensorFlow のビルドには Google が開発している Bazel2 というソフトウェアを使うのが一般的で、他にも Android SDK/NDK などが必要になります・・・。というわけで、ビルド環境の構築が面倒なので Docker を使って手順を簡素化します。

内容としては Android quickstart | TensorFlow Lite と XNNPACK backend for TensorFlow Lite をごった煮して Docker を使った最短手順をまとめたものになります。

ビルド環境構築手順

ホスト上で以下を実行し tflite-builder という Docker イメージを構築することでビルド環境を構築します。

# 1. 作業ディレクトリを作る。
mkdir tflite-cpp-api-for-android-workspace
cd tflite-cpp-api-for-android-workspace

# 2. TensorFlow 2.3.0 RC2 のコードをチェックアウトする。
git clone --depth 1 https://github.com/tensorflow/tensorflow.git -b v2.3.0-rc2

# 3. CPU 向け開発用 Docker イメージを構築する。
# 後述する tflite-android.Dockerfile が tensorflow/tensorflow:devel を要求するので、
# 止むなくタグを明示する。
docker build -t tensorflow/tensorflow:devel \
    -f tensorflow/tensorflow/tools/dockerfiles/dockerfiles/devel-cpu.Dockerfile \
    tensorflow/tensorflow/tools/dockerfiles

# 4. Android 向け TFLite 用 Docker イメージを構築する。
docker build -t tflite-builder \
    -f tensorflow/tensorflow/tools/dockerfiles/tflite-android.Dockerfile \
    tensorflow/tensorflow/tools/dockerfiles

Docker イメージのサイズは以下のようになりました。

$ docker images --format "table {{.Repository}}\t{{.Tag}}\t{{.ID}}\t{{.Size}}"
REPOSITORY              TAG                 IMAGE ID            SIZE
tflite-builder          latest              6423ae1a710c        6.92GB
tensorflow/tensorflow   devel               c3e78f366354        1.23GB

ちなみに、利用している Dockerfile は以下のとおりです。

  1. tensorflow/tools/dockerfiles/dockerfiles/devel-cpu.Dockerfile
  2. tensorflow/tools/dockerfiles/tflite-android.Dockerfile

上記の Dockerfile によると以下のような環境を構築してくれるようです。Bazel 以外、全体的に古めですね。

  1. Ubuntu 18.04
  2. Bazel 3.1.0 (2020/04/21)
  3. Android SDK r25.2.5 (2017 年 1 月)
    1. Android SDK API Level 23
  4. Android Build Tools 28.0.0 (2018 年 1 月 〜 2018 年 7 月？)
  5. Android NDK r17c (2018 年 9 月)
    1. Android NDK API Level 18

ビルド手順

ホスト上で以下を実行し、コンテナを起動します。

docker run --rm -it -v $PWD:/tmp tflite-builder bash

コンテナ上で以下を実行し、libtensorflowlite.so と必要なヘッダファイル一式をホスト側にコピーします。

# 1. TensorFlow の作業ディレクトリに移動する。
cd /tmp/tensorflow

# 2. libtensorflowlite.so をビルドする。arm64-v8a 向けのみ。
bazel build -c opt --config=android_arm64 \
    --host_crosstool_top=@bazel_tools//tools/cpp:toolchain \
    --define tflite_with_xnnpack=true \
    //tensorflow/lite:libtensorflowlite.so

# 3. 必要なファイルをコピーする。
INC_DIR=/tmp/out/include
LIB_DIR=/tmp/out/jniLibs/arm64-v8a
mkdir -p "${INC_DIR}" "${LIB_DIR}"
# 3-1. tensorflow/lite のヘッダファイルをコピーする。
find tensorflow/lite -name '*.h' -exec cp --parents \{\} ${INC_DIR}/ \;
# 3-2. flatbuffers のヘッダファイルをコピーする。
pushd .
cd bazel-tensorflow/external/flatbuffers/include
find flatbuffers -name '*.h' -exec cp --parents \{\} ${INC_DIR}/ \;
popd
# 3-3. libtensorflowlite.so をコピーする。
cp bazel-out/arm64-v8a-opt/bin/tensorflow/lite/libtensorflowlite.so ${LIB_DIR}/

exit

成功すると、作業ディレクトリ配下に以下のようなディレクトリツリーができるはずです。

out/
├── include
│   ├── flatbuffers
│   │   └── ヘッダファイル一式
│   └── tensorflow
│       └── lite
│           └── ヘッダファイル一式
├── jniLibs
│   └── arm64-v8a
│       └── libtensorflowlite.so
└── tensorflow
    └── チェックアウトした TensorFlow レポジトリ一式

あとは、これらのヘッダと共有ライブラリを利用したコードを実装してやれば OK です。

性能確認

上記で得られたヘッダと共有ライブラリを使い、TensorFlow 1.x から TensorFlow 2.3.0 rc2 にしてみました。

  • 端末: Pixel 4
  • モデル: RetinaFace
    • Backbone: MobileNet-V1
    • 入力解像度: 320x240
  • スレッド
    • TFLite スレッド数: 1
    • XNNPACK スレッド数: 2

上記構成で推論したときの FPS は、30 FPS 前後から 60 FPS 前後に向上しました。素晴らしいですね。なお、Backbone を ResNet-50 にすると 2 FPS 前後、入力解像度を 640x480 にした場合、XNNPACK スレッド数を 4 にすると 30 FPS 前後でした。

おわりに

XNNPACK が Android （というか ARM アーキテクチャ）でどれくらいの効果が出るのかが気になったので確認してみました。

オプティムでは、こうした技術に興味がある・作ってみたい・既に作っている、というエンジニアを募集しています。興味のある方は、こちらをご覧ください。

www.optim.co.jp

  1. XNNPACK の開発は活発で、Chrome M84 で WebAssembly SIMD の Origin Trial がはじまったこともあり、WebAssembly SIMD をターゲットとしたコミットが活発なようです。↩

  2. Bazel - a fast, scalable, multi-language and extensible build system" - Bazel↩