SIMD およびマルチスレッド処理で TensorFlow.js WebAssembly バックエンドを高速化する

この記事は Ann YuanとMarat Dukhanによる The TensorFlow Blog の記事 "Supercharging the TensorFlow.js WebAssembly backend with SIMD and multi-threading" を元に翻訳・加筆したものです。詳しくは元記事をご覧ください。

 投稿者: Google ソフトウェア エンジニア、Ann Yuan、Marat Dukhan



3 月に、WebAssembly（Wasm）を使って高速化した新しい TensorFlow.js 用バックエンドを紹介しました（Wasm について、そしてこれが重要である理由については、スクロールして以降をお読みください）。本日は、大幅なパフォーマンス アップデートについてお知らせします。TensorFlow.js バージョン 2.3.0 時点で、Wasm バックエンドは最大 10 倍高速になっています。これは、ニューラル ネットワーク演算用に高度に最適化されたライブラリである XNNPACKを通して SIMD（ベクター）命令とマルチスレッド処理を活用することで実現しました。

ベンチマーク

SIMD とマルチスレッド処理により、Wasm バックエンドのパフォーマンスが大幅に向上します。下の Google Chrome でのベンチマークは、BlazeFaceモデルでの改善を示すものです。BlazeFace は、10 万個のパラメータがあり、約 2000 万回の積和演算を行う軽量モデルです。

（記載されている時間は、1 回の推論あたりのミリ秒です）

それよりも大きなモデルでは、さらにスピードが向上します。350 万個のパラメータがあり、約 3 億回の積和演算を行う中規模モデルである MobileNet V2などがこれに相当します。

*注: モバイル ブラウザでのマルチスレッドのサポートは現在作業中であるため、Pixel 4 では TF.js マルチスレッド Wasm バックエンドのベンチマークは利用できません。iOS での SIMD サポートも現在開発中です。

**注: TF.js マルチスレッド Wasm バックエンドは、近日中に Node をサポートする予定です。

SIMD とマルチスレッドによるパフォーマンス向上は、それぞれ独立しています。以上のベンチマークから、SIMD を使うと Wasm のみの場合に比べて 1.7～4.5 倍パフォーマンスが向上し、マルチスレッドを使うとそこからさらに 1.8～2.9 倍スピードアップすることがわかります。

使用方法

SIMD は TensorFlow.js 2.1.0 で、マルチスレッド処理は TensorFlow.js 2.3.0 でサポートされます。

実行時に SIMD とマルチスレッド処理がサポートされているかどうかを自動的に確認し、それに基づいて適切な Wasm バイナリを取得します。現在は、以下のそれぞれの場合ごとに異なるバイナリを提供しています。

• デフォルト: このランタイムは、SIMD もマルチスレッド処理もサポートしません。
  
• SIMD: このランタイムは、SIMD をサポートしますがマルチスレッド処理はサポートしません。
  
• SIMD + マルチスレッド処理: このランタイムは、SIMD とマルチスレッド処理をサポートします。
  

マルチスレッド処理をサポートするランタイムのほとんどは SIMD もサポートするので、バンドルのサイズを小さく保つために、マルチスレッド処理のみのランタイムは省略しています。つまり、マルチスレッド処理をサポートし、SIMD はサポートしないランタイムには、デフォルトのバイナリが提供されます。Wasm バックエンドを使用する方法は 2 つあります。

1. NPM を利用
   

   // Import @tensorflow/tfjs or @tensorflow/tfjs-core
   const tf = require('@tensorflow/tfjs');
   // Add the WAsm backend to the global backend registry.
   require('@tensorflow/tfjs-backend-wasm');
   
   // Set the backend to WAsm and wait for the module to be ready.
   tf.setBackend('wasm').then(() => main());

   ライブラリは、Wasm バイナリがメイン JS ファイルに対して相対的に配置されることを前提としています。parcel や webpack などのバンドラーを使っている場合は、setWasmPathsヘルパーを使って手動で Wasm バイナリの場所を指定する必要があるかもしれません。
   

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

   詳しくは、README の “Using bundlers”セクションをご覧ください。
   
2. script タグを利用
   

   <!-- Import @tensorflow/tfjs or @tensorflow/tfjs-core -->
   <script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs"></script>
   
   <!-- Adds the WAsm backend to the global backend registry -->
   <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’)を呼び出す必要があります。
   

デモ

パフォーマンスの向上を自分の目で確かめてみたい方は、https://tfjs-wasm-simd-demo.netlify.app/から BlazeFace モデルのデモをご覧ください。このモデルは、新しい Wasm バックエンドを使うようにアップデートされています。最適化されていないバイナリと比較したい方は、こちらのバージョンのデモをお試しください。手動で SIMD とマルチスレッド処理のサポートをオフにしてあります。

