Tensorflowのような、計算グラフを構築してから実行するDefine and Runなフレームワークの利点の一つは、計算グラフを実行する前に最適化を行えることです。

計算グラフの最適化では

最終的に不要なノードを計算グラフから削除する
重複する式の計算をまとめる (common subexpression elimination)
定数で計算できるノードを畳み込む (constant folding)

等が行われます。本記事では、実行時のグラフを可視化することで最適化の具体例を示します。

最適化されたグラフの可視化 (実行時グラフ可視化)

計算グラフの最適化はsess.runの段階で行われます。このsess.runにオプションを渡すことで、最適化されたグラフを受け取ることができます。次のようにoutput_partition_graphs=Trueにすることで、run_metadataを介して実行時のグラフが渡されます。

run_options = tf.RunOptions(output_partition_graphs=True)
run_metadata = tf.RunMetadata()

sess.run([ops], run_metadata=run_metadata, options=run_options)

graph = run_metadata.partition_graphs[0]
writer = tf.summary.FileWriter(logdir="./graph/", graph=graph)
writer.flush()

common subexpression elimination

次のような計算を考えます。

v1 = a * 2 + c * 4
v2 = b * 2 + c * 4

v1とv2の計算には、どちらもc4が出現するため、c4は一度だけ計算すれば使いまわせる筈です。

Tensorboardでこの計算グラフを可視化すると次のようになります。

mulが4つあり、c*4の計算結果が使いまわされていないように見えます。しかし、上記の方法で実行時の計算グラフを可視化すると次のようになります。

mulが3つに減っていることがわかります。

constant folding

定数だけに依存するOperatorがあれば実行前に計算しておくことが可能です。
次のような計算グラフを考えます。

a = tf.constant(np.ones([4,32,32,3]), dtype=tf.float32, name="a")
W1 = tf.constant(np.ones([3, 3, 3, 8]), dtype=tf.float32, name="W1")
W2 = tf.constant(np.ones([3, 3, 8, 4]), dtype=tf.float32, name="W2")
W3 = tf.Variable(np.ones([3, 3, 4, 2]), dtype=tf.float32, name="W3")

v1 = tf.nn.conv2d(a, W1, strides=[1, 2, 2, 1], padding="SAME", name="conv1")
v2 = tf.nn.conv2d(v1, W2, strides=[1, 2, 2, 1], padding="SAME", name="conv2")
v3 = tf.nn.conv2d(v2, W3, strides=[1, 2, 2, 1], padding="SAME", name="conv3")

これを可視化すると、

となり、convolutionが3回行わていることがわかります。

v2の計算までは、定数だけで構成されるため、事前に計算が可能です。そのため、実行時の計算グラフを可視化すると

となり、convolutionが実行時には1回に減少していることがわかります。

まとめ

実行時の計算グラフを表示することで、グラフの最適化の様子を観測することができました。
実行時の計算グラフの表示は、各OperatorがどのDeviceにマッピングされたかの表示も行ってくれるため非常に便利な機能です。