「ONNX Runtime Web」でWebGPUによるブラウザ内機械学習の高速化が可能に：大規模な生成AIモデルの台頭に対応

[＠IT]

　Microsoftは2024年2月29日（米国時間）、オープンソースでクロスプラットフォームの機械学習（ML）モデルアクセラレータ「ONNX Runtime」の「ONNX Runtime Web」機能において、WebGPUを用いてWebブラウザ内でのMLを高速化できるようになったと発表した。

　ONNX Runtimeは、さまざまなハードウェア、ドライバ、OSと互換性があり、グラフの最適化や変換に加え、ハードウェアアクセラレータを適宜活用したパフォーマンスを提供する。PyTorch、TensorFlow/Keras、TensorFlow Lite、scikit-learnなどのフレームワークのモデルで使用できる。

　ONNX Runtime Webは、JavaScript APIとライブラリを使用して、WebアプリケーションでMLモデルを実行、デプロイ（展開）できる。

　ブラウザ内MLはオンデバイスAIの主要な技術として、大きな関心を集めている。プライバシー保護、ポータビリティ向上、クライアント／サーバ間の通信の排除によるコストとレイテンシの低減といった利点があるからだ。だが、大規模な生成AIモデルの台頭がブラウザ内MLに課題を投げかけている。こうしたモデルを実行するには、より多くの計算能力とメモリリソースが必要になるためだ。

WebAssembly（WASM）バックエンドとWebGPUバックエンドによる推論の高速化

　ONNX Runtime Webは、ブラウザ内でのMLモデルの推論をCPUとGPUの両方で高速化する。CPUによる推論についてはWebAssembly（WASM）バックエンドを、GPUによる推論についてはWebGLバックエンドを使用してきた。

　CPUによる推論では、ネイティブのONNX Runtime CPUエンジンをWASMバックエンドにコンパイルする。これにより、ONNX Runtime Webは一般的なMLモデルを効果的に実行でき、Transformers.jsなどのさまざまなWebアプリケーションで広く採用されている。Transformers.jsは、ブラウザで直接Transformersを実行するための人気のあるJavaScriptライブラリだ。

　一方、ONNX Runtime Webは、GPUによる推論では、WebGLバックエンドを使用してきたが、大規模で複雑な生成AIモデルの台頭によるブラウザ内MLの課題に対処するため、WebGPUバックエンドを導入した。

　WebGPUは、WebGL、WebGL 2の後継となるWebグラフィックスAPIだ。最新GPUのハードウェア機能を引き出し、高性能な計算処理を可能にする。コンピュートシェーダーなどの高度な機能により、複雑なMLワークロードをWebGLより効率的に処理できるという特徴を持つ。FP16（16bit浮動小数点数）をサポートすることで、GPUのメモリ使用量と帯域幅の要件を軽減するとともに、演算を高速化する。GPUパワーを並列計算タスクに活用することで、ブラウザ内でMLアプリケーションの推論をより効率的かつスケーラブルに実行できるようにする。

　WebGPUは、Mac、Windows、ChromeOS版のGoogle Chrome 113以降とMicrosoft Edge 113以降およびAndroid版のGoogle Chrome 121以降のブラウザでサポートされている。

WebGPUバックエンドの効果を実証するデモ

　Microsoftは、ONNX Runtime WebのパフォーマンスにおけるWebGPUバックエンドの効果を実証するため、以下の2つのデモを開発し、デモページを開設している。

画像分離のためのエンコーダー／デコーダーモデル「Segment Anything」をブラウザで実行

　Segment AnythingはMetaが発表したモデル。数回の簡単なクリック操作で、画像内の特定のオブジェクトを効率的に識別、分離できる。Microsoftは、NVIDIA GeForce RTX 3060とIntel Core i9を搭載したノートPCにおいて、ブラウザでONNX Runtime WebとWebGPUバックエンドを使用してSegment Anythingを実行したところ、ONNX Runtime WebとWASMバックエンドを使用した場合と比べて、エンコーダーが19倍、デコーダーが3.8倍高速化したと報告している。

Segment Anythingをブラウザで実行（提供：Microsoft）

画像生成AIモデル「Stable Diffusion Turbo」をブラウザで実行

　Microsoftによると、NVIDIA GeForce RTX 4090を搭載したシステムにおいて、ONNX Runtime WebとWebGPUバックエンドを使用してブラウザでStability AIの画像生成AIモデルの一つであるStable Diffusion Turboを実行すると、1秒以内に結果が出力されたという。

　またMicrosoftは、ONNX Runtime WebとWASMバックエンドが広く採用されてきたTransformers.jsを使用する際も、モデルの大部分がWebGPUバックエンドで高速化できるようになったことを明らかにしている。