TensorFlow2:MNISTやってみる（keras model class）

TensorFlowのチュートリアルにそって、MNIST問題をやってみる。

TensorFlow 2 quickstart for experts
上記のコードをベースに、関数作ったりしながらソースを作成した。
1.xの時との違いとかをメモしながら記載する。といいつつ、1.xの時に近い書き方を選んでいると思う。

ざっくり処理フロー

以下のような流れになる。流れは1.xの時と変わらない。

• MNISTデータ読み出し処理の定義
• モデルの定義
• 学習の定義
• テストの定義
• 学習ループの実行

コード

事前に

タイピングを減らすために、パッケージ名にエイリアスを付けておく。

import tensorflow as tf

# asign aliases
tfk = tf.keras
tfkl = tf.keras.layers

MNISTデータ読み出し処理の定義

def load_dataset(batch_size=32):
    (x_train, y_train), (x_test, y_test) = tfk.datasets.mnist.load_data()
    # normalize 0.0 ~ 1.0
    x_train, x_test = x_train / 255.0, x_test / 255.0
    # add newaxis for channel.
    x_train = x_train[..., tf.newaxis]
    x_test = x_test[..., tf.newaxis]
    # build pipeline: shuffle -> batch.
    ds_train = tf.data.Dataset.from_tensor_slices(
        (x_train, y_train)).shuffle(10000).batch(batch_size)
    ds_test = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(batch_size)
    # return dataset.
    return ds_train, ds_test

でました。tf.data。
昔はよくわからなかったので、結局学習ループ内で、ファイルから読み出し＆前処理やってたんですよね。
ただ、これだとデータ読み出し時に待ちが発生して遅いだとか。

シャッフルやバッチ取り出しをつないで書けるのは便利そう。
このサンプルだと一括してメモリにデータをロードしているので、でかいデータを逐次読み出す場合にどうなるかは、また調べないと。

モデルの定義

class MyModel(tfk.Model):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = tfkl.Conv2D(32, 3, padding="same", activation="relu", use_bias=True,
                                 name="conv1")
        self.pool1 = tfkl.MaxPool2D(name="pool1")
        self.conv2 = tfkl.Conv2D(64, 3, padding="same", activation="relu", use_bias=True,
                                 name="conv2")
        self.pool2 = tfkl.MaxPool2D(name="pool2")
        self.flatten = tfkl.Flatten()
        self.d1 = tfkl.Dense(128, activation="relu", name="fc1")
        self.d2 = tfkl.Dense(10, activation="softmax", name="softmax")

    def call(self, inputs, training=False):
        x = self.conv1(inputs)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.pool2(x)
        x = self.flatten(x)
        x = self.d1(x)
        return self.d2(x)

層を足したりbiasを足したりしてるけど、誤差の範囲。

1.xでは、下のように、各opごとにweightを生成して、演算のチェインを記述していたけれど、

# 1.x
W_conv1 = weight_variable([5, 5, 1, 32], "conv1_w")
b_conv1 = bias_variable([32], "conv1_b")
h_conv1 = tf.nn.relu(conv2d(x_image, W_conv1) + b_conv1)
h_pool1 = max_pool_2x2(h_conv1)

weightの生成やチャネル数の計算が演算に内包され、記述量が減った。
これはTensorFlow2というより、Kerasの恩恵か。

__init__() にて層を定義して、call() で順方向計算を記述する。
1.xでは build_model(inputs) みたいなオレオレ関数を作っていたところが、classにまとまった感じ。

学習の定義

    model = MyModel()
    loss_obj = tf.keras.losses.SparseCategoricalCrossentropy()
    optimizer = tf.keras.optimizers.Adam()

使用するモデル、損失、optimizerを定義。
1.xでは、モデルの出力値を食わせて演算をつないでいたが、これは後でやる。

    metr_train_loss = tf.keras.metrics.Mean(name='train_loss')
    metr_train_acc = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')

値を蓄積して、後で集計結果を返してくれる便利変数を定義しておく。

