【TF Advent Calendar 2018】TensorFlow Eager Executionの使い方 step by step

はじめに

この記事はTensorFlow Advent Calendar 2018 12/4の記事です。他の日付の記事は下記を参照ください。

qiita.com

今回、12/4の8:00時点で12/4に空きがあったので急遽記事を投入することにしました。このような背景もありまして、基本的に本記事の内容は、これまで自分が触ってきたTensorFlow Eager Executionの使用感を踏まえた極普通のことしか書けませんでした。申し訳ありません。

TensorFlow Eager Execution概要

What's Eager Execution

まず、TensorFlowと言えば

１．静的グラフで実現したいデータの流れと処理を定義して ２．実際にデータを流す

「Define and Run」というパラダイムを有していました。

上記の２つ目のセル print(z) の出力はx+yの計算結果ではなく、addという計算が定義されていますとしか出ない。３つ目のセルでsess.run(z)を実行することで、定義された計算が実際に実行される。

この方式では、どのような計算処理がどのような順番で行われるのかを予め宣言しておくことで、プログラム側がどの変数がどのような計算をなされたのか、その履歴を把握することができます（履歴も何もグラフとして宣言されている）。なぜ計算履歴を把握できると嬉しいのかというと、ニューラルネットワークの学習の要である後方自動微分（バックプロパゲーション）を「微分の連鎖規則」を用いて実装できるからです（もしも、少ない計算量で、複雑な合成関数の微分値を求める方法が他にあればこんなことはしなくて良いはず。もちろんそんな方法を見つけられたら凄まじい論文になるだろうと思う）。

さて、今回紹介するEager Executionでは

１．データの処理を実際に行う事によって、データの流れが定義される

「Define by Run」というパラダイムを実現しています。

tf.enable_eager_execution()によってEagerの利用を宣言します。Define and Runでは計算処理の定義を行っていただけのコードが、Eagerでは直ちに計算処理が実行されているのがprint(z)の出力からわかります。

微分の計算を行うために計算の履歴をしっかり保存することが必須要件でした。しかし、こちらのパラダイムでは「計算処理を実行したら、その都度その計算履歴を（捨て去るまで）保存する」という方式になります（それがどのように実装されているかはさておいて…）。この方式のおかげで、学習の途中でニューラルネットワークの構造を変更したりするような場合にも対応できるようになりました（後述の通り、上記の記述は計算履歴を覚えてない。計算履歴を覚えさせる場合には特別な宣言が必要となる）。

Why Eager Execution

なぜこのタイミングでEager Executionなるものを世の中に送り出すことになったのでしょうか。Googleの立場になってこたえると「PyTorchがすごい伸びてるので焦って世論に合わせた」というのが正直なところかと思います（個人の感想です）。こんなことは誰も聞いていないと思いますので、ユーザーの立場になって「なぜEager Executionを使いたいと思うのか（何が嬉しいのか）」について書きます。

計算グラフを動的に変えられる

まず「Define by Run」というパラダイムによって、既に述べた通り、学習の途中でニューラルネットワークの構造を変更したりすることが容易になりました。また、Python上で扱う変数と、機械学習モデルとしての変数が紐付いているため、同じパラメータを何度も繰り返し使うような構造のニューラルネットワーク（RNNなど）も書きやすくなります。また、「ある条件のときには畳み込み層を通し、他の条件のときには全結合層を通す」であるとかをPythonのif文で簡単に書けます。

上記は入力 x の正負に応じて処理が変わるモデルである。Pythonのif文で簡単に処理内容を分岐させられる。もちろんニューラルネットワークの計算処理の中身に if 文を入れることも可能である（これまでは if 文に相当する条件分岐そのものを「計算グラフとして定義する」必要があった）。

PythonicかつNumPyライクに使える

正直、計算グラフを動的に変えなければならないようなケースは個人的にはかなり限定的だと思っています。実際、多くのディープラーニングユーザーはDefine and RunであるTensorFlowをラッピングしたKerasを好んで利用していますし、それで事足りているケースがほとんどではないでしょうか。すなわち「Define by Runじゃないと絶対にできない」というようなものと向き合わなければならない状況の人に向けて、今回Eager Executionが導入されたというわけではないのです（個人の感想です）。