Wasm とは

WebAssembly（Wasm）は、ウェブでネイティブに近いコード実行速度を実現するクロスブラウザ バイナリ フォーマットです。Wasm は、C、C++、Go、Rust などの静的に型付けされた高レベル言語で書かれたプログラムからコンパイルします。TensorFlow.js の Wasm バックエンドは、C++ で実装し、Emscripten でコンパイルしています。XNNPACK ライブラリは、その下層として、高度に最適化されたニューラル ネットワーク演算の実装を提供します。
Wasm は、2017 年より Chrome、Safari、Firefox、Edge でサポートされており、世界中の端末の 90% で利用できます。
WebAssembly の仕様は短時間で進化し、各ブラウザは増加する試験運用版機能をサポートする努力を続けています。このサイトにアクセスすると、ランタイムがどの機能をサポートしているかを確認できます。サポートされる機能には、以下が含まれます。

1. SIMD
   SIMD は、Single Instruction, Multiple Data（単一命令複数データ）の略です。つまり、SIMD 命令は、個々のスカラーではなく、小さな固定サイズの要素のベクトルに対して実行されます。Wasm SIMD 提案は、最近のプロセッサがサポートする SIMD 命令をウェブブラウザの内部で使えるようにするものです。これにより、大幅なパフォーマンス向上を実現できます。
   
   Wasm SIMD はフェーズ 3提案で、Chrome 84～86 のオリジン トライアルで利用できます。つまり、デベロッパーはウェブサイトで Wasm SIMD をオプトインすることができ、すべてのユーザーは、ブラウザの設定で機能を明示的に有効化することなく、Wasm SIMD を利用できます。Google Chrome のほかに、Firefox Nightly もデフォルトで Wasm SIMD をサポートしています。
   
   
2. マルチスレッド処理
   最近のプロセッサには、ほぼ例外なく複数のコアが搭載されており、それぞれのコアは独立して同時に命令を実行できます。threads 提案は、WebAssembly のプログラムの作業を複数のコアに分散させることでパフォーマンスの向上を図ります。この提案によって、個別のウェブワーカー内にある複数の Wasm インスタンスが 1 つの WebAssembly.Memoryオブジェクトを共有し、ワーカー間で高速な通信を行えるようになります。
   
   Wasm threads はフェーズ 2の提案で、バージョン 74 以降の PC 版 Chrome ではデフォルトで有効になっています。モバイルのブラウザでもこの機能を利用できるようにするため、さまざまなブラウザで作業が続いています。
   

SIMD や threads、その他の試験運用版機能をサポートしているブラウザを確認するには、WebAssembly ロードマップをご覧ください。

その他の改善点

3 月に Wasm バックエンドが初めてリリースされてから、演算子のサポート範囲を広げており、現在は 70 以上の演算子をサポートしています。新しい演算子の多くは XNNPACK ライブラリを使って高速化され、HandPose モデルなどの追加モデルもサポートできるようになっています。

今後の予定

今後も Wasm バックエンドのパフォーマンス改善を続ける予定です。SIMD の範囲を広げるフレキシブル ベクトル、準融合積和演算、疑似最小および疑似最大命令など、WebAssembly でのいくつかの仕様の進化には、時間を空けずに追従しています。WebAssembly モジュールの ES6 モジュールのサポートにも期待しています。これらの機能が TF.js ユーザーコードで制限なく利用できるようになれば、SIMD やマルチスレッド処理と合わせて活用してゆきたいと考えています。

詳細情報

• WebAssembly ロードマップの確認: https://webassembly.org/roadmap/
  
• Wasm 仕様の進捗をフォロー: https://github.com/WebAssembly/spec
  
• Wasm SIMD 提案の詳細を確認: https://github.com/WebAssembly/simd
  
• Wasm threads 提案の詳細を確認: https://github.com/WebAssembly/threads
  
• フィードバックの送信や貢献は、GitHub の Issuesや PRからお願いします。
  
• 今回紹介した最適化を使って構築したものは、ソーシャル メディアで #MadeWithTFJSハッシュタグを付けて共有してください。今後のアップデートについての情報を受け取りたい方は、TensorFlow.js コミュニティのディスカッション フォーラムに参加してください。
  

謝辞

WebAssembly バックエンドおよび XNNPACK の組み込みの基礎作業を行ってくださった Daniel Smilkov 氏と Nikhil Thorat 氏、Pixel 4 ベンチマークの数値を収集してくださった Matsvei Zhdanovich 氏、XNNPACK で低レベル Wasm SIMD 最適化を実装してくださった Frank Barchard 氏に感謝いたします。