@tf.function
def step_train(model, optimizer, loss_obj, images, labels, metr_loss, metr_acc):
    with tf.GradientTape() as tape:
        predictions = model(images, training=True)
        loss = loss_obj(labels, predictions)
        gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    metr_loss(loss)
    metr_acc(labels, predictions)
    return loss

学習の1ステップに相当する関数を定義する。

@tf.functionのアノテーションを付けておくと、処理速度向上が見込める。
1.xの時のグラフのように、自動で事前にグラフにコンパイルしてくれる。

GradientTapeは、スコープ内の順方向計算を記録しておいて、gradient()で誤差を計算してくれる。
apply_gradients()で、さっき計算した誤差を反映する。

計算・反映対象の変数は、keras.Modelが集約してくれている。
keras.Model内でレイヤを定義すると、内部で作成した変数をトラッキングしているようだ。

Backpropに関しては、1.xの時にOptimizerの中でやってた処理を手動でやっている感じ。
昔もcompute_gradients()とapply_gradients()で同じことができた、はず。

あとは、便利変数に今回のlossやaccuracyを追加しておく。あとで結果は取得する。

この関数を、学習ループで呼び出すことになる。

テストの定義

metr_test_acc = tf.keras.metrics.SparseCategoricalAccuracy(name='test_accuracy')

値を蓄積して（略）

@tf.function
def step_test(model, images, labels, metr_acc):
    predictions = model(images, training=False)
    metr_acc(labels, predictions)

モデル出力を評価するのみ。

学習ループの実行

    for epoch in range(1, epochs+1):
        # do training.
        ts_train = datetime.datetime.now()
        for i, (images, labels) in enumerate(ds_train, 1):
            loss_step = step_train(model, optimizer, loss_obj, images, labels, metr_train_loss, metr_train_acc)  # type: ignore # noqa
            if i % 200 == 0:
                print("step {:8d}, loss: {:.3f}".format(i, loss_step))
        ts_train = datetime.datetime.now() - ts_train
        # do test.
        ts_test = datetime.datetime.now()
        for test_images, test_labels in ds_test:
            step_test(model, test_images, test_labels, metr_test_acc)  # type: ignore # noqa
        ts_test = datetime.datetime.now() - ts_test
        # show stats.
        template = "Epoch {} | " \
            "Train Loss: {:.3f}, Acc: {:.2f}, Test Acc: {:.2f}, " \
            "spend {}, {}"
        print(template.format(epoch, metr_train_loss.result(),
                              metr_train_acc.result()*100, metr_test_acc.result()*100,
                              ts_train, ts_test))

1.xでやっていた、placeholderの準備やSessionがなくなって素直なコードになっている。

tf.data を使うことで、通常のiterableのようにforループに渡せる。

取り出したデータ、事前に定義しといたmodel, optimizer, metricsを使って学習処理の関数を呼び出す。
学習後の評価でも同様。

metricsは1epochごとにリセットするので、reset_states()を呼び出しておく。

まとめ

• kerasはブラックボックスが多い印象で敬遠してたけど、API選べばあまり違和感なく書けそう。
• 昔は学習処理など、演算の終端のOpをSession.runしていたが、終端のOpではなく処理自体を実行する感じ。
• Eagerの導入で、昔のSessionのブロックが消えたので、どこがグラフになってるのかは不明瞭になったかな。
• 1.xに近い書き方をしてみたが、全体的に記述量は減った半面、好みは分かれそう。

注意点

前述のコードは、スクリプトの要所要所を切り貼りしたので、全部並べてコピーしても動かないです。
うまいこと並び替えて、load_dataset()を呼べば動くはず。

あと、同じスクリプト内で学習をもう一回呼ぶと、二回目の実行で落ちます。
昔はSessionで区切って変数類を生成していたが、Eager導入によって暗黙的に呼ばれるようになったためと思う。
この辺は、もうちょっと試行錯誤したので別ポストで。