tf.function内部の計算グラフ最適化処理について

TensorFlow2.0 Advent Calendar 2019 の22日目の記事になります。

本記事では、tf.functionでPythonプログラムからTensorFlowの計算グラフを構築するときに、TensorFlowが内部で行っている計算グラフの最適化処理とその最適化を制御する方法について紹介します。

なお本記事は、tf.functionの裏で行われていることを説明した記事であるため、tf.functionそのものの使い方に関しては、他の記事を参考にしてください。
TensorFlow2.0 Advent Calender 2019 でも、tf.functionに関する記事がいくつか投稿されています。

• tf.function の使い方について - https://qiita.com/Ryuichirou/items/66a75610c569a23ac493
• tf.functionの注意点とTracingについて - https://qiita.com/t_shimmura/items/1209d01f1e488c947cab

tf.function

@tf.function でデコレートされた関数は、TensorFlowの計算グラフに変換されます。
例えば、次のようなプログラムを考えてみます。

import tensorflow as tf

@tf.function
def simple_func(arg):
    a = tf.constant(7.9)
    b = tf.constant(6.3)
    c = arg + a
    d = a * b
    ret = c + d

    return ret

arg = tf.constant(8.9)
print(simple_func(arg))

simple_func 関数によって作成される計算グラフは、次のようになります。

計算グラフの最適化処理

tf.functionによって変換された計算グラフは、TensorFlowの内部（C++層）で最適化されます。
この最適化処理は、TensorFlow 1.xのデフォルトであったGraph Modeでも使われており、TensorFlow 1.xにて培われた計算グラフ最適化技術がtf.functionの中でも活用されています。

TensorFlow 1.13を対象とした古い記事になりますが、TensorFlow内部で行われている計算グラフの最適化処理に興味のある方は、次の記事も参考にしてみてください。
なお、TensorFlow 2.0では、記事で紹介している最適化に加えて、Auto Mixed Precisionなどの最適化が新たに追加されています。
また別の機会に、最新の最適化処理について紹介したいと思います。

• TensorFlow内部構造解析 (4.4) 計算グラフ最適化処理1 Grappler
• TensorFlow内部構造解析 (4.6) 計算グラフ最適化処理2 GraphOptimizationPass
• TensorFlow内部構造解析 (4.7) 計算グラフ最適化処理3 GraphOptimizer

さて、先ほどの計算グラフに話を戻すと、tf.functionによって生成された計算グラフを最適化したあとの計算グラフは、以下のようになります。

生成された計算グラフを見ると、計算グラフ内のいくつかのノードが削除されていることがわかります。
ここで行われた最適化は、Constant Folding と呼ばれるもので、ノードの入力が全て定数値であるものを計算グラフ構築時に評価し、Constノードで置き換えます。
計算グラフを構築する時点で、評価できるところを事前に評価してしまうことで、計算グラフ全体としての処理時間を短縮できます。

最適化された計算グラフを確認する

最適化された計算グラフは、TensorBoardを使って確認できます。
TensorBoard向けのSummaryデータを出力するためには、tf.functionによって構築された計算グラフを呼び出す前に tf.summary.trace_on() を呼んでおく必要があります。
なお、tf.summary.trace_on() の引数 graph と profiler に True を指定しないと、最適化された計算グラフが出力されない点に注意が必要です。
そして、確認したい計算グラフを実行したあとに tf.summary.trace_export() を呼ぶことで、最適化された計算グラフを出力できます。
最適化された計算グラフを、TensorBoard向けに出力するソースコードを次に示します。

import tensorflow as tf

@tf.function
def simple_func(arg):
    a = tf.constant(7.9)
    b = tf.constant(6.3)
    c = arg + a
    d = a * b
    ret = c + d

    return ret

# TensorBoard向けに、Summaryデータの収集を有効化する
writer = tf.summary.create_file_writer("summary")
# 引数graphとprofilerにTrueを指定することで、最適化された計算グラフを確認できる
tf.summary.trace_on(graph=True, profiler=True)

arg = tf.constant(8.9)
print(simple_func(arg))

# 収集したSummaryデータを出力する
with writer.as_default():
    tf.summary.trace_export("summary", step=0, profiler_outdir="summary")

# Summaryデータの収集を無効化する
tf.summary.trace_off()

プログラムを実行して出力されたTensorBoard向けのSummaryデータを、TensorBoardを使って読み込みます。
最初に、ユーザが定義したグラフを確認してみましょう。
TensorBoardの「GRAPHS」タブを選択した状態で、左側にあるラジオボックスから「Graph」を選択することで、ユーザが定義したグラフを表示できます。

ユーザが定義したグラフは、simple_func 内で定義した計算グラフそのものになっています。

つづいて、最適化後の計算グラフを確認してみましょう。
TensorBoardの「GRAPHS」タブを選択した状態で、左側にあるラジオボックスから「Profile」を選択することで、最適化後の計算グラフを表示できます。

最適化されたあとの計算グラフでは、1つのMulノードとその入力のConstantノードが網掛けされています。
網掛けされたノードは、TensorFlow内部で演算されなかったノードです。
TensorFlow内部で計算グラフが最適化された結果、これらの網掛けされたノードが演算されなくなったと考えられます。

このように、TensorBoardを利用することで、ユーザが定義した計算グラフと最適化後の計算グラフの違いを確認できます。

計算グラフの最適化処理を制御する

TensorFlowの内部で行われている計算グラフの最適化処理は、tf.config.optimizer.set_experimental_options() を呼び出すことで、有効/無効を切りかえることができます。
なお、tf.config.optimizer.set_experimental_options() によって有効/無効を切りかえられる最適化処理は、Grapplerによって行われる最適化処理 のみです。
GraphOptimizationPass や GraphOptimizer による最適化は、必ず実行されてしまうことに注意してください。

