TensorFlow2.0を使ってMNISTをMLPで学習する

# モチベーション
* TensorFlow 2.0を理解したい
* [ディープラーニング基礎講座](https://deeplearning.jp/lectures/deep-learning-foundation/)で学んだことのおさらいがしたい

# 当面の取り組み予定
ディープラーニング基礎講座の課題をTensorFlow 2.0で書き直す。

# 今回の課題
* MNISTデータセットを多層パーセプトロン（MLP）で学習する
    * 制約：全体の実行時間は60分以内
    * 元々の課題は「TensorFlowを使わずに」だったが、今回はTensorFlowの勉強を兼ねているため利用する
    * 元々の課題は精度も求められていたが、今回はTensorFlowの勉強に重きを置き、精度向上には取り組まない

# 参考
* [Get started with TensorFlow 2.0 for experts](https://www.tensorflow.org/alpha/tutorials/quickstart/advanced)

# 環境
* Google Colaboratory
* TensorFlow 2.0 Alpha

# コード
## TensorFlowのインストール
デフォルトではGoogle ColabにTensorFlow 2.0はインストールされていないため、自分でインストールします。今回はGPUを利用するためGPU版です。なお、初めにこのインストールを忘れ、`import tensorflow as tf`としていると、TensorFlowインストール後にランタイムの再起動が必要になりますので、忘れずに初めにインストールしておきましょう。

```Python
!pip install tensorflow-gpu==2.0.0-alpha0
```

## MNISTのダウンロードと加工
チュートリアルに載っている`tensorflow_datasets`を利用します。`as_supervised`オプションを付けると戻り値がタプルになって利用しやすそうだったため設定しています（デフォルトではディクショナリーになる）。

```Python
import tensorflow_datasets as tfds
import tensorflow as tf

dataset = tfds.load('mnist', as_supervised=True)
mnist_train, mnist_test = dataset['train'], dataset['test']
```

画像データの正規化、シャッフル、バッチ化を行います。`shuffle`の引数`buffer_size`はチュートリアル通り10,000にしていますが、訓練データのデータ数は60,000あります。データ数以上の値を設定していると確実にシャッフルされるとどこかに書いてありましたが、訓練データより少ないとどうなるのかはよくわかっていません。

```Python
def convert_types(image, label):
    image = tf.cast(image, tf.float32)
    image /= 255
    return image, label

mnist_train = mnist_train.map(convert_types).shuffle(10000).batch(32)
mnist_test = mnist_test.map(convert_types).batch(32)
```

## モデルの定義
こちらもチュートリアル通りKerasを使っています。Kerasを使わないモデル定義の方法はそのうち勉強します。あとこの場合、メソッド名を`call`とする必要があるようです（後のコードで自分では`call`を読んでいないので）。Kerasの学習はfitだと思っていましたが・・・。これもそのうちわかるでしょう。

```Python
from tensorflow.keras.layers import Dense, Flatten, Conv2D
from tensorflow.keras import Model

class MyModel(Model):
    def __init__(self):
        super(MyModel, self).__init__()
        self.flatten = Flatten() # (28, 28, 1) -> 784
        self.d1 = Dense(128, activation = 'relu') # 784 -> 128
        self.d2 = Dense(10, activation = 'softmax') # 128 -> 10
        
    def call(self, x):
        x = self.flatten(x)
        x = self.d1(x)
        return self.d2(x)
    
model = MyModel()
```

こちらもKerasを使っています。`SparseCategoricalCrossentropy()`はこれまで見たことがなかったのですが、ラベルがone hotエンコーディングされていない場合に使うといいのだそうです。今回のラベルはone hotエンコーディングされていませんので、`SparseCategoricalCrossentropy()`をそのまま使うことにしました（試しにone hotエンコーディングせず`CategoricalCrossentropy()`を使ってみましたがエラーになりました）。

```Python
loss_object = tf.keras.losses.SparseCategoricalCrossentropy()
optimizer = tf.keras.optimizers.Adam()
```

## 訓練の準備
ログ用のLossとAccuracyの定義です。

```Python
train_loss = tf.keras.metrics.Mean(name = 'train_loss')
train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name = 'train_accuracy')

test_loss = tf.keras.metrics.Mean(name = 'test_loss')
test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name = 'test_accuracy')
```

`@tf.function`を付けると、その関数はグラフにコンパイルされ速く動作するそうです。詳細は[こちら](tf.function and AutoGraph in TensorFlow 2.0)をご覧下さい。試しに時間を測定してみました。

* `@tf.function`あり：1.74分
* `@tf.function`なし：4.29分

なので、結構違いますね！基本的には`@tf.function`を付けるようにしておいた方がよさそうです。

`with tf.GradientTape() as tape`の中には微分を行う対象式を記述します。そして、その後の`tape.gradient`により自動的に微分してくれるようです。

```Python
@tf.function
def train_step(image, label):
    with tf.GradientTape() as tape:
        predictions = model(image)
        loss = loss_object(label, predictions)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    
    train_loss(loss)
    train_accuracy(label, predictions)

@tf.function
def test_step(image, label):
    predictions = model(image)
    t_loss = loss_object(label, predictions)
    
    test_loss(t_loss)
    test_accuracy(label, predictions)
```

## 学習
最後に学習です。このコードを簡単にするためにデータロード時に`as_supervised=True`としていました。

```Python
EPOCHS = 5

for epoch in range(EPOCHS):
    for image, label in mnist_train:
        train_step(image, label)
        
    for test_image, test_label in mnist_test:
        test_step(test_image, test_label)
        
    template = 'Epoch {}, Loss: {}, Accuracy: {}, Test Loss: {}, Test Accuracy: {}'
    print(template.format(epoch + 1, train_loss.result(), train_accuracy.result() * 100, test_loss.result(), test_accuracy.result() * 100))
```

結果はこちらの通りです。

```Python
Epoch 1, Loss: 0.0035276836715638638, Accuracy: 99.88829803466797, Test Loss: 0.1231490969657898, Test Accuracy: 97.93000030517578
Epoch 2, Loss: 0.0034228067379444838, Accuracy: 99.892578125, Test Loss: 0.12411412596702576, Test Accuracy: 97.92285919189453
Epoch 3, Loss: 0.0034678690135478973, Accuracy: 99.89083099365234, Test Loss: 0.12539470195770264, Test Accuracy: 97.92533874511719
Epoch 4, Loss: 0.00351154338568449, Accuracy: 99.88870239257812, Test Loss: 0.12521344423294067, Test Accuracy: 97.94062042236328
Epoch 5, Loss: 0.0034353630617260933, Accuracy: 99.89129638671875, Test Loss: 0.12637655436992645, Test Accuracy: 97.93353271484375
```

# 所感
* Placeholderや`tf.Session()`がなくなって、より直感的に書けるようになった
* `tf.keras`の出番が多い
* `tensorflow_datasets`や`tf.GradientTape()`などTensorFlow 2.0以前の関数の使い方が少しわかってよかった

# GitHubのリポジトリ
[こちら](https://github.com/shoji9x9/MNIST-By-MLP)です。