はじめに
Google が開発している 深層学習ライブラリである TensorFlow はモデルの構築や訓練ループを様々な書き方で実現できます。これは有識者にとっては便利ですが、初学者にとっては理解を妨げる要因になり得ます。
今回は TensorFlow 2.X で推奨されている書き方を網羅的に紹介し、画像認識の分野で著名なモデルである VGG16 と ResNet50 を実装しながら使い方を解説します。
※ TensorFlow 2.X とはメジャーバージョンが 2 以上の TensorFlow を指すものとします。
対象読者
- TensorFlow のチュートリアルを試したが、自力でモデルを構築できない人
- TensorFlow で書かれたソースコードを読むと、見慣れない書き方があると感じる人
- Chainer、PyTorch でモデルを書くことができるが TensorFlow でモデルを書けない人
流れ
まず、 TensorFlow の 4 つのモデル構築 API について説明します。
その後、 2 つの訓練手法について説明します。
最後にこれらの書き方を利用して VGG16 と ResNet50 の実装を行います。
検証環境
- macOS Catalina 10.15.3
- Python 3.7.7
- tensorflow 2.2.0
>>> import sys
>>> sys.version
'3.7.7 (default, Mar 10 2020, 15:43:33) \n[Clang 11.0.0 (clang-1100.0.33.17)]'
pip list | grep tensorflow
tensorflow 2.2.0
tensorflow-estimator 2.2.0
TensorFlow における 4 つのモデル構築 API
TensorFlow にはモデルを構築するために大きく分けて 2 つ、細かく分けると 4 つの API が用意されています。
- シンボリック（宣言型）API
- Sequential API
- Functional API
- Primitive API（1.X 系の書き方。非推奨）
- 命令型（モデル サブクラス化）API
- Subclassing API
まず大きな分類について簡単に紹介します。
シンボリック（宣言型） API
モデルの形状を学習実行前に宣言（≒コンパイル）する書き方です。
この API で書かれたモデルは学習中に形状を変更できません。
そのため、一部の動的に形が変わるモデル（Tree-RNN など）は実装できません。
その代わりに、データをモデルに与える前からモデル形状の確認ができます。
命令型（モデル サブクラス化）API
シンボリック API とは異なり、宣言をしない命令的（≒直感的）な書き方です。
日本（Preferred Networks）発祥の深層学習ライブラリである Chainer が最初に採用した書き方で、PyTorch もこの書き方を採用しています。
Python でクラスを書くようにモデルを実装することができるため、層の変更や拡張などのカスタマイズが容易です。
その代わりに、一度データを与えるまでモデルがどのような形状なのかをプログラム側からは認識することができません。
続いて、具体的な書き方について簡単な例を交えて紹介します。
Sequential API
その名の通り、Sequential（連続的）に層を追加してモデルを実装する API です。
Keras や TensorFlow のチュートリアルでもこの書き方が使われることが多いため、一度は見たことがあるのではないでしょうか。
以下のように、空の tensorflow.keras.Sequential クラスをインスタンス化した後に add メソッドで層を追加していく方法と、tensorflow.keras.Sequential クラスの引数にリストとして層を与えてインスタンス化する方法が一般的です。
import tensorflow as tf
from tensorflow.keras import layers
def sequential_vgg16_a(input_shape, output_size):
model = tf.keras.Sequential()
model.add(layers.Conv2D(64, 3, 1, padding="same", batch_input_shape=input_shape))
model.add(layers.BatchNormalization())
# ...(中略)...
model.add(layers.Dense(output_size, activation="softmax"))
return model
def sequential_vgg16_b(input_shape, output_size):
model = tf.keras.Sequential([
layers.Conv2D(64, 3, 1, padding="same", batch_input_shape=input_shape),
layers.BatchNormalization(),
# ...(中略)...
layers.Dense(output_size, activation="softmax")
]
return model
レイヤーを追加するメソッドのみをサポートしているため、入力、中間特徴、出力が複数になる、あるいは条件分岐が存在するような複雑なネットワークを記述することはできません。
層を順番に通していくだけの（VGG のような）シンプルなネットワークを実装する際にこの API を利用できます。
Functional API
Sequential API では記述できない複雑なモデルを実装する API です。
まず tensorflow.keras.layers.Input をインスタンス化し、最初の層に渡します。
その後、ある層の出力を次の層へと渡していくことでモデルのデータフローを定義していきます。
最後に、得られた出力と最初の入力を tensorflow.keras.Model の引数として与えることでモデルを構築できます。
from tensorflow.keras import layers, Model
def functional_vgg16(input_shape, output_size, batch_norm=False):
inputs = layers.Input(batch_input_shape=input_shape)
x = layers.Conv2D(64, 3, 1, padding="same")(inputs)
if batch_norm:
x = layers.BatchNormalization()(x)
x = layers.ReLU()(x)
# ...(中略)...
outputs = layers.Dense(output_size, activation="softmax")(x)
return Model(inputs=inputs, outputs=outputs)
上述の例では変数 batch_norm の値によって Batch Normalization 層の有無を切り替えています。
このように条件によってモデルの形状を変えるような柔軟な定義が必要な場合は Sequential API ではなく Functional API が必要になります。
なお、カッコの後にカッコが続く一見奇妙な書き方が登場しますが、これは TensorFlow 特有ではなく Python で一般的に利用できる書き方で、以下の 2 つは同じ処理を表します。
# 書き方 1
x = layers.BatchNormalization()(x)
# 書き方 2
layer = layers.BatchNormalization()
x = layer(x)
Primitive API
TensorFlow 1.X 系で主に利用されていた API です。
2.X 系の現在は非推奨となっています。
上述の Sequential API と Functional API はモデルを通るデータのフローを記述していくことでモデルを定義することができましたが、 Primitive API ではその他の計算処理を含む全体の処理フローを宣言的に記述します。
今からこの書き方を覚えるメリットはあまりないので説明は省きますが、tensorflow.Session を利用して訓練を行っている場合はこの書き方に該当します。
import tensorflow as tf
sess = tf.Session()
Subclassing API
TensorFlow 2.X へのアップデートと共に利用可能になった API です。
Chainer や PyTorch とほとんど同じ書き方であり、Python でクラスを書くようにモデルを実装することができるので直感的でカスタマイズが容易です。
まず tensorflow.keras.Model を継承してクラスを作ります。
その後、 __init__ メソッドと call メソッドを実装することでモデルを構築します。
クラス内の __init__ メソッドでは親クラスの__init__メソッド呼び出しと学習したいレイヤーの登録を行います。ここに記載していないレイヤーの重みはデフォルトでは学習対象になりません。
クラス内のcall メソッドではレイヤーの順伝播を記載します。（Chainer の __call__、PyTorch の forwardと同じようなものです。）
from tensorflow.keras import layers, Model
class VGG16(Model):
def __init__(self, output_size=1000):
super().__init__()
self.layers_ = [
layers.Conv2D(64, 3, 1, padding="same"),
layers.BatchNormalization(),
# ...(中略)...
layers.Dense(output_size, activation="softmax"),
]
def call(self, inputs):
for layer in self.layers_:
inputs = layer(inputs)
return inputs
他の書き方と比べると少々冗長にも見えますが、普通にクラスを書くようにモデルを実装できることがわかります。
なお、親クラスの初期化を行う super メソッドには引数を与えるパターンもありますが、これは 2 系の Python を考慮した書き方であり、 3 系の Python では引数なしでも同じ処理になります。
from tensorflow.keras import Model
# Python 3 系の書き方
class VGG16_PY3(Model):
def __init__(self, output_size=1000):
super().__init__()
# Python 2 系の書き方
class VGG16_PY2(Model):
def __init__(self, output_size=1000):
super().__init__(VGG16_PY2, self)
モデル構築 API 振り返り
モデル構築方法の説明は以上になります。
まとめると、以下のような使い分けができるかと思います。
- 層を一方的に通るだけのモデルを簡単に書きたい場合は Sequential API
- 複雑なモデルを学習実行前にきちんと形状確認できるように書きたい場合は Functional API
- Chainer や PyTorch 流の書き方で書きたい、あるいは動的なモデルを書きたい場合は Subclassing API
TensorFlow における 2 つの訓練方法
訓練を行う方法としては以下の 2 つが存在します。
- built-in 訓練
- カスタム訓練
※ 正式な名称があるわけではなさそうなので、本記事では便宜的に上記の名称を利用しています。
built-in 訓練
tensorflow.keras.Model の built-in function を利用して訓練を行う方法です。
Keras、 TensorFlow のチュートリアルでも利用されているためご存知の方が多いかと思います。
また、異なるライブラリですが scikit-learn でもこの方法が採用されています。
まず、上述の API で実装したモデル（tensorflow.keras.Model、あるいはこれを継承したオブジェクト）をインスタンス化します。
このインスタンスは built-in function として compile メソッドと fit メソッドを持っています。
この compile メソッドを実行して損失関数、最適化関数、評価指標を登録します。
その後 fit メソッドを実行することで訓練を行います。
import tensorflow as tf
(train_images, train_labels), _ = tf.keras.datasets.cifar10.load_data()
# 例示のため学習済みのモデルを使っています
model = tf.keras.applications.VGG16()
model.compile(
optimizer=tf.keras.optimizers.Adam(),
loss="sparse_categorical_crossentropy",
metrics=["accuracy"],
)
model.fit(train_images, train_labels)
これで訓練が実行できます。
バッチサイズの指定、エポック数、コールバック関数の登録、バリデーションデータでの評価などは、fitメソッドのキーワード引数として登録することができるため、ある程度のカスタマイズも可能です。
多くの場合はこれで十分対応できるかと思いますが、この枠にハマらないケース（例えば GAN など複数のモデルを同時に訓練するケース）は後述するカスタム訓練で記述する必要があります。
カスタム訓練
これは特別 API が用意されているというわけではなく、普通に Python の for ループで訓練する方法です。
まず、上述の API で実装したモデル（tensorflow.keras.Model、あるいはこれを継承したオブジェクト）をインスタンス化します。
次に、損失関数、最適化関数の定義に加えてデータセットのバッチ化を行います。
その後、for ループでエポック、バッチを回していきます。
for ループ内ではまず tf.GradientTape スコープの中に順伝搬の処理を記述します。
その後 gradient メソッドを呼び出して勾配を計算し、apply_gradients メソッドで最適化関数に従って重みを更新しています。
import tensorflow as tf
batch_size = 32
epochs = 10
(train_images, train_labels), _ = tf.keras.datasets.cifar10.load_data()
# 例示のため学習済みのモデルを使っています
model = tf.keras.applications.VGG16()
buffer_size = len(train_images)
train_ds = tf.data.Dataset.from_tensor_slices((train_images, train_labels))
train_ds = train_ds.shuffle(buffer_size=buffer_size).batch(batch_size)
criterion = tf.keras.losses.SparseCategoricalCrossentropy()
optimizer = tf.keras.optimizers.Adam()
for epoch in range(epochs):
for x, y_true in train_ds:
with tf.GradientTape() as tape:
y_pred = model(x, training=True)
loss = criterion(y_true=y_true, y_pred=y_pred)
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
これで訓練が実行できます。
上記の例ではバリデーションデータでの評価や TensorBoard への出力等を一切行っていませんが、普通に for ループを回しているだけなので、好きなように処理を追加していくことができます。
一方で記述量はどうしても多くなるためソースコードの品質担保がやや大変になります。
なお Chainer、PyTorch も（細かい違いはありますが）ほぼ同じ書き方ができます。
訓練方法振り返り
訓練方法の説明は以上になります。
まとめると、以下のような使い分けができるかと思います。
- 訓練中に特殊な処理を実行する必要がない場合、一般的な訓練手法の場合は built-in 訓練
- built-in の枠にハマらない場合、訓練中にいろいろな処理を追加して試行錯誤したい場合、手なりで書きたい場合は カスタム訓練
VGG16／ResNet 50 の概要
説明だけではよく分からない部分もあると思うので、実装を通して理解を深めていきます。
まず簡単に 2 つのモデルについて紹介します。
VGG16 とは
3x3 Convolution を 13 層、全結合層を 3 層持つ非常にシンプルな構造でありながら高性能なモデルです。
様々な画像認識タスクで画像特徴の抽出に利用されます。原論文は 37,000 を超える被引用数を持っており、とても有名です。
Sequential API、 Functional API、 Subclassing API で実装が可能です。
なお、原論文はこちらです。
https://arxiv.org/abs/1409.1556
ResNet50 とは
Residual 機構を持つ多層（Convolution を 49 層、全結合層を 1 層）モデルです。2020 年現在でもこの ResNet の亜種が画像分類の精度においてはトップクラスとなっており、こちらも高性能なモデルです。
VGG16 と同様に様々な画像認識タスクで画像特徴の抽出に利用されます。原論文は 45,000 を超える被引用数（BERT の約 10 倍）を持っており、こちらも非常に有名です。
Sequential API 単体では実装できません。Functional API、 Subclassing API で実装が可能です。
なお、原論文はこちらです。
https://arxiv.org/abs/1512.03385
VGG16 の実装
ではそれぞれの書き方で実装してみます。
VGG16 Sequential API
特に考えることもないので普通に書きます。
from tensorflow.keras import layers, Sequential
def sequential_vgg16(input_shape, output_size):
params = {
"padding": "same",
"use_bias": True,
"kernel_initializer": "he_normal",
}
model = Sequential()
model.add(layers.Conv2D(64, 3, 1, **params, batch_input_shape=input_shape))
model.add(layers.BatchNormalization())
model.add(layers.ReLU())
model.add(layers.Conv2D(64, 3, 1, **params))
model.add(layers.BatchNormalization())
model.add(layers.ReLU())
model.add(layers.MaxPool2D(2, padding="same"))
model.add(layers.Conv2D(128, 3, 1, **params))
model.add(layers.BatchNormalization())
model.add(layers.ReLU())
model.add(layers.Conv2D(128, 3, 1, **params))
model.add(layers.BatchNormalization())
model.add(layers.ReLU())
model.add(layers.MaxPool2D(2, padding="same"))
model.add(layers.Conv2D(256, 3, 1, **params))
model.add(layers.BatchNormalization())
model.add(layers.ReLU())
model.add(layers.Conv2D(256, 3, 1, **params))
model.add(layers.BatchNormalization())
model.add(layers.ReLU())
model.add(layers.Conv2D(256, 3, 1, **params))
model.add(layers.BatchNormalization())
model.add(layers.ReLU())
model.add(layers.MaxPool2D(2, padding="same"))
model.add(layers.Conv2D(512, 3, 1, **params))
model.add(layers.BatchNormalization())
model.add(layers.ReLU())
model.add(layers.Conv2D(512, 3, 1, **params))
model.add(layers.BatchNormalization())
model.add(layers.ReLU())
model.add(layers.Conv2D(512, 3, 1, **params))
model.add(layers.BatchNormalization())
model.add(layers.ReLU())
model.add(layers.MaxPool2D(2, padding="same"))
model.add(layers.Conv2D(512, 3, 1, **params))
model.add(layers.BatchNormalization())
model.add(layers.ReLU())
model.add(layers.Conv2D(512, 3, 1, **params))
model.add(layers.BatchNormalization())
model.add(layers.ReLU())
model.add(layers.Conv2D(512, 3, 1, **params))
model.add(layers.BatchNormalization())
model.add(layers.ReLU())
model.add(layers.MaxPool2D(2, padding="same"))
model.add(layers.Flatten())
model.add(layers.Dense(4096))
model.add(layers.Dense(4096))
model.add(layers.Dense(output_size, activation="softmax"))
return model
かなり単純に書くことができますが、層が多いので見るのが辛いことがわかります。
例えば、どこかで ReLU が抜けてても気付かなそうです。
また、例えば Batch Normalization を無くしたい、と思った場合は 1 行ずつコメントアウトしていく必要があり、再利用生やカスタマイズ性に乏しいです。
VGG16 Functional API
Sequential API と比べて柔軟に書くことができます。
今回は再利用される層のまとまり(Convolution - BatchNormalization - ReLU)を関数にしてみます。
from tensorflow.keras import layers, Model
def functional_cbr(x, filters, kernel_size, strides):
params = {
"filters": filters,
"kernel_size": kernel_size,
"strides": strides,
"padding": "same",
"use_bias": True,
"kernel_initializer": "he_normal",
}
x = layers.Conv2D(**params)(x)
x = layers.BatchNormalization()(x)
x = layers.ReLU()(x)
return x
def functional_vgg16(input_shape, output_size):
inputs = layers.Input(batch_input_shape=input_shape)
x = functional_cbr(inputs, 64, 3, 1)
x = functional_cbr(x, 64, 3, 1)
x = layers.MaxPool2D(2, padding="same")(x)
x = functional_cbr(x, 128, 3, 1)
x = functional_cbr(x, 128, 3, 1)
x = layers.MaxPool2D(2, padding="same").__call__(x) # こう書いても良い
x = functional_cbr(x, 256, 3, 1)
x = functional_cbr(x, 256, 3, 1)
x = functional_cbr(x, 256, 3, 1)
x = layers.MaxPool2D(2, padding="same").call(x) # こう書いても良い
x = functional_cbr(x, 512, 3, 1)
x = functional_cbr(x, 512, 3, 1)
x = functional_cbr(x, 512, 3, 1)
x = layers.MaxPool2D(2, padding="same")(x)
x = functional_cbr(x, 512, 3, 1)
x = functional_cbr(x, 512, 3, 1)
x = functional_cbr(x, 512, 3, 1)
x = layers.MaxPool2D(2, padding="same")(x)
x = layers.Flatten()(x)
x = layers.Dense(4096)(x)
x = layers.Dense(4096)(x)
outputs = layers.Dense(output_size, activation="softmax")(x)
return Model(inputs=inputs, outputs=outputs)
だいぶスッキリ書くことができました。
BatchNormalization を無くしたくなっても、 ReLU を LeaklyReLU に変えたくなっても数行の改修で済みます。
VGG16 Subclassing API
Functional API と同様に再利用される層のまとまり（Convolution - BatchNormalization - ReLU）をクラスにして書いてみます。
from tensorflow.keras import layers, Model
class CBR(Model):
def __init__(self, filters, kernel_size, strides):
super().__init__()
params = {
"filters": filters,
"kernel_size": kernel_size,
"strides": strides,
"padding": "same",
"use_bias": True,
"kernel_initializer": "he_normal",
}
self.layers_ = [
layers.Conv2D(**params),
layers.BatchNormalization(),
layers.ReLU()
]
def call(self, inputs):
for layer in self.layers_:
inputs = layer(inputs)
return inputs
class VGG16(Model):
def __init__(self, output_size=1000):
super().__init__()
self.layers_ = [
CBR(64, 3, 1),
CBR(64, 3, 1),
layers.MaxPool2D(2, padding="same"),
CBR(128, 3, 1),
CBR(128, 3, 1),
layers.MaxPool2D(2, padding="same"),
CBR(256, 3, 1),
CBR(256, 3, 1),
CBR(256, 3, 1),
layers.MaxPool2D(2, padding="same"),
CBR(512, 3, 1),
CBR(512, 3, 1),
CBR(512, 3, 1),
layers.MaxPool2D(2, padding="same"),
CBR(512, 3, 1),
CBR(512, 3, 1),
CBR(512, 3, 1),
layers.MaxPool2D(2, padding="same"),
layers.Flatten(),
layers.Dense(4096),
layers.Dense(4096),
layers.Dense(output_size, activation="softmax"),
]
def call(self, inputs):
for layer in self.layers_:
inputs = layer(inputs)
return inputs
__init__ でモデルの定義、call でモデルの呼び出しを担当しているため、Functional API よりも直感的に理解しやすいですが、コードは長めになります。
また命令型である Subclassing API はモデル生成時に入力の形状が必要ない（引数に input_shape が必要ない）こともポイントです。
VGG16 実装 振り返り
なるべく比較しやすいように書いたつもりですがいかがだったでしょうか。
今回の実装は Batch Normalization を Convolution 層の間に挟んでいるのと重みの初期化に He の初期化を利用していますが、原論文が提出された時はこれらのテクニックは未だ発表されていなかったので Batch Normalization 層はなく、重みの初期化には Grolot の初期化が使われていました。そのため、原論文では勾配消失を回避するために 7 層のモデルを学習してから層を徐々に付け足していくといった転移学習ライクな学習方法が採用されています。
上記の実装をより深く理解するために Batch Normalization 層をなくすと何が起きるのか、重みの初期化手法を変えると何が起きるのか、などを試してみるのも面白いと思います。
ResNet50 の実装
続いて ResNet50 を実装します。
Sequential API 単体では書けないので Functional API、 Subclassing API で書きます。
ResNet50 Functional API
再利用される Residual 機構を関数化して実装します。
from tensorflow.keras import layers, Model
def functional_bottleneck_residual(x, in_ch, out_ch, strides=1):
params = {
"padding": "same",
"kernel_initializer": "he_normal",
"use_bias": True,
}
inter_ch = out_ch // 4
h1 = layers.Conv2D(inter_ch, kernel_size=1, strides=strides, **params)(x)
h1 = layers.BatchNormalization()(h1)
h1 = layers.ReLU()(h1)
h1 = layers.Conv2D(inter_ch, kernel_size=3, strides=1, **params)(h1)
h1 = layers.BatchNormalization()(h1)
h1 = layers.ReLU()(h1)
h1 = layers.Conv2D(out_ch, kernel_size=1, strides=1, **params)(h1)
h1 = layers.BatchNormalization()(h1)
if in_ch != out_ch:
h2 = layers.Conv2D(out_ch, kernel_size=1, strides=strides, **params)(x)
h2 = layers.BatchNormalization()(h2)
else:
h2 = x
h = layers.Add()([h1, h2])
h = layers.ReLU()(h)
return h
def functional_resnet50(input_shape, output_size):
inputs = layers.Input(batch_input_shape=input_shape)
x = layers.Conv2D(64, 7, 2, padding="same", kernel_initializer="he_normal")(inputs)
x = layers.BatchNormalization()(x)
x = layers.MaxPool2D(pool_size=3, strides=2, padding="same")(x)
x = functional_bottleneck_residual(x, 64, 256)
x = functional_bottleneck_residual(x, 256, 256)
x = functional_bottleneck_residual(x, 256, 256)
x = functional_bottleneck_residual(x, 256, 512, 2)
x = functional_bottleneck_residual(x, 512, 512)
x = functional_bottleneck_residual(x, 512, 512)
x = functional_bottleneck_residual(x, 512, 512)
x = functional_bottleneck_residual(x, 512, 1024, 2)
x = functional_bottleneck_residual(x, 1024, 1024)
x = functional_bottleneck_residual(x, 1024, 1024)
x = functional_bottleneck_residual(x, 1024, 1024)
x = functional_bottleneck_residual(x, 1024, 1024)
x = functional_bottleneck_residual(x, 1024, 1024)
x = functional_bottleneck_residual(x, 1024, 2048, 2)
x = functional_bottleneck_residual(x, 2048, 2048)
x = functional_bottleneck_residual(x, 2048, 2048)
x = layers.GlobalAveragePooling2D()(x)
outputs = layers.Dense(
output_size, activation="softmax", kernel_initializer="he_normal"
)(x)
return Model(inputs=inputs, outputs=outputs)
functional_bottleneck_residual メソッド内では h1、h2、h が登場します。
このように、データのフローが途中で分岐するモデルは Sequential API では記述できません。
また、h2 は入出力のチャネル数が同じ場合は何もせず、異なる場合はチャネル数を調整する処理（Projection）を行います。このような条件分岐も Sequential API では記述できません。
このメソッドを作ってしまえば、あとは順々に記述していくだけです。
ResNet50 Subclassing API
Functional API と同様に再利用される Residual 機構をクラス化して実装します。
from tensorflow import layers, Model
class BottleneckResidual(Model):
"""ResNet の Bottleneck Residual Module です.
1 層目の 1x1 conv で ch 次元を削減することで
2 層目の 3x3 conv の計算量を減らし
3 層目の 1x1 conv で ch 出力の次元を復元します.
計算量の多い 2 層目 3x3 conv の次元を小さくすることから bottleneck と呼ばれます.
"""
def __init__(self, in_ch, out_ch, strides=1):
super().__init__()
self.projection = in_ch != out_ch
inter_ch = out_ch // 4
params = {
"padding": "same",
"kernel_initializer": "he_normal",
"use_bias": True,
}
self.common_layers = [
layers.Conv2D(inter_ch, kernel_size=1, strides=strides, **params),
layers.BatchNormalization(),
layers.ReLU(),
layers.Conv2D(inter_ch, kernel_size=3, strides=1, **params),
layers.BatchNormalization(),
layers.ReLU(),
layers.Conv2D(out_ch, kernel_size=1, strides=1, **params),
layers.BatchNormalization(),
]
if self.projection:
self.projection_layers = [
layers.Conv2D(out_ch, kernel_size=1, strides=strides, **params),
layers.BatchNormalization(),
]
self.concat_layers = [layers.Add(), layers.ReLU()]
def call(self, inputs):
h1 = inputs
h2 = inputs
for layer in self.common_layers:
h1 = layer(h1)
if self.projection:
for layer in self.projection_layers:
h2 = layer(h2)
outputs = [h1, h2]
for layer in self.concat_layers:
outputs = layer(outputs)
return outputs
class ResNet50(Model):
"""ResNet50 です.
要素は
conv * 1
resblock(conv * 3) * 3
resblock(conv * 3) * 4
resblock(conv * 3) * 6
resblock(conv * 3) * 3
dense * 1
から構成されていて, conv * 49 + dense * 1 の 50 層です.
"""
def __init__(self, output_size=1000):
super().__init__()
self.layers_ = [
layers.Conv2D(64, 7, 2, padding="same", kernel_initializer="he_normal"),
layers.BatchNormalization(),
layers.MaxPool2D(pool_size=3, strides=2, padding="same"),
BottleneckResidual(64, 256),
BottleneckResidual(256, 256),
BottleneckResidual(256, 256),
BottleneckResidual(256, 512, 2),
BottleneckResidual(512, 512),
BottleneckResidual(512, 512),
BottleneckResidual(512, 512),
BottleneckResidual(512, 1024, 2),
BottleneckResidual(1024, 1024),
BottleneckResidual(1024, 1024),
BottleneckResidual(1024, 1024),
BottleneckResidual(1024, 1024),
BottleneckResidual(1024, 1024),
BottleneckResidual(1024, 2048, 2),
BottleneckResidual(2048, 2048),
BottleneckResidual(2048, 2048),
layers.GlobalAveragePooling2D(),
layers.Dense(
output_size, activation="softmax", kernel_initializer="he_normal"
),
]
def call(self, inputs):
for layer in self.layers_:
inputs = layer(inputs)
return inputs
Functional API とそこまで大きな違いはありません。
__init__ レイヤーでリストに層をまとめるような書き方をしていますが、クラスの変数に登録されていれば良いのでこの辺は自由に書くことができます。
ResNet50 実装 振り返り
Sequential API 単体では実装できないモデルとして ResNet50 を紹介しました。
正直大きな違いはないので Functional API、 Subclassing API は好みで使い分けても問題ないかと思います。
訓練の実装
最後に訓練ループの実装を比較してみます。
全てのソースコードを載せるとかなり長くなるため部分的にメソッドを src.utils に切り出しています。
そこまで複雑なことをしているわけではないので補完しながら読んでいただければ助かります。
一応全てのソースは以下のリポジトリにあるので気になる方はご覧ください。
https://github.com/Anieca/deep-learning-models
built-in 訓練 実装
テストデータの精度算出、TensorBoard 用のログ出力などいくつかオプションを指定してみます。
import os
import tensorflow as tf
from src.utils import load_dataset, load_model, get_args, get_current_time
def builtin_train(args):
# 1. load dataset and model
(train_images, train_labels), (test_images, test_labels) = load_dataset(args.data)
input_shape = train_images[: args.batch_size, :, :, :].shape
output_size = max(train_labels) + 1
model = load_model(args.arch, input_shape=input_shape, output_size=output_size)
model.summary()
# 2. set tensorboard cofigs
logdir = os.path.join(args.logdir, get_current_time())
tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=logdir)
# 3. loss, optimizer, metrics setting
model.compile(
optimizer=tf.keras.optimizers.Adam(),
loss="sparse_categorical_crossentropy",
metrics=["accuracy"],
)
# 4. dataset config (and validation, callback config)
fit_params = {}
fit_params["batch_size"] = args.batch_size
fit_params["epochs"] = args.max_epoch
if args.steps_per_epoch:
fit_params["steps_per_epoch"] = args.steps_per_epoch
fit_params["verbose"] = 1
fit_params["callbacks"] = [tensorboard_callback]
fit_params["validation_data"] = (test_images, test_labels)
# 5. start train and test
model.fit(train_images, train_labels, **fit_params)
かなりシンプルに書けます。
コールバック関数は他にも色々あるので興味がある方はドキュメントを読んでみてください。
https://www.tensorflow.org/api_docs/python/tf/keras/callbacks
カスタム訓練 実装
上記の built-in 訓練と同じ処理を自分で実装してみます。
import os
import tensorflow as tf
from src.utils import load_dataset, load_model, get_args, get_current_time
def custom_train(args):
# 1. load dataset and model
(train_images, train_labels), (test_images, test_labels) = load_dataset(args.data)
input_shape = train_images[: args.batch_size, :, :, :].shape
output_size = max(train_labels) + 1
model = load_model(args.arch, input_shape=input_shape, output_size=output_size)
model.summary()
# 2. set tensorboard configs
logdir = os.path.join(args.logdir, get_current_time())
train_writer = tf.summary.create_file_writer(os.path.join(logdir, "train"))
test_writer = tf.summary.create_file_writer(os.path.join(logdir, "test"))
# 3. loss, optimizer, metrics setting
criterion = tf.keras.losses.SparseCategoricalCrossentropy()
optimizer = tf.keras.optimizers.Adam()
train_loss_avg = tf.keras.metrics.Mean()
train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()
test_loss_avg = tf.keras.metrics.Mean()
test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()
# 4. dataset config
buffer_size = len(train_images)
train_ds = tf.data.Dataset.from_tensor_slices((train_images, train_labels))
train_ds = train_ds.shuffle(buffer_size=buffer_size).batch(args.batch_size)
test_ds = tf.data.Dataset.from_tensor_slices((test_images, test_labels))
test_ds = test_ds.batch(args.batch_size)
# 5. start train and test
for epoch in range(args.max_epoch):
# 5.1. initialize metrics
train_loss_avg.reset_states()
train_accuracy.reset_states()
test_loss_avg.reset_states()
test_loss_avg.reset_states()
# 5.2. initialize progress bar
train_pbar = tf.keras.utils.Progbar(args.steps_per_epoch)
test_pbar = tf.keras.utils.Progbar(args.steps_per_epoch)
# 5.3. start train
for i, (x, y_true) in enumerate(train_ds):
if args.steps_per_epoch and i >= args.steps_per_epoch:
break
# 5.3.1. forward
with tf.GradientTape() as tape:
y_pred = model(x, training=True)
loss = criterion(y_true=y_true, y_pred=y_pred)
# 5.3.2. calculate gradients from `tape` and backward
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
# 5.3.3. update metrics and progress bar
train_loss_avg(loss)
train_accuracy(y_true, y_pred)
train_pbar.update(
i + 1,
[
("avg_loss", train_loss_avg.result()),
("accuracy", train_accuracy.result()),
],
)
# 5.4. start test
for i, (x, y_true) in enumerate(test_ds):
if args.steps_per_epoch and i >= args.steps_per_epoch:
break
# 5.4.1. forward
y_pred = model(x)
loss = criterion(y_true, y_pred)
# 5.4.2. update metrics and progress bar
test_loss_avg(loss)
test_accuracy(y_true, y_pred)
test_pbar.update(
i + 1,
[
("avg_test_loss", test_loss_avg.result()),
("test_accuracy", test_accuracy.result()),
],
)
# 5.5. write metrics to tensorboard
with train_writer.as_default():
tf.summary.scalar("Loss", train_loss_avg.result(), step=epoch)
tf.summary.scalar("Acc", train_accuracy.result(), step=epoch)
with test_writer.as_default():
tf.summary.scalar("Loss", test_loss_avg.result(), step=epoch)
tf.summary.scalar("Acc", test_accuracy.result(), step=epoch)
訓練開始まではそこまで変わらないですが、訓練ループ内（コメントの 5. ）の記述量がかなり多くなります。
訓練実装 振り返り
自分で TensorBoard 出力の管理やプログレスバーの作成と言ったユーティリティを管理するのはそれなりにコストがかかりますが、 built-in はかなり楽に使うことができます。
built-in に用意されていない処理を記述したい場合はカスタム訓練で書く必要がありますが、そうでなければ built-in を使った方が良さそうです。
終わりに
以上です。お疲れ様でした。
TensorFlow 2 系の色々な書き方について実装を交えて紹介しました。
あまりそれぞれの書き方について優劣はつけずにフラットに書いたつもりです。
自分で書くときは状況や好みに合わせて書けば良いと思いますが、ソースコードを探していると色々な書き方に出会うので、なんとなく全ての書き方を理解していると良いかと思います。
皆様のお役に立てれば幸いです。