低レベルAPIユーザーのためのTensorFlow2.0入門

先日(2019/10/1)、ついに[TensorFlow2.0.0が正式リリース][2]されました。当初の予定からはだいぶズレましたが、無事リリースされて何よりです。
一方で、TensorFlow1.15のRC版も公開されていて、[1.15がTF1.x系の最後のリリースとなることが明言][3]されています。
そのため、機能、サポート、stackoverflow等のweb情報の面を考えると、今後はTF2.x系に移行していくべきなのですが、TF2.0はTF1.x系から大きな変更点がいくつかあり、注意が必要です。

今回は、特にTF1.x系で`tf.Session`のような低レベルAPIを使っていた人に向けて、TF2.0で互換となりそうな基本的な書き方についてまとめたいと思います。

# TF1.xからTF2.0の変更点
TF2.0の変更点はいくつかありますが、その中でも書き方に大きく影響しそうな点をピックアップしておきます。

## Eager Executionがデフォルトに
TF2.0の一番大きな変更点は、ChainerやPyTorchのようなDefine by Run形式のEager Executionがデフォルトとなったことです。
TensorFlowは今まで計算グラフ構築→実行を明示的に行うDefine and Runがデフォルトで、最初に計算グラフを構築するおかげで最適化が行えるため、計算速度が出るというメリットがありました[^1]。
Define by Run形式となったことでこのメリットが失われるのではないか、というのが大きな懸念でしたが、TF2.0ではその辺をうまくやっていて、**Define by Run形式で書きつつ計算グラフ構築も行うハイブリッドな書き方が可能**となっています。

## keras APIの統合
TensorFlowとkerasで重複するAPIがいくつかあり、それらが統合されました。
例えば、OptimizerはTF1.x系で`tf.train.Optimizer`と`tf.keras.optimizers.Optimizer`の2種類が存在しましたが、TF2.0では`tf.optimizers.Optimizer`（内部はkerasのもの）に統合されています。
複数存在していたことで少なからず混乱の元となっていたので、この変更は素直に嬉しいですね。

# 基本的な書き方
それでは、基本的な書き方について見ていきます。
全体の流れは
　　1. ネットワークの定義
　　2. 学習ステップとコスト関数の設定
　　3. 学習を実行
となっています。

## ネットワークの定義
TF2.0では低レベルAPIライクな書き方であっても、keras APIを使ったネットワーク定義が推奨されています。
keras APIでは、重みを保持した`Layer`インスタンスをつないで`Model`クラスを作成し、`Model`インスタンスを`call`することでネットワークの出力を得ます。
keras APIを使わず今まで通りのやり方でも書くことができますが、後述のoptimizerによる学習部分の定義で必要な重みのリストを作るのが面倒なので、keras APIを使うのが無難です。

### Functional API
keras APIにもSequentialとFunctionalの2種類がありますが、より柔軟性のあるFunctionalで書くのが個人的にオススメです。

```python
import tensorflow as tf

#=========================================================================
# ネットワーク定義
#=========================================================================

input_shape = (28, 28, 1) # 入力のshape. 最初の次元（バッチサイズ）は除く.

# ネットワークの定義
# 入力層
x = tf.keras.layers.Input(input_shape)
# 畳み込み層1
h = tf.keras.layers.Conv2D(64, (3, 3), padding='SAME')(x)
h = tf.keras.layers.ReLU()(h)
# 畳み込み層2
h = tf.keras.layers.Conv2D(64, (3, 3), padding='SAME')(h)
h = tf.keras.layers.ReLU()(h)
# 畳み込み層3
h = tf.keras.layers.Conv2D(64, (3, 3), padding='SAME')(h)
h = tf.keras.layers.ReLU()(h)
# 線形層
h = tf.keras.layers.Flatten()(h)
y = tf.keras.layers.Dense(10)(h)

# モデルの作成
model = tf.keras.Model(x, y)
```

最初に入力層を定義し、その出力を各`Layer`クラスに通していきます。1つ目のカッコは`Layer`インスタンスの作成、2つ目のカッコは`Layer`インスタンスの`call`です。
重みはこの`Layer`インスタンスの中で保持していて、`Model`インスタンスを作成した時点で初期化されます。
また、重み共有は以下のように簡単に行えます。