なぜEager Executionを導入したのか（あるいはユーザーが使いたがると考えたのか）というと、間違いなくPythonユーザーが慣れ親しんだ構文で、NumPyライクに使えるという点が大きいと思われます。 TFに集まる不平の多くは、「扱うモデルが複雑になっていくにつれてデバッグが極めて困難になり、Define and Runというパラダイムによって個々の計算処理をその都度評価して、実装の確認をすることが難しい」という点にあるのではないかと思われます（生TFを書いたことがある人ならば、自分が思った通りの計算処理を実装できているか不安になり、更にその確認にとても苦労したはずだ）。

結局慣れ親しんでいる方法であるというのが大きい

Eager Execution 実践

自動微分

既に先に述べましたが、ニューラルネットワークの学習の要である自動微分をEagerではどのように使うのかを確認しましょう。「先程Eagerでは計算される度に計算履歴を構築・保持する」と言いましたが半分ウソです。実際には「これから行う計算は履歴を保持してください」と伝えないと保持しません。そのためには tf.GradientTape()というものを利用します。

tf.GradientTape()を利用することで計算履歴を保持することができる。上記のようにコンテキストを用いることで、コンテキスト内の計算処理のみを見張るようにプログラムに命令できる。tape.watch(x) は x の微分に関して計算できるように準備しておけということである。

x について見張らせて置きつつ a についての微分を要求しても None が返ってきます。

PyTorchではtorch.Tensor 自体に計算履歴を保持するような仕組みが備わっていますが、TensorFlow Eagerはそうではないことに注意しましょう（これはどちらの方法でもメリットデメリットがあります。PyTorchの方式は微分周りの処理をtorch.Tensor自体にメソッドとして持たせることができ、実装が多くの場面で楽ですが、意図しない計算履歴の保持に拠るメモリの消費が起こる可能性があります）。

実際のところ、上記のようにwith tf.GradientTape() as tape コンテキスト内部で、どの変数による微分を後々計算するのかを考えながら実装を行うのはまあまあ面倒なところです。実際のところは tape.watch()を使わずとも後々、どの変数で微分計算を行うのかを一括で指示する方法があります。

tf.GradientTape() は tf.Variable型の変数に関する微分はいつでも行えるように準備してくれます。機械学習においては「学習パラメータについての微分」を求めるのが大半の使い方なので、学習パラメータをtf.Variableで与えるようにすればOKということです。

低レベルEager

さて、ここまでの知識で簡単な線形モデルを実装してみましょう。線形モデルとは下記の形で書けるモデルを言います。

$$ f(W, b, x) = Wx + b $$

ここで $x$ は入力データです。パラメータが $W$ と $b$ になります。学習したいパラメータを tf.Variableで宣言すればよかったので、上記のモデルの実装は下記のようになります。

class Model(object):
    def __init__(self):
        self.W = tf.Variable(5.0)
        self.b = tf.Variable(0.0)
    
    def __call__(self, x):
        return self.W * x + self.b

クラスを使う必然性は特に無いのですが、後々の書き方に合わせてこうしておきましょう。線形モデルを使った回帰では、手元のデータを $D = \{(x_1, y_1), \cdots, (x_N, y_N)\}$ として、例えば下記の損失関数を最小化します（最小二乗法）。

$$ {\rm Loss} (W, b) = \sum_{i=1}^N (y_i - f(W, b, x_i))^2 $$

ここで $\rm Loss$が$W, b$だけの関数になっているのは、学習データを外から「定数として」与えているからです。別に$x$の関数とみなしても構いませんが、基本的に $x$ に関する微分を求めたりすることは無いので、着目しているパラメータだけに集中することにしましょう。

今損失関数に関して下記のように記述しておきましょう。

def loss(model, x, desired_y):
    predicted_y = model(x)
    return tf.reduce_mean(tf.square(predicted_y - desired_y))

学習では損失関数のパラメータ$W$と$b$による微分を利用することになりますので、Eagerを使う場合には、この時点で「loss関数を呼び出すときに tf.GradientTapeコンテキストで計算を見張らせればイイな」と思えるようにならなければいけません。具体的なパラメータの更新の仕方は

