TensorFlowで最適化された計算グラフをTensorBoardで確認する（TensorFlow 1.x）

TensorFlowは、TensorBoardなどを使ってユーザが定義した計算グラフを確認することができます。
しかし、次のQiitaの記事で説明している最適化処理がTensorFlow内で行われるため、実際にCPUやGPUなどのデバイスで実行される計算グラフは、ユーザが定義した計算グラフとは異なるものになっている可能性があります。

* [TensorFlow内部構造解析 (4.4) 計算グラフ最適化処理1 Grappler](https://qiita.com/nuka137/items/f1b0fe9c820e4d5f80cc)
* [TensorFlow内部構造解析 (4.6) 計算グラフ最適化処理2 GraphOptimizationPass](https://qiita.com/nuka137/items/808fbec073f97be8b872)
* [TensorFlow内部構造解析 (4.7) 計算グラフ最適化処理3 GraphOptimizer](https://qiita.com/nuka137/items/1936ad509cd4e68d2fad)

本記事では、これらの最適化処理によって変形された後の計算グラフをTensorBoard上で表示する方法を示します。
最適化後の計算グラフをTensorBoard上で表示する手順は次に示す通りです。
1については、TensorFlowのLow Level APIを使っている場合とHigh Level API（Keras、Estimator）を使っている場合とで手順が異なりますので、それぞれについて手順を示します。

1. 最適化後の計算グラフを取得する
2. デバイスごとに分割された計算グラフを結合する
3. TensorBoardで計算グラフを表示するためのデータに変換する
4. 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

```python
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形式のデータ構造](https://qiita.com/nuka137/items/8b6a9d632641fbc787d4)を参考に、計算グラフを解析することができます。


## 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として登録します。


```python
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を使用する場合](#low-level-apiを使用する場合) と同様にデバイスごとにGraphDefをProtocol Buffersのテキスト形式で出力します。


```python
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形式のファイルを解析し、分割後の計算グラフを再結合するためのプログラムを作りました。

* [nuka137/TensorFlow-Merge-Partitioned-Graph](https://github.com/nuka137/TensorFlow-Merge-Partitioned-Graph)


上記のリポジトリに置かれている `merge_partitioned_graph.py` を使って、デバイスごとに分割されたGraphDefを1つのGraphDefに結合します。
前提条件として、以下のようなプログラムなどで、デバイスごとに分割されたGraphDefをProtocol Buffers形式のファイルとして、別々のファイルに保存されてるものとします。

```python
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` に書き出すコマンドを以下に示します。

```bash
$ 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に公開しています。

* [nuka137/TensorFlow-Merge-Partitioned-Graph](https://github.com/nuka137/TensorFlow-Merge-Partitioned-Graph)

上記のリポジトリに置かれている `output_summary_from_pb.py` を使って、結合後のGraphDefが保存されたファイル `merged.pbtxt` から、TensorBoardで表示するためのデータ形式へ変換するためのコマンドを次に示します。

```bash
$ python output_summary_from_pb.py --input merged.pbtxt --output_dir summary
```

上記コマンドを実行すると、ディレクトリ `summary` の下にTensorBoardで表示するためのデータが保存されます。


# 4. TensorBoardで計算グラフを確認する

手順4で作成したTensorBoardで表示するためのデータが配置されているディレクトリを指定してTensorBoardを起動することにより、最適化後の計算グラフをTensorBoard上で確認することができます。

```bash
$ tansorboard --logdir=./summary
```

上記のコマンドを実行すると、TensorBoard上から次に示す計算グラフが確認できます。

![optimized_graph.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/229755/b1cd1774-2ade-3352-f4bf-3d40cf7a1887.png)


上記の図は、次に示すプログラムを実行して手順1から手順4までを行ったときにTensorBoard上で表示される計算グラフで、ユーザが定義した計算グラフとは異なる計算グラフがTensorBoard上で表示されていることがわかります。

例えば、計算グラフのノード名がMulであるのに対し、実際のOperationは **Const** と表示されています。これは、TensorFlowの内部で行われる計算グラフ最適化処理の1つ「Constant Folding」がユーザが定義した計算グラフに対して適用されたことを示しています。

```python
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で表示することができます。