```python
# レイヤーインスタンスの作成
layer = tf.keras.layers.Conv2D(64, (3, 3), padding='SAME')

# 同じ重みを使って2つの出力を得る
output1 = layer(input1)
output2 = layer(input2)
```
シンプルなネットワークであればこの書き方で十分だと思います。

### Custom Model
大抵のネットワークであればFunctional APIの`tf.keras.Model`でネットワークを作ることができますが、より複雑なネットワークを作りたい場合には、自分で`Model`を作ることができます。
PyTorchやChainerに一番近い書き方で、馴染みのある人が多いかもしれません。

~~こちらの場合は、`Model`インスタンスを作成した時点では重みが初期化されず、最初に`call`された時か`load_weights`メソッドで重みを復元した時に初期化されます。~~
こちらの場合は、重みによって初期化のタイミングが異なります。
各`Layer`の`__init__`メソッドで作成される重みは、Functional Modelと同様に、`Model`インスタンスを作成した時点で初期化されます。一方で、各`Layer`の`build`メソッドで作成される重みは、`Model`インスタンスを作成した時点では初期化されず、最初に`call`された時か`load_weights`メソッドで重みを復元した時に初期化されます。**(2019/10/7 修正)**

```python
class CustomModel(tf.keras.Model):
    def __init__(self, **kwargs):
        super(CustomModel, self).__init__(**kwargs)

        # Layerインスタンスの作成
        self.conv1 = tf.keras.layers.Conv2D(64, (3, 3), padding='SAME')
        self.act1 = tf.keras.layers.ReLU()
        self.conv2 = tf.keras.layers.Conv2D(64, (3, 3), padding='SAME')
        self.act2 = tf.keras.layers.Activation('sigmoid')

    def call(self, inputs):
        # ネットワークの定義
        # __call__した時に呼ばれる
        h = self.conv1(inputs)
        h = self.act1(h)
        h = self.conv2(h)
        y = self.act2(h)
        return y
```

### Custom Layer
`Model`と同じように`Layer`も自作することができます。最新の手法を試したいけど`Layer`が実装されていない場合に重宝します。

```python

class CustomLayer(tf.keras.layers.Layer):
    def __init__(self, num_features, **kwargs):
        super(CustomLayer, self).__init__(**kwargs) # 必ず最初にsuper().__init__()を呼び出す
        self.num_features = num_features
        
        # add_weightメソッドで重み作成
        # 重みはここで作成するか、buildメソッドで作成する
        self.bias = self.add_weight(
            name='bias',
            shape=(1, self.num_features),
            dtype=tf.float32,
            initializers='zeros')

    def build(self, input_shape):
        # 初めて__call__した時に呼ばれる
        # input_shapeは入力のshape
        # 入力のshapeが自動で計算されて分かるので、入力のshapeに合わせた重みを作成できる
        self.kernel = self.add_weight(
            name='kernel',
            shape=(input_shape[-1], self.num_features),
            dtype=tf.float32,
            initializers='he_normal')

        # 最後に必ずsuper().build()を呼び出す
        super(CustomLayer, self).build()

    def call(self, inputs):
        # このLayerで行う計算を定義する
        # __call__した時に呼ばれる
        # inputsがLayerの入力
        with tf.name_scope('custom_layer'):
            return tf.linalg.matmul(inputs, self.kernel) + self.bias
```

## 学習ステップとコスト関数の設定

optimizerと1回の学習で行う処理を記述した`train_step`関数を定義します。

ここでのポイントは、**`train_step`関数に`tf.function`デコレータをつけるとgraphモード(TF1.x系デフォルト)、つけないとEagerモード(TF2.0デフォルト)で実行される**ことです。
その他の部分で違いはないので、**実質`tf.function` 1行で2つのモードを切り替えることができます**。

また、gradientを計算して重みを更新する際に、更新する重みのリストが必須となりました。
ネットワーク定義に`tf.keras.Model`を使っている場合、`tf.keras.Model.trainable_variables`を呼び出すことで、そのネットワークに含まれる重みをリストとして得られます。
ネットワーク定義にkeras APIをオススメしたのはこれが理由です。

