😎
[機械学習] VAEの実装
はじめに
vaeをTensorflow2で実装してみました.
データセット、ライブラリはそれぞれ以下を使用しています.
ライブラリ: Tensorflow2
データセット: FashionMnist
本記事で紹介しているコードは以下で実行できます.
VAEとは
端的に言うと、オートエンコーダ型の生成モデルです.
生成モデルとは
データが何かしら確率的な分布から生成されていると仮定し、
データが生成される元の確率分布を学習するモデルのことをいいます.
オートエンコーダとは
データをより少ない次元に変換する次元圧縮器で、以下のような構造をとります.
入力値をエンコーダーで潜在変数に圧縮し、潜在変数からデコーダーで元の次元に復元します.
学習のゴール
VAEでは、データが正規分布から生成されていると仮定します.
正規分布は期待値と分散によって決まるのでした.
データが生成される元の変数を潜在変数と呼び、VAEでは潜在変数の期待値と分散を学習することがゴールになります.
VAEの構造
どうやって学習するのかというと、以下の形をとります.
- エンコーダは期待値と分散の２つを出力.
- 期待値と分散から潜在変数をサンプリング.
- デコーダで潜在変数から元の次元を復元する
デコーダは通常のオートエンコーダと同じですが、
エンコーダの出力と、そこからサンプリング処理があることが、オートエンコーダとの大きな違いになります.
エンコーダの実装
エンコーダは元のデータから期待値と分散の2つを出力します.
分散は必ず0以上になる必要がありますが、全結合層の出力は必ず0以上になるわけではないので、エンコーダーから出力される分散は対数を取ったものを出力します.
class Encoder(layers.Layer):
def __init__(self, latent_dim=10):
super().__init__()
self.c1 = layers.Conv2D(32, 3, activation="relu", strides=2, padding="same")
self.c2 = layers.Conv2D(64, 3, activation="relu", strides=2, padding="same")
self.f1 = layers.Flatten()
self.d1 = layers.Dense(16, activation="relu")
self.d2 = layers.Dense(latent_dim, name="u")
self.d3 = layers.Dense(latent_dim, name="var")
def call(self, x):
x = self.c1(x)
x = self.c2(x)
x = self.f1(x)
x = self.d1(x)
u = self.d2(x)
log_var = self.d3(x)
return u, log_var
サンプリング層の実装
エンコーダから出力した期待値と分散から潜在変数をサンプリングしています.
誤差逆伝播で学習するには微分可能でないといけないので、以下のような式で擬似的なサンプリングをします
この方法をreparameterization trickといいます.
潜在変数 = 期待値(u) + 標準偏差(σ) * 乱数
class Sampling(layers.Layer):
def call(self, u, log_var):
u_shape = tf.shape(u)
epsilon = tf.keras.backend.random_normal(shape=(u_shape[0], u_shape[1]))
return u + tf.exp(0.5 * log_var) * epsilon
デコーダの実装
デコーダでは、潜在変数から元のデータを復元します.
ここはオートエンコーダと同じです.
class Decoder(layers.Layer):
def __init__(self):
super().__init__()
self.d1 = layers.Dense(7 * 7 * 64, activation="relu")
self.r1 = layers.Reshape((7, 7, 64))
self.c1 = layers.Conv2DTranspose(64, 3, activation="relu", strides=2, padding="same")
self.c2 = layers.Conv2DTranspose(32, 3, activation="relu", strides=2, padding="same")
self.c3 = layers.Conv2DTranspose(1, 3, activation="sigmoid", padding="same")
def call(self, x):
x = self.d1(x)
x = self.r1(x)
x = self.c1(x)
x = self.c2(x)
y = self.c3(x)
return y
損失関数の実装
VAEの損失関数は2つの項からなります.
VAEの損失 = 再構成ロス + KLダイバージェンス
再構成ロス
- 生成されたデータと元データの差を計算
- 今回は交差エントロピーを採用
- 大体は2乗誤差か交差エントロピーのどちらかが使われるが、どちらにするかはケースバイケース.
KLダイバージェンス
- 2つの確率分布間の距離
- VAEでの定義は以下
- Q: 仮定している正規分布(平均0 分散1)
- P: 実際のエンコーダの出力
- エンコーダの出力を平均0,分散1に正規分布に近づけることが目標
- 定義をもとに計算すると最終的に以下の式になる
class VAELossFunction:
def __call__(self, u, log_var, z, y, y_true) -> (float, float, float):
# 再構成ロス
rc_loss = tf.reduce_mean(
tf.reduce_sum(
tf.keras.losses.binary_crossentropy(y_true, y),
axis=(1, 2)
)
)
# KLダイバージェンス
kl_loss = tf.reduce_mean(
tf.reduce_sum(
-0.5 * (1 + log_var - tf.square(u) - tf.exp(log_var)),
axis=1
))
return rc_loss, kl_loss, rc_loss + kl_loss
モデル全体の実装
モデル全体の計算の流れを記述しています.
- 入力値をエンコーダーに流して期待値と分散を取得
- 期待値と分散をサンプリング層に渡して潜在変数を取得
- 潜在変数からデコーダで元の次元に復元
💡Point
- 訓練用のメソッド(train_step)の出力はVAEResultとし定義
- 学習時の計算はtf.functionでデコレートしたtrain_step_functionで定義
- eager modeで動的に計算グラフを作成するため
- tf.functionはtf.Tensorに変換できる戻り値である必要があるので、train_stepとは別メソッドで定義
- 予測用にバッチサイズなしの画像を受け取って潜在変数と復元データを返すメソッドを作成(predict)
@dataclass
class VAEResult:
rc_loss: float
kl_loss: float
total_loss: float
class VAE(Model):
def __init__(self, optimizer):
super().__init__()
self.encoder = Encoder(latent_dim=2)
self.decoder = Decoder()
self.sampling = Sampling()
self.loss_function = VAELossFunction()
self.optimizer = optimizer
def call(self, x: tf.Tensor):
u, var = self.encoder(x)
z = self.sampling(u, var)
y = self.decoder(z)
return u, var, z, y
@tf.function
def train_step_function(self, y_true: tf.Tensor):
with tf.GradientTape() as tape:
u, var, z, y = self.call(y_true)
rc_loss, kl_loss, total_loss = self.loss_function(u, var, z, y, y_true)
grads = tape.gradient(total_loss, self.trainable_weights)
self.optimizer.apply_gradients(zip(grads, self.trainable_weights))
return rc_loss, kl_loss, total_loss
def train_step(self, y_true: tf.Tensor) -> VAEResult:
rc_loss, kl_loss, total_loss = self.train_step_function(y_true)
return VAEResult(
rc_loss=rc_loss,
kl_loss=kl_loss,
total_loss=total_loss
)
def test_step(self, y_test: tf.Tensor) -> (tf.Tensor, np.array):
u, var, z, y = self.call(y_test)
rc_loss, kl_loss, total_loss = self.loss_function(u, var, z, y, y_test)
return VAEResult(
rc_loss=rc_loss,
kl_loss=kl_loss,
total_loss=total_loss
)
def predict(self, image):
tensor = tf.convert_to_tensor(image, dtype=tf.float32)
tensor = tf.expand_dims(tensor, axis=0)
u, var, z, y = self.call(tensor)
z = tf.squeeze(z, [0])
y = tf.squeeze(y, [0])
return z, y
メトリクスの実装
学習時のメトリクスを定義してます
- update_on_trainでモデルで計算された各損失の平均を記録
- displayでメトリクスの情報を表示
- resetでメトリクスの状態をクリア
import tensorflow as tf
from tensorflow.keras import metrics
class VAETrainMetrics:
def __init__(self):
self.rc_loss = metrics.Mean(name='train_rc_loss')
self.kl_loss = metrics.Mean(name='train_kl_loss')
self.total_loss = metrics.Mean(name='train_total_loss')
def update_on_train(self, result: VAEResult):
self.rc_loss.update_state(result.rc_loss)
self.kl_loss.update_state(result.kl_loss)
self.total_loss.update_state(result.total_loss)
def update_on_test(self, result: VAEResult):
pass
def reset(self):
self.rc_loss.reset_states()
self.kl_loss.reset_states()
self.total_loss.reset_states()
def display(self, epoch: int):
template = 'Epoch {}, RC Loss: {:.2g}, KL Loss: {:.2g}, Total Loss: {:.2g}'
print(
template.format(
epoch,
self.rc_loss.result(),
self.kl_loss.result(),
self.total_loss.result()
)
)
学習器の実装
いつも通り、学習のイテレーションを記載してます.
以下のステップをエポック数だけ繰り返します.
- バッチ化したデータをループで回す
- モデル渡して損失を取得
- 損失をメトリクスに記録
class VAETrainer:
def __init__(self, model, metrics):
self.model = model
self.metrics = metrics
def fit(self, dataset, epochs: int = 3):
""" Train mdoel by epochs
"""
for e in range(1, epochs + 1):
# train model
for images in dataset.train_loop():
result = self.model.train_step(images)
# test model
for images in dataset.test_loop():
result = self.model.test_step(images)
self.metrics.update_on_test(result)
# show metrics
self.metrics.display(e)
self.metrics.reset()
データセットの用意
お馴染みのFashionMnistを使用しました.
学習の実行
とりあえず、10エポックほど学習してみます.
from tensorflow.keras import optimizers
# モデル初期化
vae = VAE(optimizer=optimizers.Adam())
# メトリクス初期化
vae_metrics = VAETrainMetrics()
# 学習用クラス初期化
vae_trainer = VAETrainer(vae, vae_metrics)
# データセットクラス初期化
fashion_mnist_dataset = FashionMnistDataset()
# 学習実行
vae_trainer.fit(fashion_mnist_dataset, epochs=10)
実行結果
Epoch 1, RC Loss: 2.9e+02, KL Loss: 7.5, Total Loss: 3e+02
Epoch 2, RC Loss: 2.7e+02, KL Loss: 7.5, Total Loss: 2.8e+02
Epoch 3, RC Loss: 2.7e+02, KL Loss: 7.4, Total Loss: 2.7e+02
Epoch 4, RC Loss: 2.6e+02, KL Loss: 7.3, Total Loss: 2.7e+02
Epoch 5, RC Loss: 2.6e+02, KL Loss: 7.1, Total Loss: 2.7e+02
Epoch 6, RC Loss: 2.6e+02, KL Loss: 7, Total Loss: 2.7e+02
Epoch 7, RC Loss: 2.6e+02, KL Loss: 6.8, Total Loss: 2.7e+02
Epoch 8, RC Loss: 2.6e+02, KL Loss: 6.7, Total Loss: 2.7e+02
Epoch 9, RC Loss: 2.6e+02, KL Loss: 6.6, Total Loss: 2.7e+02
Epoch 10, RC Loss: 2.6e+02, KL Loss: 6.6, Total Loss: 2.7e+02
学習結果
適当に組むましたが、割とちゃんと学習されました.
多少ぼやけてるのは、VAEだと致し方ない部分です.
より鮮明に学習したい場合は、GANなどを使う必要がありそうです.
元画像
復元した画像
おわりに
tensorflow1系で実装したことは以前ありますが、2系ではなかったので、実装してみました.
非常にスッキリかけて良いですね.
次はGANや自然言絵処理への応用など、より発展した実装をしたくなりました.
最後に、弊社では絶賛エンジニア募集中なので、気になる方は以下を覗いてみてください~
Discussion