tf.functionのメモ

tf.functionの基本

基本的な役割

TensorFlow2.0からはDefine by Runを標準としつつも、Pythonの遅さをカバーするべく、実行ごとに計算グラフが変化することのないようなものはDefine and Run的に使えるように tf.function という関数が追加されました。実態はPython関数をtf.Graphに変換して、あとは適宜このgraphを呼び出すことで利用できるというものになっています。

def python_add():
    return tf.add(1., 2.)

という関数を定義した場合TF2.0ではpython_addを呼び出した段階でtf.addが実行され結果 3.0 が得られます。

def python_add(x, y)
    return tf.add(x, y)

とすれば、この関数は x と y を引数にとり、その和を返します。これはまるっきりPythonの関数と同じように振る舞うことになります。この関数に対して

tf_add = tf.function(python_add)

としてあげることで、python_add を tf.Graph に変換してくれます。すなわち tf_add という関数は tf.Graph で静的に表現された計算グラフを呼び出す操作になっています。これはちょうど、TF1.X系統で言うところの

x = tf.placeholder(dtype=tf.float32)
y = tf.placeholder(dtype=tf.float32)
tf_add = tf.add(x, y)

としてx と y を受け取る tf_add なる計算グラフをtf.Graphに配置することに相当します。ただし、TF2.0からはtf.placeholder は廃止されているため、これはTF1.XとTF2.0とのアナロジーを見るための表現に過ぎないことには注意しておいてください。TF1.X系統ではsess=tf.Session()とセッションを立てて、sess.run(tf_add, {x: 1., y: 2.}) など tf.Graph の処理を実行するのでした。一方でTF2.0では tf.functionにより変換されたグラフは、内部的には tf.Graph であるものの、あたかも通常のPython関数のごとく

z = tf_add(1., 2.)

と値を得ることが出来ます。 また、tf.add(x, y) という関数は、x と y がtf.Tensor であれば x+y と記述することも出来ます（ただし x と y がnumpy.array ならば、それは当然、通常の numpy.array として足し算になることに注意しましょう）。

実際の使い方

仮に python_add と tf_add を別々にPythonで扱いたいという場合には上記のようにpython_add という関数を定義してから、それをtf.functionによって変換するという手続きをとることになります。しかし、多くの場合、より処理の速いtf.Graphである tf_add だけ使えれば良いというケースも多々あり、わざわざ python_add を定義してから変換するのは二度手間になることも多いです。

そのような場合はデコレータを使って、最初からtf.Graph として関数を定義してしまうことが出来ます。

@tf.function
def tf_add(x, y):
    return tf.add(x, y)

TF2.0のサンプルコードなどでもこの表記は頻繁に目にすることになりますので慣れておきましょう。

注意点

tf.function の振る舞いを見るために下記の例を見てみましょう。

@tf.function
def tf_add(x, y):
    return tf.add(x, y)

print(tf_add(1, 2))
# -> <tf.Tensor: id=49, shape=(), dtype=int32, numpy=3>

ここで、tf_add というtf.Graph で表された関数は int32型を返していることに注意してください。これは入力に int型を利用しているから起こったことです。引き続きこの関数に対して下記の処理を実行させましょう。

tf_add(1., 2.)
# -> <tf.Tensor: id=147, shape=(), dtype=int32, numpy=3>

入力に float型を与えても出力が int型になっていることに注目してください。tf.functionで変換された関数は、最初の入力に依存した型処理を行います。逆に関数を定義した直後に float型を与えた場合はどうなるでしょうか。

@tf.function
def tf_add(x, y):
    return tf.add(x, y)

print(tf_add(1., 2.))
# -> <tf.Tensor: id=156, shape=(), dtype=float32, numpy=3.0>

print(tf_add(1, 2))
# -> <tf.Tensor: id=158, shape=(), dtype=float32, numpy=3.0>

という結果になります。最初に float 型を与えた場合には、後に int型を与えても float型として処理されます。これは、関数の再利用によって思わぬ実装のミスを踏む可能性があるので注意が必要です。ただし、事態はそれほど深刻ではなく、仮に int 型を先に与えた場合でも、小数点を含むような float 型には適宜対応してくれるようです（実験する前は勝手にintで切り下げられるかと思いました）。

@tf.function
def tf_add(x, y):
    return tf.add(x, y)

print(tf_add(1, 2))
# -> <tf.Tensor: id=167, shape=(), dtype=int32, numpy=3>

print(tf_add(1.5, 2.))
# -> <tf.Tensor: id=175, shape=(), dtype=float32, numpy=3.5>

これは柔軟性があり、しっかり型を考えて実装するという手続きを省けるとも取れますし、それによって無茶なコードを書いてしまう可能性もあるということです。しっかり意図した型が出てきているかは確認したほうが良いでしょう。

制御構文

TensorFlowには tf.cond や tf.while_loop など、分岐や繰り返しを制御する関数が元々備わっています。TF1.X系統ではこれらを上手に駆使すれば、原理的にはどんなプログラムも書けるというわけなのですが、実際には使いこなしやすいものではなかったと思われます。しかし、このような関数がTFに備わっているおかげで、例えば

@tf.function
def switch_add_sub(x, y, boolian):
    if boolian:
        return x + y
    else:
        return x - y

と分岐のある処理も内部的にtf.condによる計算グラフを構築し、tf.Graphの形で定義することが出来ます。同様に

@tf.function
def train(x, y, dataset):
    for x, y in dataset:
        update(x, y)

のような繰り返し処理も実施可能です（ここでupdate という関数はどこかで定義されているとしてます）。 これは非常に有用で、いま例でみたとおり、繰り返し処理までもが tf.Graph の中に収まっているため、train の内部のループはPythonで実行される代わりに、tf.Graph として実行されます。主にこの処理によってニューラルネットワークの訓練の時間短縮がはかれることとなります。

updateがパラメータの更新をしつつ、そのときの損失 minibatch_loss を返す関数として定義されているとしましょう。すると、学習ごとに minibatch_loss がどうなっているのかを表示したいというふうに思うこともあるでしょう。その場合は tf.print 関数で表示するようにしておけば、これも tf.Graph として表現されます。

@tf.function
def train(x, y, dataset):
    for x, y in dataset:
        minibatch_loss = update(x, y)
        tf.print("loss: {}".format(minibatch_loss))

TO DO

tf.autograph.to_graph との違いを把握。