コスト関数の値や精度は関数の返り値としても良いですが、`tf.metrics`の各クラスを使うと、各iterationを平均して...といった自分で計算する部分が減って便利です。（コスト関数の値はまだ良いとして、精度の算出って正直めんどくさいですよね...）

```python

#=========================================================================
# 学習ステップの定義
#=========================================================================

optimizer = tf.optimizers.Adam(1.0e-4)
train_loss = tf.keras.metrics.Mean() # コスト記録用
train_acc = tf.keras.metrics.SparseCategoricalAccuracy() # 精度計算・記録用

@tf.function
def train_step(inputs):
    images, labels = inputs

    # tf.GtadientTapeブロックで入力からロスまで計算
    with tf.GradientTape() as tape:
        logits = model(images)
        loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels, logits)
        
    # gradientを計算
    # gradientを計算する重みをsourcesとして指定することが必須
    # keras.Modelを使っているとmodel.trainable_variablesで渡すことができて便利
    grad = tape.gradient(loss, sources=model.trainable_variables)

    # optimizerで重みを更新
    optimizer.apply_gradients(zip(grad, model.trainable_variables))

    # lossの値を記録
    train_loss.update_state(loss)
    # train_loss(loss) # このように単純に__call__しても良い

    # 精度を記録
    train_acc.update_state(labels, logits)
    # train_acc(labels, logits) # このように単純に__call__しても良い
```

## データセットを与えて実際に学習
準備はほぼ終わりました。TF2.0では、`tf.placeholder`を設定して`sess.run`で`feed_dict={...}`で渡して...といったことは必要ありません。そもそも`tf.Session`自体使いません。
実際にデータセットを与えて学習を行いましょう。

### `tf.data.Dataset` APIを用いる場合
TF1.x系では、`tf.data.Dataset` APIを用いるのがTensorFlowのオススメでしたが、その点はTF2.0でも同じです。
TF2.0では、作成した`Dataset`はミニバッチを返すiterableなオブジェクトなので、for文で取り出すのが楽です。ミニバッチ1つだけ取り出したい場合は`next(iter(dataset))`としても取り出せます。

```python
import time
import numpy as np

#=========================================================================
# データセット
#=========================================================================

# データセットをロード
# 今回はMNIST
(X_train, Y_train), (X_test, Y_test) = tf.keras.datasets.mnist.load_data()
X_train = X_train[..., np.newaxis].astype(np.float32)
Y_train = Y_train.astype(np.int32)
N = X_train.shape[0]

# tf.data.Dataset APIを使う
batch_size = 32
dataset = tf.data.Dataset.from_tensor_slices((X_train, Y_train))
dataset = dataset.shuffle(buffer_size=N)
dataset = dataset.batch(batch_size, drop_remainder=True)

#=========================================================================
# Dataset APIで学習を実行
#=========================================================================

print('train with Dataset API.')

epochs = 100
for epoch in range(epochs):
    time_start = time.time()

    for images, labels in dataset: # 1step分のデータを取り出し
        train_step((images, labels)) # 1step分の学習を実行
    
    # 平均ロスと平均精度
    # Metric.result()メソッドで取り出せる
    epoch_loss = train_loss.result()
    epoch_acc = 100 * train_acc.result()

    # epochの結果を表示
    time_epoch = time.time() - time_start
    print('epoch: {:} loss: {:.4f} acc: {:.2f}% time: {:.2f}s'.format(
        epoch + 1, epoch_loss, epoch_acc, time_epoch))

    # epoch毎にリセットしないと累積していく
    train_loss.reset_states()
    train_acc.reset_states()

# 学習済みの重みを保存
model.save_weights('weights_{:}'.format(epochs))
```

### データセットを直接与える場合

TF1.x系において、`tf.placeholder`を定義して、学習時に`feed_dict={...}`として1step分のデータセットを与えるやり方があったように、TF2.0でも`tf.data.Dataset` APIを用いずに直接与えるやり方もあります。
この場合、ミニバッチのnumpy配列を直接`train_step`関数に与えます。