$$ \begin{align} W &\leftarrow W - \epsilon \frac{\partial {\rm Loss}(W, b)}{\partial W} \\ b &\leftarrow b - \epsilon \frac{\partial {\rm Loss}(W, b)}{\partial b} \end{align} $$

となるので、これを実装しましょう。

def train(model, inputs, outputs, learning_rate):
    with tf.GradientTape() as tape:
        current_loss = loss(model, inputs, outputs)
    # 2つ目の引数をリストで渡すと、それぞれの変数での微分を返してくれる
    dW, db = tape.gradient(current_loss, [model.W, model.b])
    
    # model.W = model.W - learning_rate * dW
    # model.b = model.b - learning_rate * db
    model.W.assign_sub(learning_rate * dW)
    model.b.assign_sub(learning_rate * db)

とすれば完了です。あとは学習ループを回すだけです。

epochs = range(10)
for epoch in epochs:
    current_loss = loss(model, inputs, outputs)

    train(model, inputs, outputs, learning_rate=0.1)

学習前

学習後

もちろん実装の方法はこの限りではありません。何をpythonの関数で囲っておくかは自由ですし、ミニバッチ学習に対応したければ今回のtrain関数を改造する必要があるでしょう。

Optimizer利用

上記の記述方法は、TensorFlowの自動微分機能以外は、基本的に数式をスクラッチ気味で実装していました。機械学習の多くの場面では、同様の最適化手法を使いまわすことが多いため、当然のことながら既存の実装が存在しています。モデルと損失関数を書くところまでは先程と同様としておきます。

class Model(object):
    def __init__(self):
        self.W = tf.Variable(5.0)
        self.b = tf.Variable(0.0)
    
    def __call__(self, x):
        return self.W * x + self.b

def loss(model, x, desired_y):
    predicted_y = model(x)
    return tf.reduce_mean(tf.square(predicted_y - desired_y))

先程はパラメータの更新式をdef train()関数として自分で書き下しましたが、tf.train.モジュールを使えば基本的な更新方法を流用することが可能です。典型的にはoptimizerという変数名にして

optimizer = tf.train.GradientDescentOptimizer(learning_rate=1e-4)

などと宣言します。このクラスのメソッドapply_gradientsに現在のパラメータのリストと現在のパラメータにおける勾配のリストをzipして渡してやれば更新を行ってくれます。よってdef train()関数は下記のように書き換えられます。

optimizer = tf.train.GradientDescentOptimizer(learning_rate=1e-4)
def train(model, inputs, outputs, learning_rate):

    with tf.GradientTape() as tape:
        current_loss = loss(model, inputs, outputs)
    # 2つ目の引数をリストで渡すと、それぞれの変数での微分を返してくれる
    dW, db = tape.gradient(current_loss, [model.W, model.b])

    # 現在のパラメータのリストとその勾配のリストを（indexを対応付けて）渡す
    optimizer.apply_gradients(zip([model.W, model.b], [dW, db]))

あとは同様に学習ループを回してやれば良いです。

tf.keras.layersl + tf.train

次は、極めて基本的な使い方となるであろう方法を紹介します。 先程まではモデルとしてtf.Variable剥き出しで実装していました。また計算処理としてtf.matmulなどの低レイヤー関数も利用していました。行列演算であるとか、畳み込み計算であるとか、元々よく使われるパラメータと計算の組をセットにして実装されているtf.keras.layersモジュールを活用します。

model = tf.keras.layers.Dense(output_size)

このコードでは出力の次元がoutput_sizeとなる線形変換を定義できます（入力のサイズはまだ未定で、データが初めて入力されるタイミングで決定される）。更にtf.kerasのモジュールを使って損失関数も書いておきましょう。

def loss_fn(model, x, desired_y):
    predict_y = model(x)
    return tf.keras.losses.mean_squared_error(predict_y, desired_y)

あとの学習コードは先程と変わらず、自前で更新式を書いてもいいですしtf.trainモジュールでoptimizerを定義しても構いません。先程と同様にひとまず

optimizer = tf.train.GradientDescentOptimizer(learning_rate=1e-4)

と書いておき、学習コードについて今回は更に楽をする方法を紹介します。

