DCGAN(深層畳み込み敵対的生成ネットワーク)のコードを追ってみる

この記事は[Akatsuki Advent Calendar 2021](https://adventar.org/calendars/6566)の20日目の記事です。
機械学習を用いた **画像の自動生成** について解説していきたいと思います。

# はじめに

最近、「機械学習、面白そう！」と思って勉強をはじめました。
[TensorFlow のチュートリアルに DCGAN の実装がある](https://www.tensorflow.org/tutorials/generative/dcgan?hl=ja) のですが「機械学習少し勉強したぞ！」という状態で挑んで返り討ちにあってしまったので、できる範囲で解説していきたいと思います！
ソースコードは一番最後に載せているので、とりあえず動かしたいという方は試してみてください。
基本的にチュートリアルのコードをベースにアレンジしてます。

### 前提知識
「ニューラルネットワークはなんとなくわかる」くらいを対象にしています。
もし、ニューラルネットワークも知らない！という方は「ゼロから作るDeep Learning」はとても勉強になったのでおすすめです。

# DCGAN で何ができるのか
まずはどんなことができるのか紹介します。
DCGAN を使うことでデータセットの特徴を学習して、モデルが画像を生成できるようになります。

これは MNIST を呼ばれる画像処理の学習で広く使われているデータセットで学習してみた結果です。
28 x 28 のグレースケールの画像なので非常に小さなサイズではありますが、
学習の終盤になると、人が書いたものと見分けが付かないくらいの精度になっていることがわかります。

#### 手書きデータ
<img width="300" alt="aaa" src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/325895/f4b867dd-de47-a288-144d-b2d149b80196.png">

#### 学習過程と作成された画像
<img width="800" alt="aaa" src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/325895/632ec2a6-e0c9-ea34-5b94-f405a9189191.png">

# DCGAN

### GAN とは
GAN(Generative Adversarial Networks)とは敵対的生成ネットワークとも呼ばれ、教師なし機械学習のニューラルネットワークモデルの一種です。
通常のニューラルネットワークは1つのモデルを学習させますが、 GAN は生成器(Generator)と識別器(Discriminator)の2つのモデルを同時に学習させます。

- Discriminator は **データが本物なのかどうか** の判断を精度高く判定するように学習する
- Generator は **Discriminator が誤判断させるような** データを作成するように学習する

この2つのモデルの学習を同時に行っていくことで、本当の画像と区別がつかないようなデータを生成させることができます。
また、ランダムノイズを変更することで、全く違うデータを生成することができます。
（生成する画像をコントロールすることができないので、「1」の手書き文字を狙って生成することは通常GANではできません。CGANというモデルがありますが今回は割愛します。）

つまり、DCGANとは畳み込みを使って、データを学習し、作成するモデルのことです。

### GANの学習の流れ

1. Generator にランダムノイズを入力する
1. Generator が(画像)データを作成する
1. Discriminator に Generator が作成した(偽物の)データと(本物のデータである)データセットを入力する
1. 結果をそれぞれに誤差逆伝播する（正しい値との差分を計算して近づける）
<img width="800" alt="aaa" src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/325895/ad716ef9-54ea-8f9b-b5f6-e3c338f524a4.png">

# コード解説

ここからは TensorFlow のチュートリアルに沿ったコードの解説をしていきたいと思います。
一部 GAN に関係のない部分は割愛しています。

## 実行環境
```
Python     : 3.7.4
tensorflow : 2.7.0
keras      : 2.7.0
```

### インポートする

```python
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
import os
import time
from IPython import display
from keras.models import Sequential
from keras.layers import Activation, BatchNormalization, Dense, Dropout, Flatten, Reshape, LeakyReLU
from keras.layers.convolutional import Conv2D, Conv2DTranspose
```

### データを読み込む
```python
# MNIST は(訓練データ, 訓練データのラベル), (テストデータ, テストデータのラベル)という情報で渡される
# 今回は教師なし学習でラベルとテストデータは必要ないので、訓練データのみを読み込む
# fashion_mnist になっていますが、手書き文字の学習をしたい場合はここを mnist にしてください
(train_images, _), (_, _) = tf.keras.datasets.fashion_mnist.load_data()
# (60000, 28, 28) → (60000, 28, 28, 1) に変換する
# カラー画像は (データ数, 幅, 高さ, 3) で表現する
train_images = train_images.reshape(train_images.shape[0], 28, 28, 1).astype('float32')
# Generator で使用するtanhを使うと [-1, 1] で出力されるため、同様に合わせる
train_images = (train_images - 127.5) / 127.5

# BUFFER は訓練データの数、 BATCH は一度に訓練するデータの数
BUFFER_SIZE = 60000
BATCH_SIZE = 256

# バッチごとにデータをシャッフルして排出することができるようになります
train_dataset = tf.data.Dataset.from_tensor_slices(train_images).shuffle(BUFFER_SIZE).batch(BATCH_SIZE)
```

### 生成器(Generator)

イメージとしてはこんな感じです
出力層の活性化関数のみ、 `tanh` になっています。
<img width="800" alt="aaa" src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/325895/ad297140-e6b7-b5a9-f4af-449dbc1e13d3.png">

```python
def make_generator_model():
    # keras では Sequential を使うことでニューラルネットワークを作成することができます
    # add で層を深くしていくことができます
    model = Sequential()
    # 全結合、重みも含めた1つの行列にあらわしています
    # (*, 100) の配列を受け取り、 (*, 12,544) の配列で出力するという意味です
    model.add(Dense(7 * 7 * 256, use_bias=False, input_shape=(100,)))
    model.add(BatchNormalization())
    # 通常の ReLU は正の値はそのまま出力して、負の値は0で出力します
    # LeakyReLU 負の値は0ではなく、小さな正の値を返します
    # 少しでも勾配をつけることで学習しやすくなります
    model.add(LeakyReLU())

    # (*, 12,544) → (*, 7, 7, 256) に変換
    model.add(Reshape((7, 7, 256)))
    assert model.output_shape == (None, 7, 7, 256)

    # (*, 7, 7, 256) → (*, 7, 7, 128) に変換
    # 転置畳み込み層を使っています
    # (5, 5) はカーネルサイズ
    # strides=(1, 1) は1ピクセルごとに計算する設定、 (2, 2) にすると2マス置きに計算する
    # padding='same' は入力と同じ長さで出力する設定、0埋めして計算する
    model.add(Conv2DTranspose(128, (5, 5), strides=(1, 1), padding='same', use_bias=False))
    assert model.output_shape == (None, 7, 7, 128)
    model.add(BatchNormalization())
    model.add(LeakyReLU())

    # (*, 7, 7, 128) → (*, 14, 14, 64) に変換
    model.add(Conv2DTranspose(64, (5, 5), strides=(2, 2), padding='same', use_bias=False))
    assert model.output_shape == (None, 14, 14, 64)
    model.add(BatchNormalization())
    model.add(LeakyReLU())

    # (*, 14, 14, 64) → (*, 28, 28, 1) に変換
    # sigmoid だと [0, 1] の範囲だが、 tanh だと [-1, 1] の範囲で出力され、よりはっきりとした画像になると見込めます
    model.add(Conv2DTranspose(1, (5, 5), strides=(2, 2), padding='same', use_bias=False, activation='tanh'))
    assert model.output_shape == (None, 28, 28, 1)

    return model
```

### 識別器(Discriminator)

こちらは一般的なCNN（畳み込みニューラスネットワーク）
<img width="600" alt="aaa" src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/325895/16b56283-90d0-34e1-1327-7db19f2d2c02.png">

```python
def make_discriminator_model():
    model = Sequential()
    # (*, 28, 28, 1) → (*, 14, 14, 64) に変換
    model.add(Conv2D(64, (5, 5), strides=(2, 2), padding='same', input_shape=[28, 28, 1]))
    model.add(LeakyReLU())
    # 約3割の情報を捨てる(0にする)
    # 過学習を抑制するため
    model.add(Dropout(0.3))

    # (*, 14, 14, 64) → (*, 7, 7, 128) に変換
    model.add(Conv2D(128, (5, 5), strides=(2, 2), padding='same'))
    model.add(LeakyReLU())
    model.add(Dropout(0.3))

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

    return model
```

### モデルの作成、最適化

```python
generator = make_generator_model()

# (1, 100) の平均0、標準偏差1のランダム配列を作成する
noise = tf.random.normal([1, 100])
# Discriminator の学習では Generator のパラメータは変更されないようにする
generated_image = generator(noise, training=False)

discriminator = make_discriminator_model()
decision = discriminator(generated_image)

cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True)

plt.imshow(generated_image[0, :, :, 0], cmap='gray')

# 最適化には Adam(Adaptive Moment Estination) を使用する
# 勾配降下法ベースで GAN ではよく利用される
generator_optimizer = tf.keras.optimizers.Adam(1e-4)
discriminator_optimizer = tf.keras.optimizers.Adam(1e-4)
```

### 学習ステップ

```python

def discriminator_loss(real_output, fake_output):
    # データセットのデータのラベルは 1
    real_loss = cross_entropy(tf.ones_like(real_output), real_output)
    # Generator が作成した偽物のデータのラベルは 0
    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)

@tf.function
def train_step(images):
    noise = tf.random.normal([BATCH_SIZE, noise_dim])

    with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
        # Generator にランダムノイズを入力して、偽画像を生成する
        generated_images = generator(noise, training=True)

        # データセットの画像の判別を行う
        real_output = discriminator(images, training=True)
        # 偽画像の判別を行う
        fake_output = discriminator(generated_images, training=True)

        # Generator が作成した偽画像を Discriminator がどれくらい誤判断をしたか
        gen_loss = generator_loss(fake_output)
        # Discriminator データセットの画像と偽画像をそれぞれどれくらい誤判断をしたか
        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))
```

## コード全体

手元の環境だと1エポック2分ほどで終わったので、2分 x 50エポックの2時間くらいで学習が完了しました。
Colab でも全然間に合う量だと思います。
最初は `EPOCHS` は5, 10くらいで動かしてみると良いかと

```python
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
import os
import time
from IPython import display
from keras.models import Sequential
from keras.layers import Activation, BatchNormalization, Dense, Dropout, Flatten, Reshape, LeakyReLU
from keras.layers.convolutional import Conv2D, Conv2DTranspose

(train_images, _), (_, _) = tf.keras.datasets.fashion_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

BUFFER_SIZE = 60000
BATCH_SIZE = 256

train_dataset = tf.data.Dataset.from_tensor_slices(train_images).shuffle(BUFFER_SIZE).batch(BATCH_SIZE)

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

    model.add(Reshape((7, 7, 256)))
    assert model.output_shape == (None, 7, 7, 256)

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

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

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

    return model

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

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

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

    return model

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)

@tf.function
def train_step(images):
    noise = tf.random.normal([BATCH_SIZE, noise_dim])

    with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
        generated_images = generator(noise, training=True)

        real_output = discriminator(images, 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))

def train(dataset, epochs):
    for epoch in range(epochs):
        start = time.time()

        for image_batch in dataset:
            train_step(image_batch)

        display.clear_output(wait=True)
        generate_and_save_images(generator, epoch + 1, seed)

        if (epoch + 1) % 15 == 0:
            checkpoint.save(file_prefix = checkpoint_prefix)

        print('Time for eopch {} is {} sec'.format(epoch + 1, time.time()-start))

    display.clear_output(wait=True)
    generate_and_save_images(generator, epochs, seed)

def generate_and_save_images(model, epoch, test_input):
  predictions = model(test_input, 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.savefig('image_at_epoch_{:04d}.png'.format(epoch))

generator = make_generator_model()

noise = tf.random.normal([1, 100])
generated_image = generator(noise, training=False)

discriminator = make_discriminator_model()
decision = discriminator(generated_image)

cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True)

plt.imshow(generated_image[0, :, :, 0], cmap='gray')

generator_optimizer = tf.keras.optimizers.Adam(1e-4)
discriminator_optimizer = tf.keras.optimizers.Adam(1e-4)

checkpoint_dir = './training_checkpoints'
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")
checkpoint = tf.train.Checkpoint(generator_optimizer=generator_optimizer,
                                discriminator_optimizer=discriminator_optimizer,
                                generator=generator,
                                discriminator=discriminator)

EPOCHS = 50
noise_dim = 100
num_examples_to_generate = 16

seed = tf.random.normal([num_examples_to_generate, noise_dim])

train(train_dataset, EPOCHS)
checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))
plt.show()
```

# おわりに

初学者として苦戦したので、少しでもチュートリアルが理解できるようになったら幸いです。
GANは他にも色々な種類があるので、もし興味を持った方は色々調べてみてください！

fashion MNIST を使って同じことをやってみたらこんな感じでした。
セーターが欠けていたりしますが、だいたい服だなとわかるくらいには学習できてますね。もう少しエポック数を増やしても良かったかもしれません
<img width="300" alt="aaa" src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/325895/37bb06a1-9bbb-5e5a-6a41-12a4c50b899f.png">