```python
#=========================================================================
# データセット
#=========================================================================

# データセットをロード
# 今回はMNIST
(X_train, Y_train), (X_test, Y_test) = tf.keras.datasets.mnist.load_data()
X_train = X_train[..., np.newaxis].astype(np.float32)
Y_train = Y_train.astype(np.int32)
N = X_train.shape[0]

#=========================================================================
# numpy配列を直接与えて学習を実行
#=========================================================================

print('train with numpy array.')

epochs = 10
for epoch in range(epochs):
    time_start = time.time()

    # ここから違う========================================================

    indices = np.arange(N)
    np.random.shuffle(indices)
    for i in range(N // batch_size):
        # 1step分のデータを取り出し
        indices_batch = indices[i * batch_size:(i + 1) * batch_size]
        images = X_train[indices_batch]
        labels = Y_train[indices_batch]
        
        # 1step分の学習を実行
        # numpy配列を直接入れてOK
        train_step((images, labels))
    
    # ここまで違う========================================================

    # 平均ロスと平均精度
    # Metric.result()メソッドで取り出せる
    epoch_loss = train_loss.result()
    epoch_acc = 100 * train_acc.result()

    # epochの結果を表示
    time_epoch = time.time() - time_start
    print('epoch: {:} loss: {:.4f} acc: {:.2f}% time: {:.2f}s'.format(
        epoch + 1, epoch_loss, epoch_acc, time_epoch))

    # epoch毎にリセットしないと累積していく
    train_loss.reset_states()
    train_acc.reset_states()

# 学習済みの重みを保存
model.save_weights('weights_{:}'.format(epochs))
```

以上で基本的な書き方は全てです。最後にコードの全文を載せます。[github][1]では、Google Colaboratoryですぐ試せるipynbファイルを置いています。

<details><summary>コード全体はこちら</summary><div>