先程はtf.GradientTapeコンテキストで微分計算を見張らせていましたが、tf.contrib.eagerにはtf.keras.layersで定義されたモデルのパラメータについて自動で見張って、損失と勾配を計算してくれる関数tf.contrib.eager.implicit_value_and_gradientsが実装されています。

value_and_grads = tf.contrib.eager.implicit_value_and_gradients(loss_fn)

ここでtf.contrib.eager.implicit_value_and_gradientsには自分で定義したPython関数であるloss_fnを渡していることに注意してください。こいつは自分で記述したloss_fn関数のラッパーであり、引数はそのまま受け継ぎます。 何が嬉しいのかというと、tf.contrib.eager.implicit_value_and_gradientsは渡してあげたloss_fn関数内での計算において「tf.keras.layers内のパラメータ」を認識し、それに対応する勾配をちゃんと自動で取ってきてくれるところです。こいつに元々のloss_fnの引数を渡してやると、元々定義されていたloss_fnの返り値（損失の値）と「現在のパラメータの値と勾配のリストをzipしたもの」を返してくれます。

従って学習コードとしては下記のような形になります。

for epoch in range(num_epochs):
    loss, grads = value_and_grads(model=model, x=x_train_, y=y_train_)    
    optimizer.apply_gradients(grads)

かなりシンプルになりましたね。Keras APIとEagerのAPIを活用すると、だいぶコードが抽象化されました。

Keras での学習コード抽象化

毎度、似たような学習コードを書くのは面倒です。オリジナルの更新式や、オリジナルの損失関数を使わない限りは上記のような記述をすることは実際には必要なく、tf.keras.modelsモジュールで更に楽をできます。tf.keras.Sequentialとtf.keras.Modelの２種類があるので見ていきましょう。

class Model(tf.keras.Model):
    def __init__(self):
        super(Model, self).__init__(name='mnist_mlp')
        self.dense1 = tf.keras.layers(512, activation='relu')
        self.dense2 = tf.keras.layers(256, activation='relu')
        self.dense3 = tf.keras.layers(10, activation=None)

    def call(self, x):
        x = self.dense1(x)
        x = self.dense2(x)
        x = self.dense3(x)
        return x

model = Model()

あるいは

model = tf.keras.Sequential([
    tf.keras.layers(512, activation='relu'),
    tf.keras.layers(256, activation='relu'), 
    tf.keras.layers(10, activation=None)
])

学習中にモデルの構造をあれこれ変えたり、あるいはGANのように交互に更新が必要であったりしない場合は、典型的な学習の方法がtf.keras.Model.fit()に提供されているため、下記のコードで利用することができます。

optimizer = tf.train.AdamOptimizer(learning_rate)
model.compile(optimizer=optimizer,
              loss='categorical_crossentropy',
              metrics=["accuracy"])

model.fit(x=x_train_eager.numpy(), 
          y=y_train_eager.numpy(), 
          validation_split=0.2, 
          epochs=num_epochs,
          batch_size=batch_size)

Define by Runでなければならないモデルはさほど多くなく、モデルを構築している途中でのデバッグが楽というのがEagerの大きな特徴でもあります。複雑なモデルでも学習方法は共通であるケースが多いので、Kerasの高レベルAPIが使えるときは積極的に使っていきましょう。

最後に

TensorFlow Eagerは使いやすい反面、Static Graphに比べて処理速度は遅いと言わざるを得ません。Pythonは決して処理の速い言語ではなく、PyTorchも含めなるべく処理をCやC++等のより低レベルな言語に任せるのが主流になっています。今後はTensorFlow Eagerで人間に優しいAPIが整備されながらも、裏では高速処理を実現するためのStatic Graphとの互換性に関しても開発が進んでいくはずです（tf.autographについて触れたかったのですが、勉強する時間が取れませんでした…）。

いずれにしても今年でTensorFlow 2.0として大幅に変動を遂げるはずですので見守っていきましょう。

TensorFlow Eager Tutorial

下記のレポジトリにTensorFlow Eagerでの基本的なモデルの実装を置いておきます（随時更新）。

github.com

また公式より tf.kerasの利用を勧めるアナウンスが入りました。

medium.com

www.hellocybernetics.tech