TensorFlow.js の WebAssemblyバックエンドの紹介

2020年6月12日 21:40

以下の記事を参考に書いてます。

・Introducing the WebAssembly backend for TensorFlow.js

1. はじめに

TensorFlow.js がブラウザと Node.js の両方に WebAssembly（WASM）バックエンド を提供するされたことをお知らせいたします。

このバックエンドは WebGLバックエンドの代替であり、最小限のコード変更で高速なCPU実行をもたらします。幅広いデバイスセット、特にWebGLのサポートがないか、GPUが遅いローエンドのモバイルデバイスでのパフォーマンス向上に役立ちます。XNNPackライブラリ を使用して操作を高速化します。

2. インストール

新しいWASMバックエンドを使用するには、2つの方法があります。

◎ NPM

// @tensorflow/tfjs または @tensorflow/tfjs-core のインポート
const tf = require('@tensorflow/tfjs');
// グローバルバックエンドレジストリへの WASM バックエンドの追加
require('@tensorflow/tfjs-backend-wasm');

// バックエンドをWASMに設定し、モジュールの準備できるまで待つ
tf.setBackend('wasm').then(() => main());

ライブラリは、WASMバイナリがメインのJSファイルに対して相対的であることを想定しています。parcel や webpack などのバンドラを使用している場合は、setWasmPathヘルパーを使用してWASMバイナリの場所を手動で指定する必要がある場合があります。

import {setWasmPath} from '@tensorflow/tfjs-backend-wasm';
setWasmPath(yourCustomPath);
tf.setBackend('wasm').then(() => {...});

詳細については、READMEの「Using bundlers」を参照してください。

◎ スクリプトタグ

<!-- @tensorflow/tfjs または @tensorflow/tfjs-core のインポート -->
<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs"></script>

<!-- グローバルバックエンドレジストリへの WASM バックエンドの追加 -->
<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs-backend-wasm/dist/tf-backend-wasm.js"></script>

<script>
 tf.setBackend('wasm').then(() => main());
</script>

【注意】 TensorFlow.jsは各バックエンドの優先度を定義し、特定の環境でサポートされている最適なバックエンドを自動的に選択します。現在はWebGLが最も優先度が高く、次にWASM、次にバニラJSバックエンドが続きます。常にWASMバックエンドを使用するには、明示的に tf.setBackend('wasm') を呼び出す必要があります。

3. デモ

WASMバックエンドで実行される顔検出デモ（MediaPipe BlazeFaceモデルを使用）を確認してください。モデルの詳細については、このブログ投稿を参照してください。

4. なぜWASMなのか？

WASMは、クロスブラウザ、ポータブルアセンブリ、およびWeb用のバイナリ形式であり、ネイティブコードに近い実行速度をWebにもたらします。2015年に新しいWebベースのバイナリフォーマットとして導入され、C、C ++、Rustで記述されたプログラムを提供します。これは、Web上で実行するためのコンパイルターゲットです。WASMは2017年からChrome、Safari、Firefox、Edgeでサポートされており、世界中のデバイスの90％でサポートされています。

5. WASMのパフォーマンス

◎ JavaScriptとの比較
WASMは、機械学習タスクで一般的な数値ワークロードのJavaScriptよりはるかに高速です。さらに、WASMは、JavaScriptの解析よりも最大20倍速くネイティブにデコードできます。JavaScriptは動的に型指定され、ガベージコレクションが行われるため、実行時に重大な非決定性の速度低下を引き起こす可能性があります。さらに、最新のJavaScriptライブラリ（TensorFlow.jsなど）は、TypeScriptやES6トランスパイラなどのコンパイルツールを使用して、バニラES6 JavaScriptよりも実行が遅い（ワイドブラウザサポート用の）ES5コードを生成します。

◎ WebGLとの比較
ほとんどのモデルでは、WebGLバックエンドはWASMバックエンドよりも優れていますが、ウルトラライトモデル（3MB未満および60Mの積和演算）ではWASMの方が高速です。このシナリオでは、GPU並列化の利点は、WebGLシェーダーの実行による固定オーバーヘッドコストよりも重要です。以下に、この行を見つけるためのガイドラインを示します。ただし、SIMD命令を追加して、複数の浮動小数点演算をベクトル化して並列に実行できるWASM拡張提案があります。予備テストでは、これらの拡張機能を有効にすると、今日のWASMを2〜3倍高速化できることが示されています。これがブラウザに適用されることに注意してください。TensorFlow.jsでは自動的にオンになります。

