[論文レビュー] TensorFlow Eager: A Multi-Stage, Python-Embedded DSL for Machine Learning

Akshay Agrawal, Akshay Naresh Modi+9 more|arXiv (Cornell University)|Feb 27, 2019

Parallel Computing and Optimization Techniques参考文献 26被引用数 52

ひとこと要約

TensorFlow Eager は Python 関数をデータフロウグラフへステージングできる JIT トレーサを備えた TensorFlow への命令型フロントエンドを提供し、機械学習ワークロードのための命令型とグラフ実行のシームレスな補間を可能にします。

ABSTRACT

TensorFlow Eager is a multi-stage, Python-embedded domain-specific language for hardware-accelerated machine learning, suitable for both interactive research and production. TensorFlow, which TensorFlow Eager extends, requires users to represent computations as dataflow graphs; this permits compiler optimizations and simplifies deployment but hinders rapid prototyping and run-time dynamism. TensorFlow Eager eliminates these usability costs without sacrificing the benefits furnished by graphs: It provides an imperative front-end to TensorFlow that executes operations immediately and a JIT tracer that translates Python functions composed of TensorFlow operations into executable dataflow graphs. TensorFlow Eager thus offers a multi-stage programming model that makes it easy to interpolate between imperative and staged execution in a single package.

研究の動機と目的

• ML ワークフローのために、命令型の使いやすさとグラフベースの最適化を組み合わせた DSL の必要性を動機づける。
• デフォルトで命令型実行をサポートし、オプションのグラフステージングを可能とする Python 埋め込みのマルチステージプログラミングモデルを提案する。
• Python 関数を実行可能なデータフローグラフに変換するオプトイン tracing ベースの JIT の実装を説明する。
• 実行モード間のシームレスな移行を可能にするために、状態、デバイス、および分散がどのように表現されるかを説明する。
• 代表的な ML モデルに対して、命令型実行とステージド実行のパフォーマンスのトレードオフを評価する。

提案手法

• TensorFlow に層状に組み込まれた Python 埋め込み DSL を導入し、デフォルトで命令型に実行でき、ステージング時に Python 関数をトレースしてデータフローグラフを構築する。
• Python 関数をトレースして名前付き入力/出力を持つグラフ関数を作成し、実行用の C++ データフロー ランタイムを提供する関数デコレータを提供する。
• ネストされたテープを用いた高階導関数のための tracing ベースの逆モード微分 (AD) による自動微分を説明する。
• シリアライズと復元のための変数とグラフベースの状態マッチングによる状態管理を説明する。
• 自動デバイス配置、デバイス間データ転送、および TPU/XLA 支援のステージドパスを含むデバイスと分散のサポートについて議論する。
• グラフ内に命令型コードを埋め込み、トレース時の状態作成を管理するためのエスケープおよび統合メカニズム（py_func、tf.init_scope）を概説する。

実験結果

リサーチクエスチョン

• RQ1 データフローグラフの利点を犠牲にすることなく、命令型に優しい TensorFlow のフロントエンドを、グラフベースの実行と補完するように設計するにはどうすればよいか？
• RQ2 命令型実行とステージドグラフ構築を混在させるマルチステージプログラミングモデルの設計・実装・性能のトレードオフは何か？
• RQ3 両方の実行モードで、状態、デバイス配置、分散を一貫して管理するにはどうするか？
• RQ4 tracing ベースの自動微分は、部分的にステージドな計算や Python の制御フローとどのように相互作用するか？

主な発見

• 命令型 TensorFlow Eager は、命令型実行を用いると単一 GPU で ResNet-50 のグラフ性能に匹敵できる。
• ステージド TensorFlow Eager は TPU 加速トレーニングの ResNet-50 を、TPU 上で命令型 TensorFlow Eager が達成できるよりはるかに高速に実行できる。
• ステージングは、操作が小さなモデルに対して顕著な速度向上をもたらし、最小限のコード変更で済む。
• このフレームワークは、実行モードに依存しない一貫した単一 API 表面を提供し、ツールとデプロイメントのために TensorFlow エコシステムを活用する。
• 実装は軽量で、Python レベルのステージング約2000行、コアは Python と C で差異を持ち、クロスプラットフォーム。