WebAssemblyでの機械学習モデルデプロイの動向

Dec 2, 2020 19:30 · 321 words · 2 minute read rust wasm deeplearning mlops

本記事は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についてはこの記事がわかりやすいです。図を見るだけでもわかります。

Standardizing WASI: A system interface to run WebAssembly outside the web - Mozilla Hacks - the Web developer blog

Wasmを機械学習で使うモチベーション

Wasmを機械学習分野で使うモチベーションは、以下が考えられます。これらはWasm自体のメリットの延長にあるものです。

  1. ブラウザでの推論をJavaScriptベースのライブラリより高速化するため
  2. マルチプラットフォームへのデプロイを容易にするため

後者について補足します。開発とプロダクション、あるいは学習と推論で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-WASMWASI-NNというプロジェクトがあります。

これらは、Deep Learningの実行部分はWasmではなくCPU、GPU、DSP、FPGAなどで実行することを目指しているものです。技術的な詳細は後述します。

また調べきれていないのですが、MediaPipeのWasmは、ブラウザだけでなくスタンドアロンでデバイス等で動くでしょうか?

各プロジェクトのWasm対応状況

Deep Learningフレームワークや、各プロジェクトでのWasm対応がどうなっているかを調べました。他にもあるかもしれないです。見つけたら教えて下さい。

TensorFlow

TensorFlowでWasmをサポートするのは、TensorFlow LiteTensorFlow.jsですが、いずれもXNNPACKをバックエンドとして利用しています。先程のMediaPipeも内部ではTensorFlow Liteを呼び出しています。これらは、ライブラリ全体をWasmにコンパイルしています。

XNNPACKは奥村さんの記事が非常にわかりやすいです。

物体検出器 EfficientDet をブラウザで高速推論 - OPTiM TECH BLOG

TensorFlow.jsがWasm対応して高速化していく変遷は、公式のブログ参照。

さらに、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 より引用

https://tvm.apache.org/2020/05/14/compiling-machine-learning-to-webassembly-and-webgpu より引用

TVMは、Deep Learningの計算のカーネル自体もpure Wasmで記述されるか、カーネルだけWasmからWebGPUの実装を呼び出すか、の2通りが選択できます。

WebGPUを使う話はこちら。数ヶ月前の時点で、ブラウザでもネイティブGPUの性能に近づいている!ということですね。

Pure Wasmはこちらが参考になります。LLVMでビルド。

Pure な WebAssembly で推論処理をしてみる. JavaScript レイヤを必要としない、pure な WASM… | by Kazutaka Morita | nttlabs | Medium

WASI-NN

WASI-NNはIntelのメンバーが主に推進している、非常に抽象度の高いアプローチです。

Wasmには、「なんらかのフォーマットのモデルで推論するAPI」だけが含まれ、具体的な実装はすべてWASI-NNの実装であるネイティブ側に隠蔽されます。 特にハードウェアアクセラレーションをする場合、Wasmのレイヤーで諸々を吸収するのは大変そうなので、WASI-NNのようにかなり抽象化してほとんどをネイティブに投げてしまうのは現実的だと思いました。

https://www.w3.org/2020/Talks/mlws/ms_ab_wasinn.pdf より引用

https://www.w3.org/2020/Talks/mlws/ms_ab_wasinn.pdf より引用

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 より引用

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のコードを動かしたことがあるのですが、色々辛かったので別の機会にお話できればと思います。

その他

結論

  • 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に統一されていく?

ML部分がWasmになるから嬉しいというよりは、アプリケーション全体をWasmにしてデプロイすることに利点が大きく、ML部分もそのエコシステムに乗れるようにしていくモチベーションがあると思います。

実際、冒頭で話したように開発と本番で動く環境が異なることでシステムが複雑化したり開発工数が膨らんだりすることはあるので、少なくともアプリケーションレイヤーはなんらかマルチプラットフォームに実装できると嬉しいなあと思っています。

また、ワクワクする話としては、リソースの関係でDockerが利用できないようなエッジデバイスのプラットフォームにおいて、Wasmで機械学習のアプリケーションをデプロイできるようなマイクロサービスを設計するのは夢があります。やってみたい。

文章多めでしたが、最後まで読んでいただきありがとうございました。


おまけ

Wasmにコンパイルせず、PythonスクリプトをそのままWasmのVMで動かしたいあなたへ。

Pythonを用いたTutorial環境や、Kubeflow Pipelineで実行した結果をその場でブラウザでインタラクティブに可視化する(サーバーを立てることなく)や、これはオフィシャルのモチベーションですが、Jupyter Notebookの分析結果を共有しやすくするといったところではうまく使えそうですね。

tweet