本記事はMLOps Advent Calendar 2020の 2 日目の記事です。
WebAssembly（Wasm）は機械学習モデルをデプロイする新たな手段になりうるでしょうか。 この記事では、機械学習モデル（特に Deep Learning）を Wasm でデプロイする周辺技術の動向や内部の仕組みをざっくりと説明します。
tkat0 です。WebAssembly(Wasm)面白いですね。 私も最近はyewでフロントエンドを書いて遊んでいます。Rust で React っぽいことできるのは新鮮で面白いです。
Wasm は、なんとなく速い JavaScript？とか機械学習で何に役立つの？とか思ってる方も多いと思います。 しかし、Wasm はブラウザでの推論時に使えるだけでなく、機械学習モデルのサービングやエッジデバイスでの推論にも使えると知ったら驚きませんか？
WebAssembly(Wasm) とは
MDN の説明を引用します
https://developer.mozilla.org/ja/docs/WebAssembly
WebAssembly はモダンなウェブブラウザーで実行できる新しいタイプのコードです。ネイティブに近いパフォーマンスで動作するコンパクトなバイナリー形式の低レベルなアセンブリ風言語です。さらに、 C/C++ や Rust のような言語のコンパイル対象となって、それらの言語をウェブ上で実行することができます。 WebAssembly は JavaScript と並行して動作するように設計されているため、両方を連携させることができます。
また最近では、ブラウザだけでなく、FaaS での利用やプラグインとしての利用なども広がっています。 これらの応用は、Wasm の特長である、ランタイム側からアプリケーションに制約をかけられるのでセキュアであることや、実行がネイティブ並に速いこと、そして OS やハードウェアに依存しないポータビリティを活用しています。
ブラウザ以外で実行する場合、ファイルやネットワーク、環境変数などの OS の機能を呼び出したいことが多いです。 Wasm 自体はポータブルに設計しつつもこれらを実現するために、Wasm からこれらの機能を呼び出す共通のインターフェースであるWASI(WebAssembly System Interface)が策定されています。
WASI についてはこの記事がわかりやすいです。図を見るだけでもわかります。
Wasm を機械学習で使うモチベーション
Wasm を機械学習分野で使うモチベーションは、以下が考えられます。これらは Wasm 自体のメリットの延長にあるものです。
- ブラウザでの推論を JavaScript ベースのライブラリより高速化するため
- マルチプラットフォームへのデプロイを容易にするため
後者について補足します。開発とプロダクション、あるいは学習と推論で Deep Learning を実行するプラットフォームが異なる場合、Deep Learning フレームワーク間の変換をするケースが多いです。その中間フォーマットや推論側のランタイムは複数あり、色々と大変です。
そしてそれらをうまく使い分けるためにアプリケーション側もうまく抽象化して設計することが求められるケースもあります。 この辺りは今年の RustFest で話をしました。
そのため、Wasm のポータブルであるという特長をうまく生かして、学習済みモデルやそれを含む機械学習のアプリケーションを OS やプログラミング言語を超えてどこにでもデプロイしやすくなると嬉しいよね、というのがモチベーションの一つです。
しかし、Wasm 自体は thread や SIMD などサポートされつつありますが、Deep Learning の高性能な計算に使うような GPU、TPU、DSP といったアクセラレーターを動かすことはできません。 なぜならハードウェアを抽象化してポータブルにしているためです。
ではどうやっているかと言うと、現在大きく 2 つのアプローチがあります。
- pure Wasm での推論
- 主にブラウザ向け、超軽量モデルなら十分ネイティブ同等の性能
- Wasm + ネイティブコードの呼び出しての推論
- 主にスタンドアロン向け、既存のミドルウェアなどを使いたいケース
ここでのスタンドアロンは、サーバーサイドで推論エンドポイントを立てたり、エッジデバイスに組み込んだりする形式を指します。
また、Wasm 自体の特長ではないですが、Wasm を使ってブラウザやデバイスにデプロイする、エッジ AI（クライアントサイド ML）としては以下のような利点があります。
- サーバーサイドのコストダウン
- 通信不要
- レイテンシ削減 → UX 向上を図れる
- 通信できない/したくないユースケースでの ML 利用ができる
- データアップロード不要でプライバシー保護ができる
Wasm の機械学習モデルデプロイの動向
ブラウザへのデプロイ
Google AI Blog: Background Features in Google Meet, Powered by Web ML
王道的な応用ですが、Google Meet の背景ぼかしは、ブラウザ上で Wasm としてデプロイしたモデルで行っています。
技術的にはMediaPipeを利用していて、TensorFlow Lite や C++で記述されたアルゴリズムを含む、ML パイプラインをまるごと Wasm に変換しブラウザで動かしています。 Wasm での SIMD 利用や、JavaScript と Wasm の行き来のオーバーヘッドを避ける設計とし、ネイティブレベルの速度を pure Wasm で実現しています。
サーバーサイドへのデプロイ
Wasm をデプロイできる FaaS の中でも、機械学習モデルをデプロイする事例がいくつかありそうです。
SecondState
High performance and safe AI as a Service in Node.js
SecondStateの事例では、Deep Learning の実行はネイティブコードを Wasm から呼び出しています。 これはパフォーマンス向上の目的です。 具体的には、TensorFlow を利用した CLI アプリを Wasm から呼んでいます。そんなのあり？
呼び出しているCLI アプリの実装を見ると、TensorFlow オフィシャルの Rust bindingを使っていますね。
Scailable
Object recognition on the edge. Using PyTorch, ONNX, WebAssembly, and… | by Maurits Kaptein | Medium
Scailable も任意の Deep Learning フレームワークのモデルを Wasm にコンパイルしてサービングできるサービスを作っているようです。
エッジデバイスへのデプロイ
ONNC-WASMやWASI-NNというプロジェクトがあります。
これらは、Deep Learning の実行部分は Wasm ではなく CPU、GPU、DSP、FPGA などで実行することを目指しているものです。技術的な詳細は後述します。
また調べきれていないのですが、MediaPipe の Wasm は、ブラウザだけでなくスタンドアロンでデバイス等で動くでしょうか？
各プロジェクトの Wasm 対応状況
Deep Learning フレームワークや、各プロジェクトでの Wasm 対応がどうなっているかを調べました。他にもあるかもしれないです。見つけたら教えて下さい。
TensorFlow
TensorFlow で Wasm をサポートするのは、TensorFlow LiteとTensorFlow.jsですが、いずれもXNNPACKをバックエンドとして利用しています。先程の MediaPipe も内部では TensorFlow Lite を呼び出しています。これらは、ライブラリ全体を Wasm にコンパイルしています。
XNNPACK は奥村さんの記事が非常にわかりやすいです。
物体検出器 EfficientDet をブラウザで高速推論 - OPTiM TECH BLOG
TensorFlow.js が Wasm 対応して高速化していく変遷は、公式のブログ参照。
- Introducing the WebAssembly backend for TensorFlow.js — The TensorFlow Blog
- Supercharging the TensorFlow.js WebAssembly backend with SIMD and multi-threading — The TensorFlow Blog
さらに、MediaPipe。これはリアルタイムな ML パイプラインを構築するアプリケーションフレームワークですが、この中では TensorFlow Lite のモデルや、動画の前処理、フィルタなどの後処理をすべてを Wasm にコンパイルしてブラウザで動かしています。MediaPipe 自体の価値は、メルカリ Edge AI チームからの記事が非常にわかりやすいです。YouTube も実は公開当日から見ていましたがよかったです。
Mediapipe を活用したストリーミング推論の事例紹介-カメラをかざして家の中から売れるアイテムを探そう | メルカリエンジニアリング
MediaPipe はすぐにブラウザで動かせるデモが複数あります。実際に動かしてみると、手軽過ぎて驚くし、この UX の良さは Wasm ならではだなと期待せざるを得ません。
ところで、MediaPipe もリビルド不要でオリジナルの Calculator を Wasm でプラグインできるようになると嬉しいな～。できる？
TVM
Deep Learning コンパイラ、推論エンジンとして有名なTVMにも既存のインフラの上でWasm サポートが導入されました。ブラウザでもスタンドアロンでも動くようになっています。
RPC サーバーもサポートしているので、Python から Wasm にコンパイルしたモデルをベンチマークしたり AutoTVM を用いてブラウザで動かしながらモデルを最適化できるのは既存のインフラを生かした強みですね。
https://tvm.apache.org/2020/05/14/compiling-machine-learning-to-webassembly-and-webgpu より引用
TVM は、Deep Learning の計算のカーネル自体も pure Wasm で記述されるか、カーネルだけ Wasm から WebGPU の実装を呼び出すか、の 2 通りが選択できます。
WebGPU を使う話はこちら。数ヶ月前の時点で、ブラウザでもネイティブ GPU の性能に近づいている！ということですね。
- Compiling Machine Learning to WASM and WebGPU with Apache TVM
- WebGPU powered machine learning in the browser with Apache TVM | by Jason Knight | OctoML | Medium
Pure Wasm はこちらが参考になります。LLVM でビルド。
WASI-NN
WASI-NNは Intel のメンバーが主に推進している、非常に抽象度の高いアプローチです。
Wasm には、「なんらかのフォーマットのモデルで推論する API」だけが含まれ、具体的な実装はすべて WASI-NN の実装であるネイティブ側に隠蔽されます。 特にハードウェアアクセラレーションをする場合、Wasm のレイヤーで諸々を吸収するのは大変そうなので、WASI-NN のようにかなり抽象化してほとんどをネイティブに投げてしまうのは現実的だと思いました。
API の定義をみると load, init_execution_context, set_input, compute, get_output などがあります。
現在ネイティブ側は、OpenVINOをつかった PoC を行っているようです。
このようにして、同じ WASI-NN のインターフェースでも、ネイティブ側で TensorFlow が動くものや PyTorch が動くものなどを作ることができます。プラットフォームに応じてモデルファイルや WASI-NN の実装を用意しておけば、アプリケーション側は共通の.wasm ファイルでデプロイできるのは良いのかも知れません。
ONNC-WASM
最後はONNC-WASMです。これもエッジで動かすことを目指しています。ONNCは ONNX からなんらかのプラットフォーム向けのコード生成をするコンパイラです。ONNC-WASM では、ONNX を C のソースコードに変換した上で、Wasm にコンパイルしています。
ONNC-WASM はISCA2020 のスライドが一番詳しいです。
https://github.com/ONNC/onnc-tutorial/blob/master/ISCA2020-slides/ISCA2020_ONNC_WASM_Project.pdf より引用”
これも WASI-NN と同様に、Deep Learning の計算部分はネイティブ側で実行する設計です。WASI-NN がすべてをネイティブ側に隠蔽しているのに対して、ONNC-WASM はグラフ自体を実行する部分は Wasm で実装され、各レイヤーのカーネル呼び出しを ONNC Wasm Runtime Interface でネイティブ側を呼んでいます。
実は今年の夏に ONNC-WASM のコードを動かしたことがあるのですが、色々辛かったので別の機会にお話できればと思います。
その他
- NimTorch: https://github.com/sinkingsugar/nimtorch
- PyTorch の Wasm 動向を調べていて見つけた。“WASM support” とある
- ONNX.js
- Wasm, WebGL を backend として使えます
結論
- Wasm を ML で使うモチベーション
- ブラウザでの推論を JavaScript ベースのライブラリより高速化するため
- マルチプラットフォームへのデプロイを容易にするため
- 今後アプリケーションを Wasm でデプロイすることが一般的になると、Wasm から ML を呼ぶのは必然
- Wasm の ML 動向
- ブラウザでの Wasm 推論は実用レベル
- エッジやサーバーサイドでの Wasm 推論は今まさに開発中
- 技術的には 2 アプローチ
- ブラウザに実装された Wasm の機能だけ（pure Wasm）で頑張る
- TensorFlow.js, TensorFlow Lite, MediaPipe など
- BlazeFace のような超軽量モデルは十分
- HW アクセラレーターやミドルウェアを使うため、Wasm からネイティブコードを呼び出す
- 各ライブラリで独自に設計（TVM, WASI-NN, ONNC-WASM）
- ネイティブ呼び出しの界面をどこにするかで設計の違いがあっておもしろいですね
- アプリケーション的には API は統一されていると嬉しいので、今後は WASI-NN に統一されていく？
- ブラウザに実装された Wasm の機能だけ（pure Wasm）で頑張る
ML 部分が Wasm になるから嬉しいというよりは、アプリケーション全体を Wasm にしてデプロイすることに利点が大きく、ML 部分もそのエコシステムに乗れるようにしていくモチベーションがあると思います。
実際、冒頭で話したように開発と本番で動く環境が異なることでシステムが複雑化したり開発工数が膨らんだりすることはあるので、少なくともアプリケーションレイヤーはなんらかマルチプラットフォームに実装できると嬉しいなあと思っています。
また、ワクワクする話としては、リソースの関係で Docker が利用できないようなエッジデバイスのプラットフォームにおいて、Wasm で機械学習のアプリケーションをデプロイできるようなマイクロサービスを設計するのは夢があります。やってみたい。
文章多めでしたが、最後まで読んでいただきありがとうございました。
おまけ
Wasm にコンパイルせず、Python スクリプトをそのまま Wasm の VM で動かしたいあなたへ。
pyodide: CPythonをwasmビルドし、そこでPythonスクリプトを動かしている。numpyやmatplotlibなどはパッチを当ててブラウザで動くようにしている。ブラウザ上でPyPIからもインストールできる...https://t.co/eDMroQppWq
— tkat0 (@_tkato_) November 20, 2020
Python を用いた Tutorial 環境や、Kubeflow Pipeline で実行した結果をその場でブラウザでインタラクティブに可視化する（サーバーを立てることなく）や、これはオフィシャルのモチベーションですが、Jupyter Notebook の分析結果を共有しやすくするといったところではうまく使えそうですね。