MediaPipe v0.8.9 の custom operator を TensorFlow に追加して wheel をビルドする

目的

MediaPipe v0.8.9 の custom operator を TensorFlow に追加してビルドし、Pythonのwheelパッケージを生成する。

理由は、MediaPipe Face Mesh の face_landmark_with_attention のtfliteモデルをTensorFlowで読み込んで、PINTO様の神ツール tflite2tensorflow を使ってTensorFlowの標準演算に置き換えた上で、最終的にはONNX形式等で書き出したいから、等です。

検証環境

• Ubuntu 20.04.3 x86_64
• Python 3.8.10
• TensorFlow v2.8.0
• MediaPipe 0.8.9
• Bazel 5.0.0 (MediaPipe用), 4.2.1 (TensorFlow用)
• Docker 20.10.14

方針

ベースにした方法

基本的な方法は、PINTO様の下記のドキュメントの通りです。
本ドキュメントは、下記を執筆時点で最新のバージョンで実行するための事項を整理したものです。

• Add a custom OP to the TFLite runtime to build the whl installer

大まかな流れ

• MediaPipeのパッチを適用したTensorFlowのソースを取得
  MediaPipeのソースをビルドすることにより、MediaPipe専用のパッチを適用済みのTensorFlowのソースを生成します。
• MediaPipeのcustom operatorのコードを修正
  TensorFlowで使えるように、MediaPipeのcustom operatorのコードの名前空間を変更します。
• MediaPipeのcustom operatorのコードをTensorFlowに追加
  上記のコードをTensorFlowの所定のパスに配置し、TensorFlowのビルドに含まれるようビルド設定を修正します。
• TensorFlowをビルド
  TensorFlowのwheelパッケージを生成し、MediaPipeのtfliteが読み込めることを確認します。

手順

MediaPipeのパッチを適用したTensorFlowのソースを取得

MediaPipe公式ドキュメントにしたがって、下記のようにMediapPipeをインストールします。

このとき、裏ではBazelがTensorflowのリポジトリから特定のコミットを取得し、MediapPipe用のパッチを当てたり一部のコードを追加したりしてくれています。

cd ${HOME}
git clone https://github.com/google/mediapipe.git
cd mediapipe

sudo apt-get install mesa-common-dev libegl1-mesa-dev libgles2-mesa-dev
wget https://github.com/bazelbuild/bazel/releases/download/5.0.0/bazel-5.0.0-installer-linux-x86_64.sh
sudo chmod +x bazel-5.0.0-installer-linux-x86_64.sh
sudo ./bazel-5.0.0-installer-linux-x86_64.sh
sudo bazel clean --expunge

export GLOG_logtostderr=1
sudo bazel run --define MEDIAPIPE_DISABLE_GPU=1 mediapipe/examples/desktop/hello_world:hello_world

ビルド中に、ログに

DEBUG: /root/.cache/bazel/_bazel_root/f7852c81ed1d3784a691d7bc8b21f673/external/org_tensorflow/third_party/repo.bzl:124:14: 

などが出ていると思います。このorg_tensorflow以下のフォルダが欲しかったものです。
下記コマンド等で、このリポジトリを使いやすい場所に移動しておきます。

sudo cp -r /root/.cache/bazel/_bazel_root/f7852c81ed1d3784a691d7bc8b21f673/external/org_tensorflow ${HOME}/org_tensorflow

ここまでで、下記の状態のソースが取得できているはずです。

https://github.com/KenjiAsaba/tensorflow/tree/c285bdaa93f8aa9245cc66f7926ced4ab6a04ec9

MediaPipeのcustom operatorのコードを修正

${HOME}/mediapipe/mediapipe/util/tflite/operationsにある下記12ファイルに対して、下記修正内容の通り名前空間の変更を行います。

対象ファイル：

• landmarks_to_transform_matrix.cc
• landmarks_to_transform_matrix.h
• max_pool_argmax.cc
• max_pool_argmax.h
• max_unpooling.cc
• max_unpooling.h
• transform_landmarks.cc
• transform_landmarks.h
• transform_tensor_bilinear.cc
• transform_tensor_bilinear.h
• transpose_conv_bias.cc
• transpose_conv_bias.h

修正内容：

https://github.com/KenjiAsaba/mediapipe/commit/f96b3697f050b17c27973c09ba7d65618e4c592f

MediaPipeのcustom operatorのコードをTensorFlowに追加

${HOME}/mediapipe/mediapipe/util/tflite/operationsにある先述の12ファイルを、${HOME}/org_tensorflow/tensorflow/lite/kernelsにコピーします。

また、これらがTensorFlowのビルド時に組み込まれるように、${HOME}/org_tensorflow/tensorflow/lite/kernelsにある下記3ファイルについて、下記修正内容の通りの変更を行います。

対象ファイル：

• BUILD
• register.cc
• register_ref.cc

修正内容：

• (BUILDについては、上記12ファイルおよび依存先のソースが読み込まれるよう指定します)
• (register.cc、register_ref.ccについては、ノードの名前と実際の関数との関連付けを行います)

https://github.com/KenjiAsaba/tensorflow/commit/35845189a59e5e815516106abb76f41fc85e93fd

TensorFlowをビルド

TensorFlow公式ドキュメントのDockerを使ったビルド方法にしたがって、下記のようにビルドします。

まずホスト側で、下記コマンドの通り、Bazel 4.2.1のインストーラとTensorFlowのDockerイメージを取得し、コンテナを起動します。

cd ${HOME}/org_tensorflow
wget https://github.com/bazelbuild/bazel/releases/download/4.2.1/bazel-4.2.1-installer-linux-x86_64.sh
sudo chmod +x bazel-4.2.1-installer-linux-x86_64.sh

sudo docker pull tensorflow/tensorflow
sudo docker run -it -w /tensorflow -v ${HOME}/org_tensorflow:/tensorflow -v $PWD:/mnt \
    -e HOST_PERMS="$(id -u):$(id -g)" tensorflow/tensorflow bash

Dockerコンテナに入るので、下記コマンドの通り、必要なツールをインストールします。

apt install git -y
apt install unzip -y

./bazel-4.2.1-installer-linux-x86_64.sh
bazel clean --expunge

続いて、下記コマンドでBazelビルドのコンフィグを行います。Enter連打でOKです。

./configure

最後に、TensorFlowのビルドを実行します。半日くらいかかります。

bazel build --config=opt //tensorflow/tools/pip_package:build_pip_package
/bazel-bin/tensorflow/tools/pip_package/build_pip_package /mnt
chown $HOST_PERMS /mnt/tensorflow-2.8.0-cp38-cp38-linux_x86_64.whl

ホスト側の${HOME}に、tensorflow-2.8.0-cp38-cp38-linux_x86_64.whlが生成されているはずです。

動作確認

下記コマンドでインストールします。

pip3 install tensorflow-2.8.0-cp38-cp38-linux_x86_64.whl

MediaPipeのtfliteを読み込んで、煮るなり焼くなりしましょう。

python3

from tensorflow.lite.python.interpreter import Interpreter
interpreter = Interpreter(model_path="face_landmark_with_attention.tflite")
interpreter.allocate_tensors()

input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

import pprint
pprint.pprint(input_details)
pprint.pprint(output_details)