Pure な WebAssembly で推論処理をしてみる
こんにちは、NTT の森田です。
今月 Microsoft 社から、WebAssembly (WASM) を OS 上のサンドボックスとして Kubernetes で利用できるようになる Krustlet がリリースされました。新しい標準バイナリフォーマットとして、WASM の認知度がさらに上がったように思います。
ところで WASM で深層学習の推論処理を行おうとすると、どうすれば良いのでしょう？ TensorFlow.js や ONNX.js などは WASM バックエンドを持っていますが、JavaScript から実行することしかできません。今回、 JavaScript レイヤを必要としない、pure な WASM から推論処理を実行する方法を紹介したいと思います。
学習済みモデルから WASM ライブラリへコンパイル
LLVM は 8.0 から WASM ターゲットに対応しているので、LLVM をバックエンドに使っている深層学習コンパイラであれば、原理的には学習済みモデルから WASM バイナリを生成することが可能です。今回は TVM を使った例を見てみましょう。全体の処理フローは以下の図のようになります。
まずは学習済みモデルから WASM のライブラリを生成します。コンパイルを行うスクリプトはこちらのような感じになります。
LLVM のターゲットとして wasm32-unknown-unknown を指定してコンパイルすると、生成されるオブジェクトの形式はWASM フォーマットになるので、あとは普通に ar でアーカイブしてあげれば、clang や rustc でリンク可能な WASM の静的ライブラリの完成です。
WASM 推論プログラムの生成
TVM で生成した WASM ライブラリを実行するためには TVM のランタイムが必要です。C/C++ 言語だと標準ライブラリの関係で pure WASM を生成するのはハードルが高いのですが、TVM は Rust のランタイム実装 (tvm_runtime crate) を持っているので、WASM の推論プログラムが容易に生成できます。ここでは Krustlet で対応している、WASI と waSCC の WASM プログラムを、コンパイルした WASM ライブラリを使って作成してみましょう。
WASI
WASM は WASI というインターフェイスでファイルシステムへのアクセスができるようになります。このインターフェイスを使って、コマンドライン引数で指定した画像ファイルを読み込み、推論結果をコンソールへ出力する WASM プログラムを作成してみます。
サンプルプロジェクトはこちらになります。ビルドのターゲットは wasm32-wasi です。学習済みパラメータやグラフ構造は全てバイナリに埋め込んでいますが、WASI を使えるのでファイルから読み込んでも構いません。
実行結果はこちら。WASM のランタイムには wasmtime を使いました。ちゃんと推論ができていますね。
$ wget -O cat.png https://github.com/dmlc/mxnet.js/blob/master/data/cat.png?raw=true
$ wasmtime target/wasm32-wasi/release/tvm-wasi.wasm --dir . cat.png
original image dimensions: (256, 256)
resized image dimensions: (224, 224)
input image belongs to the class `tabby, tabby cat`waSCC
waSCC はアクターモデルの WASM ランタイムで、作成した WASM アクターを読み込んで、例えば HTTP 経由で WASM の実行をできたりします。ここでは推論処理を行う WASM アクターを作成してみます。
サンプルプロジェクトはこちらです。これをビルドすると、HTTP リクエストで画像ファイルを受け取り、推論処理結果を HTTP レスポンスの BODY に入れて返す WASM アクターができます。あとは waSCC のチュートリアルに従って、Hello world アクターの代わりに今回ビルドしたアクターを使えば、簡単な推論処理 HTTP サーバの完成です。
curl で画像ファイルを送りつけると、推論結果が返ってきます。
$ curl -s -T cat.png localhost:8081
"tabby, tabby cat"性能
私が WASM を試してみたいと思った背景でもあるのですが、TVM の開発で「WebGL は捨てて WebGPU に移行しよう。WebGL 使うぐらいなら最適化した WASM の方が早いよ。」といった話がありました。
実際のところどうなのでしょう？先ほどの wasmtime の実行時間を測定してみます。
$ time wasmtime target/wasm32-wasi/release/tvm-wasi.wasm --dir . cat.png > /dev/nullreal 0m3.235s
user 0m2.839s
sys 0m0.397s
モデルは ResNet50、CPU は Xeon CPU (E5–2660 v4) です。あまり速くないですね。原因は色々ありますが、原因の一つに WASM がマルチスレッドの native サポートをできていないというのがあります。マルチスレッドはまだ実装は提案中の段階です。WASM の性能に関しては、今後まだまだ改善の余地がありそうです。
おわりに
TVM は WASM 生成を想定した Rust 実装のランタイムを持っており、WASM との親和性が高い深層学習コンパイラです。さらに TVM は自動チューニング機構を持っているので、将来的には WASM バイナリの自動チューニングまで出来るようになると考えています。
また、最近 TVM コミュニティでは MLIR frontend の提案が出ました。これが実現すると、もはや TVM は Tensorflow のバックエンドコンパイラの一つとして位置づけられるとも言えます。TVM のエコシステムが拡大していって、楽しみです。
私たちNTTは、オープンソースコミュニティで共に活動する仲間を募集しています。ぜひ弊社 ソフトウェアイノベーションセンタ紹介ページ及び、採用情報ページをご覧ください。
