TFLiteをWasm化(WasmSIMD化)してGoogle Meet 仮想背景機能を動かす
以前、Google Meetの仮想背景のモデル(Segmentation Model)をTensorflowjsで動作させた記事を投稿しました。今回はさらなるパフォーマンス改善を目指しWasm化したTFLiteで動かしてみようと思います。
結果、かなりの好成績でした。
前回の振り返りとパフォーマンス改善方針
以前の記事では、仮想背景機能で一般によく使われるMediaPipeのBodyPixと比較して大体1.6倍くらいの高速化が期待できそうということがわかりました。しかし、本家のGoogle Meetでの驚くほどの速度（軽さ）の実現は難しそうでした。詳細は以前の記事でご確認いただきたいですが、どうやらGPUからCPUへのメモリコピーで時間がかかっているように見えました。Google AIのブログによると、本家のGoogle MeetではWasm化したTFLiteで動かしているようです。この方式ではGPUからCPUへのコピーのオーバヘッドが発生しないため高速なのではないかと考えています。ということで、今回我々もTFLiteをWasm化して動かしてみようと思います。
なお、Google Meetの仮想背景モデルは、ライセンスが変更となっており、今現在はAPACHE-2.0で入手、利用することはできなくなっています(経緯はこちら)。一般にライセンス変更はそれ以前の成果物に遡及して適用されないと考えられていますが(参考)、もし本ブログを参考に再現を行う場合は各自の責任においてモデルの入手、利用をするようにしてください。
WasmとEmscripten
Wasm(WebAssembly)はウェブブラウザ上で実行可能なバイナリフォーマットのコードです(参考)。基本的に直接記述はせずに、C言語やC++, Rust, Goといった高水準言語で記述されたプログラムをコンパイルして生成します。特にC言語とC++で記述されたプログラムはEmscriptenと言われる一連のツールチェインを用いてコンパイルされることが多いです。TFLiteはC++で記述されていますので、このEmscriptenを用いてWasm化することになります。
TFLiteのWasm化
TFLiteはBazelというビルドツールを用いてビルドするように構成されています。私も知りませんでしたが、EmscriptenがBazelをサポートするようになったのは結構最近で、昨年の9月のようです(参考)。いずれにせよEmscriptenでもBazelを使えるということなので、これを使ってビルドしていきたいと思います。また、ついでに前処理や後処理で使うOpenCVもWasm化します。
手順は次のとおりです。
- Emscriptenのインストール
- OpenCVのWasmビルド
- Tensorflow, MediaPipeのソース取得
- コーディング
- TFLiteのWasmビルド
それでは、作業内容を説明していきます。なお、今回使用するソースコード等は下記のリポジトリにおいてあります。詳細はそちらをご参照ください。以下の作業はリポジトリ内のDockerfileに記載されている処理となりますので、実際に再現する場合はDockerfileからコンテナをビルドすることをお勧めします。
Emscriptenのインストール
公式ドキュメントでは、Emscriptenはemsdkを用いてインストールすることが推奨されています。我々もこれに従ってインストールをします。なお、今回は最新版のv2.0.14を使います。
$ git clone https://github.com/emscripten-core/emsdk.git -b 2.0.14 --depth 1
$ cd emsdk
$ ./emsdk install latest
$ ./emsdk activate latest
OpenCVのWasmビルド
githubからOpenCVのソースコードを取得し、emsdkを用いてビルドします。バージョンは最新版のv4.5.1を使います。上記の作業によりemscriptenはemsdk/upstream/emscriptenにインストールされたようです。(1)でこのパスを指定してconfigを作成します。このコマンドによりbuild_wasmフォルダが作成されるので、そこに入ってmakeします。
$ git clone https://github.com/opencv/opencv.git -b 4.5.1 --depth 1
$ cd opencv
$ python3 platforms/js/build_js.py build_wasm --emscripten_dir=<emsdk_dir>/upstream/emscripten --config_only # (1)
$ export OPENCV_JS_WHITELIST=/opencv/platforms/js/opencv_js.config.py
$ cd build_wasm && <emsdk_dir>/upstream/emscripten/emmake make -j && <emsdk_dir>/upstream/emscripten/emmake make install
Tensorflow, MediaPipeのソース取得
TFLiteのソースコードはTensorflowのリポジトリに含まれています。Tensorflowのソースを入手しましょう。バージョンは最新版のv2.4.1とします。また、MediaPipeのソースコードも入手します。こちらも最新版の0.8.3.1を使います。TensorflowのBUILDファイルにtoolchainの設定をする必要があります。なお、今回使用するTensorflowが依存するモジュールの中には、今回使うEmscripteのバージョンを想定していない（コンフリクトする）モジュールがあります。このため、一部の依存モジュールを最新のものに置き換える必要があります。これらの変更はリポジトリ内のDockerfileを参照ください。
$ git -C /tensorflow_src checkout refs/tags/v2.4.1 -b v2.4.1
$ git clone https://github.com/google/mediapipe.git -b 0.8.3.1 --depth 1
コーディング
EmscriptenでBazelを用いたビルドを行う雛形が<emsdk_dir>/bazelに用意されています。このフォルダの中身をコピーして編集していきます。Bazelのビルド環境の作成方法は公式をご参照ください。依存モジュールを記載するWORKSPACEやコンパイルオプションを記載するBUILD、.bazelrcを作成しておく必要があります。ソースコードは基本的にC++でTFLiteを使用するときのように記載すれば良いです。Javascriptと情報をやり取りするために、メモリアドレスを取得するための関数を用意しておくことがポイントです。Javascriptから呼び出す必要のある関数はEMSCRIPTEN_KEEPALIVEを付与しておきます。
下記の例では(1)はモデルのバイナリを格納するバッファのアドレスを返します。(2)は入力画像を格納するバッファのアドレスを返します。
namespace{
<snip>
///// Buffer for model
char modelBuffer[1024 * 1024 * 1];
///// Buffer for image processing
unsigned char inputImageBuffer[3 * MAX_WIDTH * MAX_HEIGHT]; // Input image Buffer
<snip>
}
extern "C"
{
<snip>
EMSCRIPTEN_KEEPALIVE
char *getModelBufferMemoryOffset(){ // ---(1)
return modelBuffer;
}
EMSCRIPTEN_KEEPALIVE
unsigned char *getInputImageBufferOffset(){// ---(2)
return inputImageBuffer;
}
<snip>
}
ビルド
<emsdk_dir>/bazelをコピーしたフォルダで下記のコマンドを実行します。
$ bazel build --config=wasm -c opt :tflite
simdを有効にしたい場合は、次のようにオプションをつけてビルドします。
$ bazel build --config=wasm -c opt --copt='-msimd128' :tflite-simd
出力されたファイルはアーカイブされているのでお好きな場所に展開しましょう。
$ tar xvf /tflite_src/bazel-bin/tflite -C <output_dir>
以上でTFLiteのWasm化は完了です。
Javascriptからの呼び出し
モジュールのロード
index.html等で生成された.jsファイルをloadします。
<head>
...
<script src="tflite/tflite.js"></script>
<script src="tflite/tflite-simd.js"></script>
...
</head>
そして、BUILDファイルに記載したEXPORT_NAMEをcallするとTFLiteのwasmがロードされます。
(1) Buildファイルで-s EXPORT_NAME=createTFLiteModuleと設定した想定。
createTFLiteModule().then(tflite => { // ---(1)
...
})
SIMDを使う場合
SIMDはChromeのExperimentalな機能です。利用するためにはorigin-trailを行う必要があります。
ここでSIMDのtokenを生成して、index.htmlのmetaタグに記載してください。
<meta
http-equiv="origin-trial"
content="xxxxx"
/>
評価
精度の評価
Google Meetの仮想背景のモデルは基本的に前回の使用したものと同じです[1]。結果もほぼ同じになりますので、96x160のモデルの結果だけ代表例として示します。
処理時間の評価
それでは今回の本当目的の処理時間について評価をしていきたいと思います。画像のサイズや使用する計算機の条件は前回の記事と同じです。
また、今回の計測する処理の範囲は前回に倣い下記のようにします。ただ、実装の都合上、後処理の一部[2]を(b)の中に入れてしまっているので、(b)についてはTFLite版は少々不利となっています。また、同じく実装の都合上(a)単体の計測は行っていません[3]。
今回は下記のような結果になりました。
(1)-(4)は前回の結果のコピーです。(5)-(7)がTFLite Wasmを用いた処理時間です。カッコ内はSIMDを有効化したものです。
(b) の範囲では、(1)~(3)のTensorflowjsを使用した場合に比べ、大幅に処理時間を短縮できています。比較的軽量なモデルである(5)(6)では、 10ms以下で処理出てきます。これと比較して(7)はだいぶ処理時間が伸びていますが、それでも(3)と比較して65%程度の処理時間で処理できています。SIMDを有効化すると更に短縮されます。(5)(6)では5~6msで処理できています。驚異的です。(7)でも10ms以下で処理できています。１秒間で考えると100フレーム以上の処理を軽々できてしまっているということになります。
(c) の範囲では、(5)-(7)ともに(b)の処理に+6msした処理時間になっています。SIMDを有効にすると、(7)と比較して(5),(6)で処理時間の増加量が多くなっています。この理由は今の所わかっていません。(c)の後処理部分はHTMLCanvasエレメントのメソッドを使って作っているのですが、この内部でwaitがかかっているのかもしれません。マスクのみ行ってHTMLCanvasに表示する(c)-p.p.についても、同様に(7)と(5)(6)の処理時間の差が若干縮まっています。HTMLCanvasの処理を行うごとに一定の待ち合わせが行われているのかもしれません(完全に推測です)。
いずれにせよ、Tensorflowjsのバージョンと比較して大幅な高速化が実現できていることがわかりました。
(b)の部分をグラフで比較するとこのような感じになります。
処理時間の評価(評価機を変えてみる)
同じ評価をM1チップ搭載のMac Book Airでも試してみました。全体的に高速化されていますが、基本的な傾向は同じでした。
より大きなモデルではどうなるのか？
Google Meet 仮想背景のモデルを用いた実験では、上記の通りTensorflowjsを用いるよりも高速に処理ができることがわかりました。しかし、144x256のモデルの結果をみるとわかるとおり、モデルの規模が大きくなると処理時間の短縮幅が小さくなります。これは大きなモデルほど並列処理が得意なGPUを使うほうが有利であるという直感的に理解できる結果かと思います。ということで、実際により大きなモデルをTFLite wasmを用いて処理をするとどうなるのかを簡易的にではありますが確認してみたいと思います。
今回は、こちらの記事で用いたWhite-box-Cartoonizationのモデルで実験してみました。
計測する処理範囲は画像をTensorflowjs/TFLiteに渡すところから、処理結果を表示するところまでです。要は上記のGoogle Meet 仮想背景のモデルの(c)に相当する範囲です。画像の解像度は192x192で試しています。次のような結果になりました。
やはりTensorflowjs(WebGL)が一番早いです。TFLite Wasmに対しては圧倒的に早いです。しかし、意外なことにTFLite WasmSimdもかなり早く、MBPでは圧倒的というほどの差は出せていないです。GoogleのBlogによるとSIMDのMultithread化でさらなる高速化ができるという情報もありますので、もしかしたらそれなりの規模のモデルであってもWebGLを使わずに使用できる日が来るかもしれません。
Demo
こちらからデモを参照していただけます。
- Google Meet 仮想背景のモデル
- White-box-Cartoonizationのモデル
まとめ
今回は、TFLiteをWasm化(Wasm SIMD化)してGoogle Meet の仮想背景のモデルを動かしました。処理時間を計測したところ、軽量なモデルであるGoogle Meet 仮想背景のモデルであればTensorflowjsよりもかなり高速に動かすことができることがわかりました。また、より大きなモデルでも実験したところ、SIMDを有効化することでTensorflowjsには及ばないまでも、かなり高速な推論が可能であることがわかりました。SIMDのMultiThreadによる高速化の検討も進められているようなので、今後のさらなる高速化が期待できそうです。
リポジトリ
本ブログで使用したソースコードはこちらのリポジトリにおいてあります。tfl001_google-meet-segmentationとtfl002_white-box-cartoonizationのフォルダをご覧ください。
謝辞
人物の画像、人物の動画、背景画像はこちらのサイトの画像を使わせていただきました。
Discussion