【実装編】敵対的生成ネットワーク（GAN）とは？

Tensorflowのチュートリアルをカスタマイズして、なるべく少ないコードでGANを実装しました。

必要なライブラリのインストール

# ライブラリのインストール
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
from tensorflow.keras import layers

データセットのダウンロード

# データセットのロード
(train_images, train_labels), (_, _) = tf.keras.datasets.mnist.load_data()
train_images = train_images.reshape(train_images.shape[0], 28, 28, 1).astype('float32')
train_images = (train_images - 127.5) / 127.5  # Normalize the images to [-1, 1]
# Batch and shuffle the data
train_dataset = tf.data.Dataset.from_tensor_slices(train_images).shuffle(60000).batch(256)

PCのスペックが低くても動くように、サイズの小さいMNIST（手書き数字文字）を使用します。
後半は、GANで使用できるようにデータの形式を整えています。

モデルの定義

ジェネレータとディスクリミネータを作成していきます。

・ジェネレータ
一次元ベクトルから、最終的に28×28の画像を生成できるように逆畳み込み層を定義しています

def make_generator_model():
    model = tf.keras.Sequential()
    model.add(layers.Dense(7*7*256, use_bias=False, input_shape=(100,)))
    model.add(layers.BatchNormalization())
    model.add(layers.LeakyReLU())

    model.add(layers.Reshape((7, 7, 256)))
    assert model.output_shape == (None, 7, 7, 256)  # Note: None is the batch size

    model.add(layers.Conv2DTranspose(128, (5, 5), strides=(1, 1), padding='same', use_bias=False))
    assert model.output_shape == (None, 7, 7, 128)
    model.add(layers.BatchNormalization())
    model.add(layers.LeakyReLU())

    model.add(layers.Conv2DTranspose(64, (5, 5), strides=(2, 2), padding='same', use_bias=False))
    assert model.output_shape == (None, 14, 14, 64)
    model.add(layers.BatchNormalization())
    model.add(layers.LeakyReLU())

    model.add(layers.Conv2DTranspose(1, (5, 5), strides=(2, 2), padding='same', use_bias=False, activation='tanh'))
    assert model.output_shape == (None, 28, 28, 1)

    return model

入力：1×12544
　　↓
一層目（全結合層）：7×7×256
　　↓
二層目（逆畳み込み層）：7×7×128
　　↓
三層目（逆畳み込み層）：14×14×64
　　↓
出力層（逆畳み込み層）：28×28

・ディスクリミネータ
28×28の画像を、0(本物) or 1(偽物)で判別するように畳み込み層と全結合層がつながっています

def make_discriminator_model():
    model = tf.keras.Sequential()
    model.add(layers.Conv2D(64, (5, 5), strides=(2, 2), padding='same', input_shape=[28, 28, 1]))
    model.add(layers.LeakyReLU())
    model.add(layers.Dropout(0.3))

    model.add(layers.Conv2D(128, (5, 5), strides=(2, 2), padding='same'))
    model.add(layers.LeakyReLU())
    model.add(layers.Dropout(0.3))

    model.add(layers.Flatten())
    model.add(layers.Dense(1))

    return model

入力：28×28
　　↓
一層目（畳み込み層）：14×14×64
　　↓
二層目（畳み込み層）：7×7×128
　　↓
三層目（ベクトル変換）：6272
　　↓
出力層（全結合層）：1 (0 or 1)

モデルの設定

損失関数や最適化関数の設定を行います。

# モデルの定義
generator = make_generator_model()
discriminator = make_discriminator_model()

# 損失関数の定義
cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True)

# ディスクリミネータ用の損失関数
def discriminator_loss(real_output, fake_output):
    real_loss = cross_entropy(tf.ones_like(real_output), real_output)
    fake_loss = cross_entropy(tf.zeros_like(fake_output), fake_output)
    total_loss = real_loss + fake_loss
    return total_loss

# ジェネレータ用の損失関数
def generator_loss(fake_output):
    return cross_entropy(tf.ones_like(fake_output), fake_output)

# 最適化関数の定義
generator_optimizer = tf.keras.optimizers.Adam(1e-4)
discriminator_optimizer = tf.keras.optimizers.Adam(1e-4)

ディスクリミネータとジェネレータに、個別に損失関数を定義します。
どちらも交差エントロピー関数を使用しているため、最初に宣言をしています。

学習

モデルの学習を行います。
チュートリアルではtrain関数を作成していますが、一目見て理解しやすいように少しカスタマイズしています。

# 学習ループ
for epoch in range(100):
  for image_batch in train_dataset:
      # ランダムノイズを生成
      noise = tf.random.normal([image_batch.shape[0], 100])  # バッチサイズに合わせてノイズを生成
  # 勾配計算のためのテープを作成
  with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
      # Generatorで画像生成
      generated_images = generator(noise, training=True)
      # Discriminatorで画像判定
      real_output = discriminator(image_batch, training=True)
      fake_output = discriminator(generated_images, training=True)
      # 損失を計算
      gen_loss = generator_loss(fake_output)
      disc_loss = discriminator_loss(real_output, fake_output)
  # 勾配を計算
  gradients_of_generator = gen_tape.gradient(gen_loss, generator.trainable_variables)
  gradients_of_discriminator = disc_tape.gradient(disc_loss, discriminator.trainable_variables)

  # 勾配を適用
  generator_optimizer.apply_gradients(zip(gradients_of_generator, generator.trainable_variables))
  discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, discriminator.trainable_variables))
  
  # 途中経過を表示
  if epoch % 10 == 0:
     print(f'Epoch {epoch + 1}, Generator Loss: {gen_loss.numpy()}, Discriminator Loss: {disc_loss.numpy()}')

1. ランダムベクトルを生成
   ・ジェネレータに入力するためのベクトルを生成します

2. ジェネレータで画像生成
   ・ベクトルから画像を生成します

3. ディスクリミネータで画像判別
   ・データセット内の真の画像が本物か偽物かを判別します
   ・ジェネレータから生成された偽の画像が本物か偽物かを判別します

4. それぞれの損失を計算
   ・ディスクリミネータの損失は、それぞれの画像を見破れたかどうかで計算されます
   ・ジェネレータの損失は、自身の生成画像でディスクリミネータを騙せたかどうかで計算されます

5. 勾配降下法にて学習を行う
   ・それぞれの損失をフィードバックし、モデルを学習します

まとめ

学習後のモデルを使用して、画像生成を行います
Google Colaboratoryで実行しやすいように、モデルを保存せずにそのまま画像を生成するようコードを変更しています

# 学習後のモデルで生成タスク
seed = tf.random.normal([4, 100])
predictions = generator(seed, training=False)

fig = plt.figure(figsize=(4, 4))
for i in range(predictions.shape[0]):
  plt.subplot(4, 4, i+1)
  plt.imshow(predictions[i, :, :, 0] * 127.5 + 127.5, cmap='gray')
  plt.axis('off')

plt.show()

左から、学習毎の途中経過です
回数を重ねるごとに、オリジナルのデータに近い画像が生成できていることが分かります

学習ごとの生成画像

以上、GANによる画像生成の実装を行いました。
データセットやモデルをカスタマイズすることで、よりリアルが画像生成を行えるようになります。

是非参考にしてみてください。