カスタムオペレーションを追加したTensorFlow Liteのビルド

2019年10月19日 17:02

以下の記事が面白かったので、ざっくり訳してみました。

・Compiling a TensorFlow Lite Build with Custom Operations

1. はじめに

「TensorFlow Lite」は素晴らしいです。
一方、カスタムオペレーションを含むモデルを実行しようとすると、次のような例外が発生します。TensorFlow Lite」に存在しないオペレーション「Normalize」「ExtractFeatures」「Predict」を使用したためです。

W/System.err: Caused by: com.google.firebase.ml.common.FirebaseMLException: Internal error has occurred when executing Firebase ML tasks
   Caused by: java.lang.IllegalArgumentException: Cannot create interpreter: Didn't find custom op for name 'Normalize'
07-23 18:01:16.018 21726-21726/io.github.the_dagger.mlkit W/System.err: Didn't find custom op for name 'ExtractFeatures'
   Didn't find custom op for name 'Predict'
   Registration failed.

この問題を解決するには、これら「カスタムオペレーション」を含む独自の「カスタムTensorFlow Lite」をビルドし、Googleが提供する「TensorFlow Lite」の代わりに使用することです。

簡単に聞こえますが、関連するドキュメントはかなり曖昧でまとまっていません。「Android NDK」や「Bazel」について何も知らず、「C++」に何年も触れていない身にこれは大仕事で、この問題を解決するのにかなりの時間を要しました。

幸いなことに成功したので、どのように行ったかについての詳細な手順を共有することにしました。

2. セットアップ

はじめに、必要なツールとパッケージをマシンにセットアップします。

(1) Python 2.7/3
https://www.python.org/downloads/

(2) TensorFlowリポジトリ
TensorFlowリポジトリをGitHubからクローンします。

$ git clone https://github.com/tensorflow/tensorflow/

(3) Android Studio
Android Developerサイトにアクセスし、ダウンロードおよびインストールを行います。

(4) Android NDK (バージョン15cを推奨)
このプロジェクトをビルドするには、Android NDKも必要です。Bazelはバージョン16bまでサポートしていますが、ビルドに関する多くの問題がありました。次のリンクに移動して、バージョン15cをダウンロードしてください。
https://developer.android.com/ndk/downloads/older_releases

(5) Bazel
Bazelは、カスタムTensorFlowパッケージのビルドに使用するツールです。
次のリンクに移動して、システムにインストールしてください。
https://docs.bazel.build/versions/master/install.html

3. コンフィギュレーション

クローンしたTensorFlowリポジトリのルートに移動し、次のコマンドを入力します。

./configure

これにより、対話型のコンフィギュレーションウィザードが起動し、ダウンロードしたツールとパッケージの場所に関する質問が表示されます。
提案されたオプションのほとんどは必要ないため、ビルドの時間を節約するために、次のように回答することをお勧めします。

(1) Pythonインストールの場所を指定するか、Bazelが自動的に検出できる場合はEnter/Returnを押します。
(2) 後続の質問に対して次のように回答します。

1. GCPサポート: n
2. Hadoopサポート: n
3. AWSサポート: n
4. Apache Kafkaサポート: n
5. XLA JITサポート: n
6. GDRサポート: n
7. VERBSサポート: n
8. nGraphサポート: n
9. OpenCL SYCLサポート: n
10. CUDAサポート: n
11. clangの最新リリースのダウンロード : y
12 : MPIサポート: n

(3) 最適化フラグを指定するよう求められます。そのままにEnter/Returnを押します。
(4) Androidビルド用にワークスペースを構成するかどうかを尋ねられます。「y」を押します。
(5) Android NDKへのパスを求められます。確認後にEnter/Returnを押します。
(6) Android SDKへのパスを求められます。確認後にEnter/Returnを押します。
(7) Android Build Toolsのバージョンが求められます。確認後にEnter/Returnを押します。
(8) 同じディレクトリにある「.tf_configure.bazelrc」を開き、「build --action_env ANDROID_NDK_API_LEVEL = x」という行を見つけ、「x」に「21」を指定します。

4. カスタムオペレーションをTensorFlow Liteに追加

カスタムオペレーションのヘッダファイル「predict.h」「normalize.h」「extract_features.h」を、クローンしたTensorFlow Liteリポジトリの
「tensorflow/lite/java/src/main/native」にコピーします。

オペレーションのヘッダファイル(.h)がなく、ソース(.cc)のみがある場合は、.ccファイルの名前を.hに変更するだけで機能します。C++のビルドシステムは、これら2つの違いを認識していません。

次に、「builtin_ops_jni.cc」を開き、上記で追加したカスタムオペレーションのヘッダを追加します。

...
#include "tensorflow/lite/java/src/main/native/normalize.h"
#include "tensorflow/lite/java/src/main/native/predict.h"
#include "tensorflow/lite/java/src/main/native/extract_feature.h"
namespace tflite {
...

次に、テキストファイルと同じフォルダにある「BUILD」ファイルを編集します。「hdrs」という名前のブロックを見つけ、さきほどコピーしたヘッダファイルを追加します。

...
hdrs = [
        "exception_jni.h",
        "nativeinterpreterwrapper_jni.h",
        "tensor_jni.h",
        "tensorflow_lite_jni.h",
        "normalize.h",
        "predict.h",
        "extract_feature.h",
    ],
...
----

5. カスタムTensorFlowののビルド

TensorFlowリポジトリのルートに戻り、次のコマンドを入力してビルドします。

bazel build --cxxopt='--std=c++11' -c opt        \
    --fat_apk_cpu=x86,x86_64,arm64-v8a,armeabi-v7a   \
    //tensorflow/lite/java:tensorflow-lite

ビルドが完了するまでには、13インチのMacBook Proで約15〜20分かかります。成功すると、以下のパスに「カスタムTensorFlow Lite」(*.aar)が生成されています。

bazel-genfiles/tensorflow/contrib/lite/java/tensorflow-lite.aar

6. Androidプロジェクトに.aarファイルをインポート

.aarファイルをAndroidプロジェクトにインポートする手順は次の通りです。

(1) Android Studioでプロジェクトを開き、「File → New → New Module」を選択。
(2) 「Import .JAR/.AAR」を選択。
(3) 「tensorflow-lite.aar」を指定し、「Finish」を押す。

成功すると、プロジェクトに「tensorflow-lite」という名前のモジュールが追加されます。

このモジュールをアプリに含めるには、アプリの「build.gradle」を開き、「dependencies」に次の行を追加します。

implementation project(':tensorflow-lite')

過去にデフォルトの「TensorFlow Lite」を追加したことがある場合は、.aarファイルを追加する前に、削除するのを忘れないでください。

この記事が気に入ったらサポートをしてみませんか？