TensorFlowは、TensorBoardなどを使ってユーザが定義した計算グラフを確認することができます。
しかし、次のQiitaの記事で説明している最適化処理がTensorFlow内で行われるため、実際にCPUやGPUなどのデバイスで実行される計算グラフは、ユーザが定義した計算グラフとは異なるものになっている可能性があります。
- TensorFlow内部構造解析 (4.4) 計算グラフ最適化処理1 Grappler
- TensorFlow内部構造解析 (4.6) 計算グラフ最適化処理2 GraphOptimizationPass
- TensorFlow内部構造解析 (4.7) 計算グラフ最適化処理3 GraphOptimizer
本記事では、これらの最適化処理によって変形された後の計算グラフをTensorBoard上で表示する方法を示します。
最適化後の計算グラフをTensorBoard上で表示する手順は次に示す通りです。
1については、TensorFlowのLow Level APIを使っている場合とHigh Level API（Keras、Estimator）を使っている場合とで手順が異なりますので、それぞれについて手順を示します。
- 最適化後の計算グラフを取得する
- デバイスごとに分割された計算グラフを結合する
- TensorBoardで計算グラフを表示するためのデータに変換する
- TensorBoardで計算グラフを確認する
1. 最適化後の計算グラフを取得する
Low Level APIを使用する場合
Session.run の引数に tf.RunOptions を渡して最適化後の計算グラフを出力するように設定し、Session.run の引数に tf.RunMetadata を渡します。
Session.run 実行後、最適化された後の計算グラフがProtocol Buffers形式で tf.RunMetadata に保存されるため、保存された情報をファイルなどに出力して計算グラフを確認することができます。
参考：https://qiita.com/jack_ama/items/1835372064106de84357
import re
import tensorflow as tf
run_opts = tf.RunOptions()
run_opts.output_partition_graphs = True # ★最適化後の計算グラフを出力する設定
run_md = tf.RunMetadata()
ops = ...
with tf.Session() as sess:
sess.run(opts, run_metadata=run_md, options=run_opts)
# ★最適化後の計算グラフを出力（計算グラフはデバイスごとに分割されている）
for pg in run_md.partition_graphs:
if len(pg.node) == 0:
continue
device = re.sub(r"[/:]", "_", pg.node[0].device)
with open("partitioned_graphs/{}.pbtxt".format(device), "w") as f:
f.write(repr(pg))
tf.RunMetadata には、各演算にかかった実行時間やメモリ使用量など、TensorFlowで行った演算の実行情報が保存されています。
最適化後の計算グラフは tf.RunMetadata.partition_graphs にGraphDefのリストとして保存されています。
上記のプログラムでは、デバイスごとにGraphDefをProtocol Buffersのテキスト形式で出力します。
例えば、デバイス /job:localhost/replica:0/task:0/device:CPU:0 に割り当てられた計算グラフであれば、partitioned_graphs/_job_localhost_replica_0_task_0_device_CPU_0.pbtxt に保存されます。
node {
name: "Mul"
op: "Const"
device: "/job:localhost/replica:0/task:0/device:CPU:0"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_FLOAT
tensor_shape {
}
tensor_content: "\000\000\020A"
}
}
}
experimental_debug_info {
original_node_names: "Mul"
}
}
node {
name: "_retval_Mul_0_0"
op: "_Retval"
input: "Mul"
device: "/job:localhost/replica:0/task:0/device:CPU:0"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "index"
value {
i: 0
}
}
experimental_debug_info {
original_node_names: "_retval_Mul_0_0"
}
}
library {
}
versions {
producer: 55
}
なお、ここでファイルに出力したGraphDefの情報について、TensorFlow内部構造解析 (1.1) Protocol Buffers形式のデータ構造を参考に、計算グラフを解析することができます。
Estimatorを使用する場合
Estimator利用時に最適化後の計算グラフを出力するには、Session.run 実行後に最適化後の計算グラフを出力するための設定を、EstimatorのHookとして登録する必要があります。
最適化後の計算グラフを出力するためには、tf.RunOptions を Session の引数に渡す必要がありますが、これは Session.run 実行前に呼び出される before_run メソッドで実現できます。
サンプルコードでは、最適化後の計算グラフを出力するように設定した tf.RunOptions を Session.run の引数に渡すように tf.train.SessionRunArgs を作成し、戻り値として返しています。
また、Session.run 実行後に最適化後の計算グラフをProtocol Buffers形式でファイル出力するため、Session.run 実行後に呼び出される after_run メソッドに処理を記述します。
最後に、作成したHook RunMetadataHook をEstimatorのHookとして登録します。
import os
import re
import tensorflow as tf
from tensorflow.python.training import training_util
def train_input_fn(features, labels, batch_size):
...
classifier = tf.estimator.LinearClassifier(...)
class RunMetadataHook(tf.train.SessionRunHook):
def begin(self):
# ★Global Stepごとにディレクトリを作成するために、Global Step用のテンソルを作成
self._global_step_tensor = training_util.get_global_step()
if self._global_step_tensor is None:
raise RuntimeError("Failed to create global step tensor")
def before_run(self, run_context):
# ★最適化後の計算グラフを出力する設定
opts = tf.RunOptions(output_partition_graphs=True)
requests = {"global_step": self._global_step_tensor}
return tf.train.SessionRunArgs(requests, options=opts)
def after_run(self, run_context, run_values):
global_step = run_values.results["global_step"]
run_md = run_values.run_metadata
# ★Step毎にディレクトリを作成し、最適化後の計算グラフをファイルに出力
dirpath = "partitioned_graphs/global_step_{}".format(global_step)
if not os.path.exists(dirpath):
os.makedirs(dirpath)
for pg in run_md.partition_graphs:
if len(pg.node) == 0:
continue
device = re.sub(r"[/:]", "_", pg.node[0].device)
with open("{}/{}.pbtxt".format(dirpath, device), "w") as f:
f.write(repr(pg))
# ★分割後の計算グラフをファイルに出力するためのHookを登録
run_md_hook = RunMetadataHook()
classifier.train(input_fn=lambda: train_input_fn(train_x, train_y, 100),
hooks=[run_md_hook], steps=2000)
上記のサンプルを実行すると、Global Stepが10の時かつデバイス /job:localhost/replica:0/task:0/device:CPU:0 に割り当てられた計算グラフは、 partitioned_graphs/global_step_10/_job_localhost_replica_0_task_0_device_CPU_0.pbtxt に保存されます。
Kerasを使用する場合
Keras利用時は、model.compile の引数に tf.RunMetadata や tf.RunOptions を渡すことができます。
model.fit 後、tf.RunMetadata に最適化後の計算グラフが保存されているため、Low Level APIを使用する場合 と同様にデバイスごとにGraphDefをProtocol Buffersのテキスト形式で出力します。
import re
import tensorflow as tf
import keras
model = keras.Sequential()
...
run_opts = tf.RunOptions()
run_opts.output_partition_graphs = True # ★最適化後の計算グラフを出力する設定
run_md = tf.RunMetadata()
model.compile(..., options=run_opts, run_metadata=run_md)
model.fit(...)
# ★最適化後の計算グラフを出力（計算グラフはデバイスごとに分割されている）
for pg in run_md.partition_graphs:
if len(pg.node) == 0:
continue
device = re.sub(r"[/:]", "_", pg.node[0].device)
with open("partitioned_graphs/{}.pbtxt".format(device), "w") as f:
f.write(repr(pg))
注意事項
Kerasでは、model.fit の中で複数回 Session.run 相当の処理が行われます。
tf.RunMetadata に保存されるGraphDefは、model.fit 内の最後の Session.run で実行された計算グラフであることに注意が必要です。
現在のところ、ステップ毎に tf.RunMetadata を出力する方法は存在しないようです。
2. デバイスごとに分割された計算グラフを結合する
tf.RunMetadata.partition_graphs に保存されたGraphDefは、デバイスごとに分割された計算グラフであるため、計算グラフ全体を確認するためには分割後のそれぞれの計算グラフを確認する必要があり、面倒です。
そこで、GraphDefが保存された複数のProtocol Buffers形式のファイルを解析し、分割後の計算グラフを再結合するためのプログラムを作りました。
上記のリポジトリに置かれている merge_partitioned_graph.py を使って、デバイスごとに分割されたGraphDefを1つのGraphDefに結合します。
前提条件として、以下のようなプログラムなどで、デバイスごとに分割されたGraphDefをProtocol Buffers形式のファイルとして、別々のファイルに保存されてるものとします。
for pg in run_md.partition_graphs:
if len(pg.node) == 0:
continue
device = re.sub(r"[/:]", "_", pg.node[0].device)
with open("partitioned_graphs/{}.pbtxt".format(device), "w") as f:
f.write(repr(pg))
上記のプログラムは、仮にCPUと1つのGPUが利用できる環境であれば、2つのProtocol Buffers形式のファイル partitioned_graphs/_job_localhost_replica_0_task_0_device_CPU_0.pbtxt と partitioned_graphs/_job_localhost_replica_0_task_0_device_GPU_0.pbtxt を出力します。
これらのファイルを、前述の計算グラフ再結合プログラムを使って1つのGraphDefに結合し、ファイル merged.pbtxt に書き出すコマンドを以下に示します。
$ python merge_partitioned_graph.py \
--inputs partitioned_graphs/_job_localhost_replica_0_task_0_device_CPU_0.pbtxt \
partitioned_graphs/_job_localhost_replica_0_task_0_device_GPU_0.pbtxt \
--output merged.pbtxt
3. TensorBoardで計算グラフを表示するためのデータに変換する
手順2で結合した、GraphDefが保存されたProtocol Buffers形式のファイルを、TensorBoardで表示するためのデータ形式に変換します。
GraphDefが保存されたProtocol Buffers形式のファイルを、TensorBoardで表示するためのデータ形式に変換するプログラムもGitHubに公開しています。
上記のリポジトリに置かれている output_summary_from_pb.py を使って、結合後のGraphDefが保存されたファイル merged.pbtxt から、TensorBoardで表示するためのデータ形式へ変換するためのコマンドを次に示します。
$ python output_summary_from_pb.py --input merged.pbtxt --output_dir summary
上記コマンドを実行すると、ディレクトリ summary の下にTensorBoardで表示するためのデータが保存されます。
4. TensorBoardで計算グラフを確認する
手順4で作成したTensorBoardで表示するためのデータが配置されているディレクトリを指定してTensorBoardを起動することにより、最適化後の計算グラフをTensorBoard上で確認することができます。
$ tansorboard --logdir=./summary
上記のコマンドを実行すると、TensorBoard上から次に示す計算グラフが確認できます。
上記の図は、次に示すプログラムを実行して手順1から手順4までを行ったときにTensorBoard上で表示される計算グラフで、ユーザが定義した計算グラフとは異なる計算グラフがTensorBoard上で表示されていることがわかります。
例えば、計算グラフのノード名がMulであるのに対し、実際のOperationは Const と表示されています。これは、TensorFlowの内部で行われる計算グラフ最適化処理の1つ「Constant Folding」がユーザが定義した計算グラフに対して適用されたことを示しています。
import re
import tensorflow as tf
run_opts = tf.RunOptions()
run_opts.output_partition_graphs = True
run_md = tf.RunMetadata()
with tf.Session() as sess:
a = tf.constant(1.0)
b = tf.constant(2.0)
c = tf.add(a, b)
d = tf.add(a, b)
out = tf.multiply(c, d)
result = sess.run(out, run_metadata=run_md, options=run_opts)
for pg in run_md.partition_graphs:
if len(pg.node) == 0:
continue
device = re.sub(r"[/:]", "_", pg.node[0].device)
with open("partitioned_graphs/{}.pbtxt".format(device), "w") as f:
f.write(repr(pg))
このように本記事で説明した手順を踏むことで、最適化後の計算グラフをTensorBoardで表示することができます。