TensorFlow : Performance : XLA JIT コンピレーションを使用する（翻訳/解説）
翻訳 : (株)クラスキャット セールスインフォメーション
作成日時 : 03/22/2017
* 本ページは、TensorFlow の本家サイトの Performance – Using JIT Compilation を翻訳した上で
適宜、補足説明したものです：
https://www.tensorflow.org/performance/xla/jit
* ご自由にリンクを張って頂いてかまいませんが、sales-info@classcat.com までご一報いただけると嬉しいです。
【注意】XLA を含めるためには TensorFlow をソースからコンパイルしなければなりません。
何故 just-in-time (JIT) コンピレーションを使用するのか？
TensorFlow/XLA JIT コンパイラは XLA を通して TensorFlow グラフのパーツをコンパイルして実行します。標準の TensorFlow 実装を超えるこれのメリットは XLA が複数の演算を少数のコンパイル済みのカーネルに融合できることです（カーネル・フュージョン）。演算の融合は、TensorFlow executor が行なっているような、演算を一つずつ (one-at-a-time) 実行することに比べてメモリ帯域幅の要求を減少させてパフォーマンスを改善することができます。
XLA を通して TensorFlow グラフを実行する
XLA を通して TensorFlow 計算を実行するには２つの方法があります、CPU または GPU デバイス上に置かれた JIT-コンパイリング演算子によるか、XLA_CPU または XLA_GPU TensorFlow デバイス上に演算子を置くことによるかです。演算を TensorFlow XLA デバイス上に直接置くことは演算にそのデバイス上で実行されることをことを強制し主としてテストのために使用されます。
【注意】 XLA CPU バックエンドは（多くの場合）高速な単一スレッドのコードを生成しますが、TensorFlow CPU バックエンドと同様にまだ並列化はしていません。XLA GPU バックエンドは標準 TensorFlow 実装と競争力があり、ときにより速く、ときにより遅いです。
JIT コンピレーションを有効にする
JIT コンピレーションは session レベルでもあるいは手動で選択した演算子のためにでも有効にすることができます。これらのアプローチの両者ともゼロコピーです — データはデータが渡されるときに同じデバイス上に置かれた、コンパイル済みの XLA カーネルと TensorFlow 演算子の間でコピーされる必要はありません。
Session
session レベルで JIT コンピレーションが有効にされると全ての可能な演算子は貪欲に XLA 計算にコンパイルされる結果になるでしょう。各 XLA 計算は基礎となるデバイスのために一つあるいはそれ以上のカーネルにコンパイルされます。
２、３の制約のもとで、グラフに２つの隣接する演算がありそれらが両者とも XLA 実装を持つのであれば、それらは単一の XLA 計算にコンパイルされます。
JIT コンピレーションは、global_jit_level config に tf.OptimizerOptions.ON_1 をセットして session 初期化の間に config を渡すことによって session レベルで有効になります。
# Config to turn on JIT compilation config = tf.ConfigProto() config.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.ON_1 sess = tf.Session(config=config)
【注意】JIT を session レベルで有効にすることは演算が CPU のためにコンパイルされる結果にはなりません。CPU 演算のための JIT コンピレーションは後述されるような手動の方法を通して行なわれなければなりません。この決定は単一スレッドである CPU バックエンドによってなされました。
マニュアル (手動)
JIT コンピレーションはまた一つあるいはそれ以上の演算のために手動で有効にすることもできます。これは、_XlaCompile=true 属性でコンパイルするように演算にタグ付けすることでなされます。これを行なうもっとも単純な方法は、tensorflow/contrib/compiler/jit.py で定義された tf.contrib.compiler.jit.experimental_jit_scope() スコープを通すことです。使用例は :
jit_scope = tf.contrib.compiler.jit.experimental_jit_scope x = tf.placeholder(np.float32) with jit_scope(): y = tf.add(x, x) # The "add" will be compiled with XLA.
_XlaCompile 属性は現在はベストエフォート・ベースでサポートされています。演算がコンパイルできないときは、TensorFlow は通常の実装に静的にフォールバックします。
演算を XLA デバイス上に置く
XLA を通して計算を実行する他の方法は特定の XLA デバイス上に演算を置くことです。この方法は通常はテストのためにだけ使用されます。有効なターゲットは XLA_CPU または XLA_GPU です。
with tf.device("/job:localhost/replica:0/task:0/device:XLA_GPU:0"):
output = tf.add(input1, input2)
標準の CPU と GPU デバイス上の JIT コンピレーションと異なり、これらのデバイスは（データが）デバイスへ/から転送される時にデータのコピーを作成します。余分なコピーは、同じグラフ内で XLA と TensorFlow 演算をミックスすることをコスト高にします。
チュートリアル
このチュートリアルは、JIT を有効にした MNIST softmax の単純なバージョンのトレーニングをカバーします。現在は session レベルの JIT は、これはチュートリアルのために使用しますが、GPU のみをサポートします。
チュートリアルを始める前に、LD_LIBRARY 環境変数あるいは ldconfig が $CUDA_ROOT/extras/CUPTI/lib64 を含むことを確認してください、これは CUDA Profiling Tools Interface ( CUPTI ) のためのライブラリを含みます。
TensorFlow は GPU から tracing 情報を引っ張るために CUPTI を使用します。
Step #1: サンプル・スクリプトを準備する
mnist_softmax_xla.py をダウンロードするか、TensorFlow ソース・トリーの外側のフォルダに移します。
Step #2: XLA なしで実行する
XLA なしでモデルをトレーニングするために python スクリプトを実行します。
python mnist_softmax_xla.py --xla=''
Chrome Trace Event Profiler (browse to chrome://tracing) を使用して、スクリプトが終了したときに作成される timeline ファイルを開きます: timeline.ctf.json。
表示されたタイムラインは複数のグリーン・ボックスを持つ下の図と類似したものになるでしょう、(CPU の場合には) たぶん複数の CPU に渡って。
訳注: 以下は CPU 環境で確認したものです :
Step #3 XLA で実行する
XLA グラフを出力する環境変数を通して XLA のデバッグ機能を有効にして XLA でモデルをトレーニングするための python スクリプトを実行します。
TF_XLA_FLAGS=--xla_generate_hlo_graph=.* python mnist_softmax_xla.py
作成された timeline ファイル (timeline.ctf.json) を開きます。表示されたタイムラインは _XlaLaunch とラベル付けされた一つの長いバーを持つ下の図と類似したものになるでしょう。
_XlaLaunch で何が起きているかを理解するためには、以下に類似したステートメントのコンソール出力を見てみましょう :
computation cluster_0[_XlaCompiledKernel=true,_XlaNumConstantArgs=1].v82 [CPU: pipeline start, before inline]: /tmp/hlo_graph_0.dot
コンソール・ステートメントは hlo_graph_xx.dot ファイルの位置を指し示し、これは XLA で作成されたグラフについての情報を含みます。Ops を融合するために XLA が取るプロセスは hlo_graph_0.dot から開始して続く各図を見ることで可視化されます。
.dot ファイルを png に展開するためには、 GraphViz をインストールして実行します :
dot -Tpng hlo_graph_80.dot -o hlo_graph_80.png
結果は次のようなものです :
以上