```python

import time
import numpy as np
import tensorflow as tf

#=========================================================================
# ネットワーク定義
#=========================================================================

input_shape = (28, 28, 1) # 入力のshape. 最初の次元（バッチサイズ）は除く.

# ネットワークの定義
# 入力層
x = tf.keras.layers.Input(input_shape)
# 畳み込み層1
h = tf.keras.layers.Conv2D(64, (3, 3), padding='SAME')(x)
h = tf.keras.layers.ReLU()(h)
# 畳み込み層2
h = tf.keras.layers.Conv2D(64, (3, 3), padding='SAME')(h)
h = tf.keras.layers.ReLU()(h)
# 畳み込み層3
h = tf.keras.layers.Conv2D(64, (3, 3), padding='SAME')(h)
h = tf.keras.layers.ReLU()(h)
# 線形層
h = tf.keras.layers.Flatten()(h)
y = tf.keras.layers.Dense(10)(h)

# モデルの作成
model = tf.keras.Model(x, y)


#=========================================================================
# 学習ステップの定義
#=========================================================================

optimizer = tf.optimizers.Adam(1.0e-4)
train_loss = tf.keras.metrics.Mean() # コスト記録用
train_acc = tf.keras.metrics.SparseCategoricalAccuracy() # 精度計算・記録用

@tf.function
def train_step(inputs):
    images, labels = inputs

    # tf.GtadientTapeブロックで入力からロスまで計算
    with tf.GradientTape() as tape:
        logits = model(images)
        loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels, logits)
        
    # gradientを計算
    # gradientを計算する重みをsourcesとして指定することが必須
    # keras.Modelを使っているとmodel.trainable_variablesで渡すことができて便利
    grad = tape.gradient(loss, sources=model.trainable_variables)

    # optimizerで重みを更新
    optimizer.apply_gradients(zip(grad, model.trainable_variables))

    # lossの値を記録
    train_loss.update_state(loss)
    # train_loss(loss) # このように単純に__call__しても良い

    # 精度を記録
    train_acc.update_state(labels, logits)
    # train_acc(labels, logits) # このように単純に__call__しても良い


import time
import numpy as np

#=========================================================================
# データセット
#=========================================================================

# データセットをロード
# 今回はMNIST
(X_train, Y_train), (X_test, Y_test) = tf.keras.datasets.mnist.load_data()
X_train = X_train[..., np.newaxis].astype(np.float32)
Y_train = Y_train.astype(np.int32)
N = X_train.shape[0]

# tf.data.Dataset APIを使う
batch_size = 32
dataset = tf.data.Dataset.from_tensor_slices((X_train, Y_train))
dataset = dataset.shuffle(buffer_size=N)
dataset = dataset.batch(batch_size, drop_remainder=True)

#=========================================================================
# Dataset APIで学習を実行
#=========================================================================

print('train with Dataset API.')

epochs = 10
for epoch in range(epochs):
    time_start = time.time()

    for images, labels in dataset: # 1step分のデータを取り出し
        train_step((images, labels)) # 1step分の学習を実行

    # 平均ロスと平均精度
    # Metric.result()メソッドで取り出せる
    epoch_loss = train_loss.result()
    epoch_acc = 100 * train_acc.result()

    # epochの結果を表示
    time_epoch = time.time() - time_start
    print('epoch: {:} loss: {:.4f} acc: {:.2f}% time: {:.2f}s'.format(
        epoch + 1, epoch_loss, epoch_acc, time_epoch))

    # epoch毎にリセットしないと累積していく
    train_loss.reset_states()
    train_acc.reset_states()

# 学習済みの重みを保存
model.save_weights('weights_{:}'.format(epochs))

#=========================================================================
# numpy配列を直接与えて学習を実行
#=========================================================================

print('train with numpy array.')

epochs = 10
for epoch in range(epochs):
    time_start = time.time()

    # ここから違う========================================================

    indices = np.arange(N)
    np.random.shuffle(indices)
    for i in range(N // batch_size):
        # 1step分のデータを取り出し
        indices_batch = indices[i * batch_size:(i + 1) * batch_size]
        images = X_train[indices_batch]
        labels = Y_train[indices_batch]
        
        # 1step分の学習を実行
        # numpy配列を直接入れてOK
        train_step((images, labels))
    
    # ここまで違う========================================================
    
    # 平均ロスと平均精度
    # Metric.result()メソッドで取り出せる
    epoch_loss = train_loss.result()
    epoch_acc = 100 * train_acc.result()

    # epochの結果を表示
    time_epoch = time.time() - time_start
    print('epoch: {:} loss: {:.4f} acc: {:.2f}% time: {:.2f}s'.format(
        epoch + 1, epoch_loss, epoch_acc, time_epoch))

    # epoch毎にリセットしないと累積していく
    train_loss.reset_states()
    train_acc.reset_states()

# 学習済みの重みを保存
model.save_weights('weights_{:}'.format(epochs))
```

</div></details>

# まとめ
TF1.x系で低レベルAPIを使ってネットワークを書いていた人に向けて、TF2.0において互換になりそうな基本的な書き方を見ていきました。

TF1.x系で必要だった`tf.placeholder`や`tf.Session`、変数の初期化等が必要なくなり、コード量が少なくなってスッキリした印象です。また、実行速度の面で優位な今までのgraphモードに、実質1行の変更で切り替えられるのは嬉しいですね。

TF2.0になってどういう書き方になるのか、個人的に大きな不安がありましたが、難解な部分は全くなく、とても書きやすかったです。1.x系ではとっつきにくさがあり、他人にもあまり積極的にオススメできませんでしたが、2.0では自信を持ってオススメできると思います。

Google Colaboratoryで実行できるipynbファイルを[github][1]に置いているので、気軽に試してみてください。

### 追記

[2019/10/7] Custom Modelにおける重みの初期化タイミングについて、間違いがあったため修正しました。
[2019/10/8] Define by RunとDefine and Runが入れ替わっている部分があったため修正しました。

[1]: https://github.com/mgmk2/TensorFlow_v2_MNIST_sample
[2]: https://github.com/tensorflow/tensorflow/releases/tag/v2.0.0
[3]: https://github.com/tensorflow/tensorflow/releases/tag/v1.15.0-rc2

[^1]: Define by Run形式のPyTorchがTensorFlowと同等以上の実行速度が出ているというテスト結果もあります。