計算グラフの最適化処理について、有効/無効を切りかえるプログラムについて説明する前に、デフォルトの最適化の設定を確認してみましょう。
計算グラフの最適化に関する設定は、tf.config.optimizer.get_experimental.options() を呼び出すことで確認できます。

import tensorflow as tf

tf.config.optimizer.get_experimental_options()

上記を実行すると、以下の結果が得られます。

{'disable_meta_optimizer': False, 'disable_model_pruning': False}

disable_meta_optimizer は、Grapplerによって行われる最適化処理 を無効化する設定で、デフォルトでは False が指定されています。
このことから、デフォルトでGrapplerによる最適化処理が有効化されていることがわかります。
また、そのほかの最適化に関しては設定されていないため、それぞれ デフォルトの最適化設定 が適用されています。

それでは実際に、最適化を有効化/無効化する例を挙げながら、その効果を確認していきたいと思います。

最適化「Debug Stripper」を無効化する

Debug Stripper は、デバッグ目的に使われるノード（Assertなど）を削除する最適化処理です。
Debug Stripperは、デフォルトで無効化されているため、tf.Assert によって追加されたAssertノードは削除されません。
このため、次に示すコードは tf.Assert で例外が発生します。

import tensorflow as tf

@tf.function
def assert_func():
    a = tf.constant(1.2)
    computation_graph = tf.Assert(tf.less_equal(a, 1.0), [a])   # 例外「InvalidArgumentError」が発生する
    return a

print(assert_func())

一方でDebug Stripperを有効化して実行すると、tf.Assert によって追加されたAssertノードは削除され、上記で発生していた例外は発生しなくなります。

import tensorflow as tf

# 「Debug Stripper」を有効化する
tf.config.optimizer.set_experimental_options({'debug_stripper': True})

@tf.function
def assert_func():
    a = tf.constant(1.2)
    computation_graph = tf.Assert(tf.less_equal(a, 1.0), [a])   # 例外は発生しない
    return a

print(assert_func())

デバッグ用途で追加したAssertなどは、テンソルデータを確認する必要があるなど処理に時間がかかるため、計算グラフの実行時間に影響を与えます。
デバッグが完了したあとは、デバッグ目的で追加した tf.Assert などを1つ1つ削除するのもよいですが、ここで示した方法でDebug Stripperを有効化するだけで、デバッグ用途の演算がすべて削除されるため、ぜひ活用してみてはいかがでしょうか。

Grapplerで行われる全ての計算グラフの最適化を無効化する

Grapplerで行われる全ての最適化処理は、disable_meta_optimizer を True に設定することで無効化されると書きましたが、ここではその効果を確認してみたいと思います。

最初に、デフォルトの最適化設定で最適化後の計算グラフを確認してみましょう。
次に示すソースコードは、Transposeが連続した計算グラフを構築します。

import tensorflow as tf
import numpy as np

@tf.function
def optimized(arg):
    a = arg * 2

    # 「Arithmetic Optimizer」により、削除される
    b = tf.transpose(a, perm=[1, 0])
    ret = tf.transpose(b, perm=[1, 0])

    return ret

writer = tf.summary.create_file_writer("summary")
tf.summary.trace_on(graph=True, profiler=True)

arg = tf.constant(np.random.normal(size=(30, 40)))
optimized(arg)

with writer.as_default():
    tf.summary.trace_export("summary", step=0, profiler_outdir="summary")

tf.summary.trace_off()

TensorBoardで最適化されたあとの計算グラフを確認すると、Transposeのノードが削除されていることが確認できます。
これは、Arithmetic Optimizer の RemoveIdentityTranspose によって互いに打ち消しあう転置のペアを削除したためです。
また、演算に影響を与えないIdentityノードも、最適化によって削除されていることがわかります。

続いて、disable_meta_optimizer を True にした状態で同じ計算グラフを実行し、TensorBoardで最適化されたあとの計算グラフを確認してみましょう。

import tensorflow as tf
import numpy as np

# Grapplerによって行われる、全ての計算グラフ最適化処理を無効化する
tf.config.optimizer.set_experimental_options({'disable_meta_optimizer': True})

@tf.function
def not_optimized(arg):
    a = arg * 2
    b = tf.transpose(a, perm=[1, 0])
    ret = tf.transpose(b, perm=[1, 0])

    return ret

writer = tf.summary.create_file_writer("summary")
tf.summary.trace_on(graph=True, profiler=True)

arg = tf.constant(np.random.normal(size=(30, 40)))
not_optimized(arg)

with writer.as_default():
    tf.summary.trace_export("summary", step=0, profiler_outdir="summary")

tf.summary.trace_off()

最適化されたあとの計算グラフを見てみると、Transposeのノードがそのまま残っており、Arithmetic Optimizerが無効化されていることがわかります。
また、Identityノードが削除されていないことにも注目です。

まとめ

tf.functionによって構築された計算グラフが、TensorFlow内でどのように最適化されるかを紹介し、その最適化を制御する方法を説明しました。
最近でも新たな最適化が追加されるなど、計算グラフの最適化機能は継続して開発される見込みです。今後に期待しましょう。

なお、本記事で紹介した内容は、英語版のドキュメントとして Pull Request を出し、先日無事にマージされました。
TensorFlowのドキュメントとして正式に公開されると思いますので、実際に計算グラフの最適化が動作していることを、Google Colab上で体験してもらえればと思います。
今回マージされたドキュメントは、いずれ 日本語訳も行ってTensorFlowのドキュメントにContributeしたい と考えていますので、ご期待ください。