[論文レビュー] Auto-Vectorizing TensorFlow Graphs: Jacobians, Auto-Batching And Beyond

Ashish Agarwal, Igor Ganichev|arXiv (Cornell University)|Mar 8, 2019

Scientific Computing and Data Management参考文献 18被引用数 6

ひとこと要約

本論文は、TensorFlowのハイレベルデータフローIRにおける静的ループベクトル化最適化を導入し、新しい記号的並列for（pfor）抽象化を通じて、自動バッチ処理、例ごとの勾配、ヤコビ行列計算を効率的に行えるようにしている。このアプローチは、TensorFlowのネイティブなループ実装に対してGPUで最大60倍の高速化を達成し、DyNetの動的バッチ処理よりも38倍以上高速であり、特に大規模なバッチサイズと高次元出力の場合に顕著である。

ABSTRACT

We propose a static loop vectorization optimization on top of high level dataflow IR used by frameworks like TensorFlow. A new statically vectorized parallel-for abstraction is provided on top of TensorFlow, and used for applications ranging from auto-batching and per-example gradients, to jacobian computation, optimized map functions and input pipeline optimization. We report huge speedups compared to both loop based implementations, as well as run-time batching adopted by the DyNet framework.

研究の動機と目的

• 自動バッチ処理、例ごとの勾配、ヤコビ行列計算などの操作における、TensorFlowにおけるループベース実装の性能ボトルネックを解消すること。
• DyNetのようなフレームワークにおける動的バッチ処理に伴うランタイムオーバーヘッドを低減し、静的でコンパイル時ベクトル化を可能にする。
• ハードウェア利用効率の向上を図り、機械学習モデルにおけるネストされたおよび動的制御フローのパターンを効率的に実行可能にする。
• TensorFlowのハイレベルIRにおける任意のループをベクトル化するための汎用的かつ再利用可能な抽象化を提供すること。
• ヤコビ行列とヘッセ行列の効率的計算をTensorFlowがネイティブにサポートできるように拡張すること。これらは現在、最適化が不十分または非対応である。

提案手法

• グラフ構築時にTensorFlowのハイレベルデータフローIR上で動作する新しい静的ベクトル化並列for（pfor）抽象化を導入する。
• ループ本体をベクトル化された演算に統合することで、静的ループベクトル化を適用し、明白に並列な計算に対してループを排除する。
• 記号的テンソルを用いてベクトル化されたループ反復の出力を表現することで、ネストされた制御フローおよび動的形状の最適化を可能にする。
• バッチ要素に対するループで前方伝搬をラップし、その全体をベクトル化することで、自動バッチ処理用のモデルにpfor抽象化を適用する。
• 例ごとの勾配を計算するために、前方伝搬と後方伝搬の両方を1つのpfor構造内でベクトル化する。
• ヤコビ行列計算を可能にするために、各出力要素を独立したスカラー計算として扱い、勾配計算を各出力要素にわたってベクトル化する。

実験結果

リサーチクエスチョン

• RQ1 TensorFlowのハイレベルIRにおける静的ループベクトル化が、自動バッチ処理や勾配計算において、動的ループベース実装を著しく上回る性能を発揮できるか。
• RQ2 pforベースのベクトル化は、非スカラー出力やRNNのような複雑なモデルにおいて、ヤコビ行列計算の性能をどの程度向上できるか。
• RQ3 DyNetの実行時バッチ処理と比較して、pforの性能は、さまざまなバッチサイズや出力次元でどの程度優れているか。
• RQ4 可変長シーケンスやネストされたループの状況において、pforの現在の実装における性能ボトルネックは何か。
• RQ5 pforは、ベクトル化された制御フローを通じて、従来不可能または非効率的であった操作、例えばヘッセ行列の計算や入力パイプライン最適化を可能にするか。

主な発見

• pforベースの実装は、出力サイズが128の高次元出力において、TensorFlowのネイティブなtf.while_loopに対して最大60倍の高速化を達成しており、特にVGG16とLSTMモデルで顕著である。
• LSTMにおける例ごとの勾配計算では、バッチサイズ256でpforがtf.while_loopを38倍の速度で上回り、GPUの利用効率が優れていることが示された。
• MNISTでは、pforがtf.while_loopに比べて一貫したスループット向上を示し、特に大規模なバッチサイズで、より良いベクトル化とメモリアクセスパターンのおかげで顕著な向上が見られた。
• 出力次元が高くなるにつれて、pforとtf.while_loopの性能差が拡大しており、ベクトル化アプローチの優れたスケーラビリティが裏付けられた。
• LSTMモデルでは、可変長シーケンスとtf.while_loopのベクトル化に伴う残留オーバーヘッドのため、相対的な高速化が低かったが、動的形状の処理における最適化の余地が示唆された。
• 本手法により、従来非対応であった動的RNNのヤコビ行列や非スカラー出力のヘッセ行列の計算が可能となり、新たな研究分野の開拓が可能になった。