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様の下記のドキュメントの通りです。
本ドキュメントは、下記を執筆時点で最新のバージョンで実行するための事項を整理したものです。
大まかな流れ
-
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
ここまでで、下記の状態のソースが取得できているはずです。
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
修正内容:
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については、ノードの名前と実際の関数との関連付けを行います)
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)
Discussion
記事ありがとうございますm(_ _)m
Tensorflowのソース生成ですが、どこにdiffが出るか分からないと辛いので、Gitリポジトリに直接パッチを当てるほうが楽だと思います(Git管理したい場合は特に)。
あと、一長一短ありますが、Bazelはバージョンが上がると割と頻繁に壊れるので、特にこだわりがなければ、Bazeliskを使ってバージョンを固定したほうが良いと思います。
情報ありがとうございます!
gitコマンドでパッチ適用できるんですね…!git使いこなせていない人間でして、今知りました!
Bazeliskも、自動で適切なバージョンを使ってくれるとか…、これも今知りました笑
全般的に雰囲気でツールを使っているので、こういった改善コメントをいただけると大変助かります!🙏