Tensorflow2.Xの5Stepのネットワーク実装

この記事を書こうと思ったきっかけ

世の中にTensorflow 2.X実装のネットワークが増えてほしいと思ったから。
最近ではTensorflow 2.X実装のネットワークも増えつつあるが、今でもTensorflow 1.X実装のネットワークは少なくない。
PyTorchを使っている人が多いということもあるが、昔から使い続けている、Tensorflow 1.Xを使い続けているという人も少なくないのだろう。
しかしながら、RTX3000シリーズがCUDA10をサポートしないことを追い風にTensorflow 2.Xへの以降が進むことを期待して、Tensorflow 2.Xでネットワークを定義する方法を5Stepに分けて解説する。
たしか公式のドキュメントにもこの5種類の実装方法を全部1ページに解説しているページはなかったはずなので、参考になると嬉しい。
（相当にパフォーマンスが求められない限り、`tf.keras.layers.Layer`を用いた実装は基本的には`tf.keras.Model`で代替可能なので今回は解説していない。）

同じくGoogleがリリースしているJAXが使えそうになりつつある現在、今更Tensorflowかとなる気もするけど既存のリソース上で書く場合は十分使えるので、ぜひともという気がする。

tf.keras.Sequential

最も簡単で最も使う書き方（自分調べ）。

import tensorflow as tf

model = tf.keras.Sequential([
    tf.keras.layers.Dense(100, activation=tf.nn.relu),
    tf.keras.layers.Dense(100, activation=tf.nn.relu),
    tf.keras.layers.Dense(1, activation=tf.nn.relu),
]) # 中間層3層の全結合層のみのネットワーク
print(model(tf.random.uniform(4, 100))) #(4, 1)のTensorが帰る

以上のように引数にlayerのリストを与えることでネットワークの構造を記述する。
学習には`fit`関数, 推論には`predict`関数などを使用することで、細かいTensorflowの仕様をあまり気にせずにscikit learnのインタフェースで使用できる。

基本的に1方通行のネットワーク構造しか記述できないため、Skip connectionは記述できない。
しかしながら、通常の用途であればこれで問題ないことが多い。

tf.keras Model

一番複雑なモデル構築ができる。
このオブジェクトはクラスの規定継承

tf.keras.Modelとtf.keras.layers.Layerの違い

基本的にはtf.keras.Modelよりもtf.keras.layers.Layerはスリムなインタフェースしか持っていない点が大きな違いである。
例えば`fit`, `predict`, `evaluate`などの学習や推論に用いるインタフェースはtf.keras.Modelしか持っていない。
tf.keras.layers.LayerはModel内で多数使用されることが前提とされていて、よりスリムであることが求められるためである。
今回はグラフ構築を動的にすることを前提にしているため、特別`tf.keras.layers.Layer`について解説しないが、静的に出力shapeの計算を行うことで高速なグラフ構築を自動で行ってくれる。

tf.keras.Model+add_loss

誤差関数以外のlossの記述はこの関数を使う。
tf.keras.Model+add_lossと書いたが別に`Sequential`にも使える。
使い方は非常に簡単。`add_loss`の第一引数に追加したい誤差を記述するだけである。

import tensorflow as tf

model = tf.keras.Sequential([
    tf.keras.layers.Dense(100, activation=tf.nn.relu),
    tf.keras.layers.Dense(100, activation=tf.nn.relu),
    tf.keras.layers.Dense(1, activation=tf.nn.relu),
]) # 中間層3層の全結合層のみのネットワーク
model.add_loss(tf.reduce_mean(tf.abs(model.layers[0].weights[0]),axis=-1)) # この場合、1層目の全結合層のy=Wx+bのWに対するL1正則化
model.fit(x, y) # いつもの学習

一応上記はDenseの引数のkernel_regularizerにregularizerを追加すればできることなので、上のような方法でL1正則化をかけるケースはない。
それでも有用なケースはいくらでもあって、例えば以下のようにAutoencoderの潜在変数にL1正則化を書けたい場合だ。

class LatentL1Autoencoder(keras.Model):
    def __init__(self):
        self.encoder = tf.keras.Dense(10)
        self.decoder = tf.keras.Dense(10)

    def call(self, x):
        latent= self.encoder(x, training=True)  # Encoding
        y_pred = self.decoder(latent, training=True) # Decoding
        l1_loss = tf.reduce_mean(tf.abs(latent), axis=-1)
        self.add_loss(l1_loss) # Lossの追加

        return y_pred

model = LatentL1Autoencoder()
model.fit(x, x)

同様にl1_lossの計算を標準正規分布とのkl-divergenceに置き換えればVAEの出来上がりだ。
多分、VAEを書くならこれだけで十分であると言える。

tf.keras.Model+train_step

`train_step`はtf.keras.Modelのメンバ関数でこれをオーバーロードすることで記述する。
これは上記と同じことを記載している。
あまりこれだけの用途であれば意味はないが、2回微分とかより複雑な誤差関数を有するモデルを構築したい場合はこういう書き方が有用だ。

class LatentL1Autoencoder(keras.Model):
    def __init__(self):
        self.encoder = tf.keras.Dense(10)
        self.decoder = tf.keras.Dense(10)

    def call(self, x):
        latent= self.encoder(x, training=True)  # Encoding
        y_pred = self.decoder(latent, training=True) # Decoding

        return y_pred

    def train_step(self, data):
        x, y = data

        with tf.GradientTape() as tape:
            l1_loss = tf.reduce_mean(tf.abs(latent), axis=-1)
            self.add_loss(l1_loss)

            #regularization_lossesにself.lossesを入れないと正則化項が一切含まれなくなるので注意
            loss = self.compiled_loss(y, y_pred, regularization_losses=self.losses) 

        trainable_vars = self.trainable_variables
        gradients = tape.gradient(loss, trainable_vars) # 勾配の計算
        self.optimizer.apply_gradients(zip(gradients, trainable_vars)) # パラメータの更新
        self.compiled_metrics.update_state(y, y_pred) # 評価値の計算
        return {m.name: m.result() for m in self.metrics} # 計算した評価値を返す

model = LatentL1Autoencoder()
model.fit(x, x)

一応これで終わりだが、厳密にはdataにsample_weightが含まれる場合の対応やtest_stepの実装を行わないとevaluation時のlossにl1正則化項が含まれないなどいくらか問題がある。
抜粋してきたコードだがtest_step自体の実装は以下になる。sample_weightを使う場合それとは別にdataの系列長を確認して3ならsample_weightが含まれていて、2なら含まれていないという解釈になる。

Customize what happens in Model.fit  |  TensorFlow Core

python

>

class CustomModel(tf.keras.Model):
def test_step(self, data):
# Unpack the data
x, y = data
# Compute predictions
y_pred = self(x, training=False)
# Updates the metrics tracking the loss
self.compiled_loss(y, y_pred, regularization_losses=self.losses)
# Update the metrics.
self.compiled_metrics.update_state(y, y_pred)
# Return a dict mapping metric names to current value.
# Note that it will include the loss (tracked in self.metrics).
return {m.name: m.result() for m in self.metrics}
|

tf.keras.Model+自前のtraining loop

今や、これを使うのはGANを書くときぐらいだろうか。
`tf.keras.Model`はoptimizerを一つしか使えないが`fit`関数を使わずにtraining loopを書けば複数個使うことができる。
非常に長くなるので、その書き方は是非君の目で確かめてほしい。

Writing a training loop from scratch  |  TensorFlow Core

基本的にはtrain_stepの中身を外に出してその外側にdatasetをfor文で回してtrain_stepの引数に与えるような実装になる。