◎ 移植性と安定性
機械学習に関しては、数値精度が重要です。WASMは浮動小数点演算をネイティブにサポートしますが、WebGLバックエンドはOES_texture_float拡張を必要とします。すべてのデバイスがこの拡張機能をサポートしているわけではありません。つまり、一部のデバイス（WASMがサポートされている古いモバイルデバイスなど）では、GPUアクセラレーションTensorFlow.jsがサポートされていません。

さらに、GPUドライバーはハードウェア固有であり、デバイスごとに精度の問題が発生する可能性があります。iOSでは、32ビットフロートはGPUでサポートされていないため、16ビットフロートにフォールバックし、精度の問題が発生します。WASMでは、計算は常に32ビットの浮動小数点数で行われるため、すべてのデバイスで精度が同等になります。

6. いつWASMを使用すればよいか？

一般にWASMは、モデルが小さい場合、幅広いデバイスサポートを重視する場合、プロジェクトが数値の安定性に敏感な場合に適しています。ただし、WASMはWebGLバックエンドと同等ではありません。WASMバックエンドを使用していて、opを実装する必要がある場合は、Githubに問題を報告してください。本番ユースケースのニーズに対処するために、訓練よりも推論を優先しました。ブラウザでモデルを訓練するには、WebGLバックエンドの使用をお勧めします。

Node.jsのWASMバックエンドは、TensorFlowバイナリをサポートしていない、またはソースからビルドしたくないデバイス向けの優れたソリューションです。以下の表は、WebGL、WASM、およびプレーンJS（CPU）バックエンドで公式にサポートされているいくつかのモデルの2018 MacBook Pro（Intel i7 2.2GHz、Radeon 555X）上のChromeにおける推論時間（ミリ秒単位）を示しています。

WASMバックエンドは、モデル全体でプレーンなJS（CPU）バックエンドよりも10〜30倍高速であることがわかります。WASMをWebGLと比較すると、主に2つのポイントがあります。

(1) WASMは、MediaPipeのBlazeFaceやFaceMeshのような超軽量モデルの場合、WebGLよりも同等または高速です。

(2) WASMは、MobileNet、BodyPix、PoseNetなどの中規模エッジモデルの場合、WebGLより2〜4倍遅くなります。

7. 今後の展望

WASMはますます好まれるバックエンドになると信じています。昨年、WASMバックエンドが理想的に適合するエッジデバイス（MediaPipeのBlazeFaceやFaceMeshなど）用に設計された生産品質の超軽量モデルの波が見られました。

さらに、SIMDやスレッドなどの新しい拡張機能が積極的に開発されており、将来さらに加速することができます。

◎ SIMD / QFMA
SIMD命令を追加するWASM拡張の提案があります。現在、Chromeは実験的なフラグの下でSIMDを部分的にサポートしています。FirefoxとEdgeのステータスは開発中ですが、Safariからは何も公開されていません。SIMDは非常に有望です。一般的なMLモデルでSIMD-WASMを使用したベンチマークは、SIMD以外のWASMに比べて2〜3倍の速度向上を示しています。

元のSIMD提案に加えて、LLVM WASMバックエンドは最近、カーネルのパフォーマンスをさらに改善するはずの実験的なQFMA SIMD命令をサポートしました。一般的なMLモデルのベンチマークは、QFMA SIMDが通常のSIMDよりもさらに26〜50％高速化することを示しています。

TF.js WASMバックエンドは、WASM SIMD用に最適化されたマイクロカーネルを含むXNNPACKライブラリを通じてSIMDを利用します。SIMDが着地すると、これはTensorFlow.jsユーザーには見えなくなります。

◎ マルチスレッド
WASM仕様は最近、マルチスレッドアプリを高速化することを目的としたスレッドとアトミックの提案を得ました。この提案は初期段階であり、将来のW3Cワーキンググループのシードを目的としています。特に、Chrome 74以降ではWASMスレッドのサポートがデフォルトで有効になっています。

スレッド化の提案が出たら、TensorFlow.jsのユーザーコードを変更することなく、XNNPACKライブラリを介してスレッドを利用できるようになります。

8. 詳細につて

・詳細については、WebAssemblyガイドを参照してください。
・Mozilla Developer Networkによるこのリソースのコレクションをチェックして、WebAssemblyの詳細を学んでください。
・GitHubの問題やPRを通じてフィードバックや貢献をお願いします

この記事が気に入ったらサポートをしてみませんか？