Zenn
🛠

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

2022/03/30に公開2件
Linux
TensorFlow
Ubuntu
MediaPipe
tech

目的

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

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

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)

Discussion

ピン留めされたアイテム
homulerhomuler

記事ありがとうございますm(_ _)m
Tensorflowのソース生成ですが、どこにdiffが出るか分からないと辛いので、Gitリポジトリに直接パッチを当てるほうが楽だと思います（Git管理したい場合は特に）。

git clone https://github.com/google/mediapipe.git
git clone https://github.com/tensorflow/tensorflow.git
cd tensorflow
git checkout [commit_id] # mediapipe/WORKSPACEの `_TENSORFLOW_GIT_COMMIT` をチェック

# WORKSPACEに書かれている順にパッチを適用する（本当は `patch -p1` するのが正しい）
git apply ../mediapipe/third_party/org_tensorflow_compatibility_fixes.diff
git apply ../mediapipe/third_party/org_tensorflow_custom_ops.diff

あと、一長一短ありますが、Bazelはバージョンが上がると割と頻繁に壊れるので、特にこだわりがなければ、Bazeliskを使ってバージョンを固定したほうが良いと思います。

npm install -g @bazel/bazelisk
株式会社空き家総合研究所株式会社空き家総合研究所

情報ありがとうございます！
gitコマンドでパッチ適用できるんですね…！git使いこなせていない人間でして、今知りました！
Bazeliskも、自動で適切なバージョンを使ってくれるとか…、これも今知りました笑
全般的に雰囲気でツールを使っているので、こういった改善コメントをいただけると大変助かります！🙏