はじめに
コチラにGANについてまとめながら勉強しています。
アウトプット練習ということで、ポケモン画像を題材に
DCGANモデルとLightweight GANでそれぞれ新しいポケモンを生成してみました。
実装にあたって
・[今さら聞けないGAN （2）DCGANによる画像生成]
(https://qiita.com/triwave33/items/35b4adc9f5b41c5e8141)
・[GAN色々試してみた(元祖,DCGAN,CGAN,LightweightGAN)～AIが良い感じに仕事してくれる日を夢見て～]
(https://qiita.com/o93/items/96ca7dbcc82a8dd873ad)
・GANについて概念から実装まで ～DCGANによるキルミーベイベー生成～
を参考にさせて頂きました。ありがとうございます。
またどういうモデルなのかの説明、理論的な部分に関してはコチラの自分の記事に別でまとめていく予定なので、この記事では省略させて頂くことご了承ください。
実行環境
学習コストがかかるのでGoogle ColaboratoryのProに課金しました。
ローカルフォルダの扱いの利便性の低さや、制限はあるもののほぼ無限にGPUが使えて月額1072円は安すぎませんかね...。
GCPとかのクラウドサービスもいいけど、やっぱり常に**「何時間使ったから料金これくらいだな」っていう意識があって精神的にストレスなのでGoogle Colab Pro**に浮気しそうですね。
Google ColabのT4で回しています。（GPUリセマラめんどくさかった）
使ったデータセット
KaggleのPokemon Images Datasetを使わさせて頂きました。
合計819枚のポケモン画像がセットになっています。
ただComplete Pokemon Image Datasetの方がデータ数が多くていいかも。
ちなみに自分はポケモンの記憶がダイパまでで消えてるので、最新のポケモンわかんないです笑
まずはDCGANから
1.下準備
import numpy as np
import time
import os
import glob
import cv2
import matplotlib.pyplot as plt
import datetime as dt
import tensorflow as tf
from tensorflow.keras.models import Sequential, Model
from tensorflow.keras.layers import *
from tensorflow.keras.optimizers import Adam
必要なモジュールをインポートして
base_dir = './pokemon_img'
X_train = []
for image_path in glob.glob(base_dir + '/*'):
img = cv2.imread(image_path)
img_rgb = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
img_resize = cv2.resize(img_rgb , (128, 128))
numpy_img = np.array(img_resize)
numpy_img = numpy_img.reshape(128, 128, 3)
numpy_img = numpy_img.astype('float32')
numpy_img = (numpy_img - 127.5) / 127.5
X_train.append(numpy_img)
X_train = np.array(X_train)
コチラの記事でも扱ったようにndarray形式でTensorFlowに読み込める形にします。
このデータセットでは元々256×256サイズなのですが、さすがそのままのサイズで入れるのは学習の負担になるかなと思って半分にしています。
またDCGANではGeneratorネットワークの最後の活性化関数にtanh(双曲正接)を用いるため、生成される画像のスケールが[-1,1]になります。
そのためよくCNNモデルで用いられる正規化の手段の1つである255で割ることをしてしまうと生成される画像スケールは[-1,1]なのにDiscriminatorに与えられる画像スケールが[0,1]でズレてしまいます。
そのため
numpy_img = (numpy_img - 127.5) / 127.5
上記のように処理をしています。
画像スケールの変形はミスると「真っ黒い画像しか生成されんやんけ」とかってことになりがちなので気をつけたいところです。
2. Generator
def build_generator():
noise_shape = (z_size,)
model = Sequential(name='generator')
model.add(Dense(128 * 32 * 32, activation="relu", input_shape=noise_shape))
model.add(Reshape((32, 32, 128)))
model.add(UpSampling2D())
model.add(Conv2D(128, kernel_size=3, padding="same"))
model.add(BatchNormalization(momentum=0.8))
model.add(Activation("relu"))
model.add(UpSampling2D())
model.add(Conv2D(64, kernel_size=3, padding="same"))
model.add(BatchNormalization(momentum=0.8))
model.add(Activation("relu"))
model.add(Conv2D(3, kernel_size=3, padding="same"))
model.add(Activation("tanh"))
model.summary()
return model
ほぼ参考記事のまんまです。カラー画像なので最後の畳み込み層のフィルター数を3にしています。
コチラにもあるように
畳み込み→バッチ正規化→活性化関数の順番で構築すると良さそうです。
3. Discriminator
def build_discriminator():
model = Sequential(name='discriminator')
model.add(Conv2D(32, 5, strides=(2, 2), padding="same", input_shape=img_shape))
model.add(LeakyReLU())
model.add(Conv2D(128, 5, strides=(2, 2)))
model.add(LeakyReLU())
model.add(Dropout(0.25))
model.add(Conv2D(128, 5, strides = (2, 2)))
model.add(LeakyReLU())
model.add(GlobalAveragePooling2D())
model.add(Dense(256))
model.add(LeakyReLU())
model.add(Dropout(0.5))
model.add(Dense(1))
model.add(Activation("sigmoid"))
model.summary()
return model
これもほぼ参考コードのまんまです。
一応Dropout層をちょこちょこ入れて、過学習を抑えようとしているのと、全結合層の前にGAP層を入れることで計算コストを下げてます。
→GAPについてはコチラが分かりやすかったです。
4. 2つのモデルをつなげる
def build_combined():
discriminator.trainable = False
model = Sequential([generator, discriminator], name='Combined')
model.summary()
return model
Generatorを学習する時には、Discriminatorの学習をオフにします。
5. 各種パラメータの定義、モデルの立ち上げ
# 入力画像のサイズ
img_shape = X_train.shape[1:]
# ノイズの次元数
z_size = 100
# 最適化関数の定義
d_optimizer = Adam(lr=0.0001, beta_1=0.1)
g_optimizer = Adam(lr=0.0002, beta_1=0.5)
# Discriminatorモデルの生成・コンパイル
discriminator = build_discriminator()
discriminator.compile(loss="binary_crossentropy", optimizer=d_optimizer, metrics=["accuracy"])
# Generatorモデルの生成・コンパイル
generator = build_generator()
# ネットワーク作成
combined_model = build_combined()
combined_model.compile(loss="binary_crossentropy", optimizer=g_optimizer)
Model: "discriminator"
_________________________________________________________________
Layer (type) Output Shape Param #
=================================================================
conv2d_19 (Conv2D) (None, 64, 64, 32) 2432
_________________________________________________________________
leaky_re_lu_12 (LeakyReLU) (None, 64, 64, 32) 0
_________________________________________________________________
conv2d_20 (Conv2D) (None, 30, 30, 128) 102528
_________________________________________________________________
leaky_re_lu_13 (LeakyReLU) (None, 30, 30, 128) 0
_________________________________________________________________
dropout_6 (Dropout) (None, 30, 30, 128) 0
_________________________________________________________________
conv2d_21 (Conv2D) (None, 13, 13, 128) 409728
_________________________________________________________________
leaky_re_lu_14 (LeakyReLU) (None, 13, 13, 128) 0
_________________________________________________________________
global_average_pooling2d_3 ( (None, 128) 0
_________________________________________________________________
dense_9 (Dense) (None, 256) 33024
_________________________________________________________________
leaky_re_lu_15 (LeakyReLU) (None, 256) 0
_________________________________________________________________
dropout_7 (Dropout) (None, 256) 0
_________________________________________________________________
dense_10 (Dense) (None, 1) 257
_________________________________________________________________
activation_12 (Activation) (None, 1) 0
=================================================================
Total params: 547,969
Trainable params: 547,969
Non-trainable params: 0
_________________________________________________________________
Model: "generator"
_________________________________________________________________
Layer (type) Output Shape Param #
=================================================================
dense_11 (Dense) (None, 131072) 13238272
_________________________________________________________________
reshape_3 (Reshape) (None, 32, 32, 128) 0
_________________________________________________________________
up_sampling2d_6 (UpSampling2 (None, 64, 64, 128) 0
_________________________________________________________________
conv2d_22 (Conv2D) (None, 64, 64, 128) 147584
_________________________________________________________________
batch_normalization_6 (Batch (None, 64, 64, 128) 512
_________________________________________________________________
activation_13 (Activation) (None, 64, 64, 128) 0
_________________________________________________________________
up_sampling2d_7 (UpSampling2 (None, 128, 128, 128) 0
_________________________________________________________________
conv2d_23 (Conv2D) (None, 128, 128, 64) 73792
_________________________________________________________________
batch_normalization_7 (Batch (None, 128, 128, 64) 256
_________________________________________________________________
activation_14 (Activation) (None, 128, 128, 64) 0
_________________________________________________________________
conv2d_24 (Conv2D) (None, 128, 128, 3) 1731
_________________________________________________________________
activation_15 (Activation) (None, 128, 128, 3) 0
=================================================================
Total params: 13,462,147
Trainable params: 13,461,763
Non-trainable params: 384
_________________________________________________________________
Model: "Combined"
_________________________________________________________________
Layer (type) Output Shape Param #
=================================================================
generator (Sequential) (None, 128, 128, 3) 13462147
_________________________________________________________________
discriminator (Sequential) (None, 1) 547969
=================================================================
Total params: 14,010,116
Trainable params: 13,461,763
Non-trainable params: 548,353
_________________________________________________________________
コチラにも述べられているようにGeneratorとDiscriminatorの最適化アルゴリズムは両方ともにAdamを用いるわけですが、全く一緒にはせずに学習率と減衰率beta1は変えています。
パラメータ数多いですね、、、
6. いざ学習
# 学習結果の表示・保存をする関数を定義
def save_imgs(log_path, epoch):
r, c = 5, 5
noise = np.random.normal(0, 1, (r * c, z_size))
gen_imgs = generator.predict(noise)
# [0,1]スケールに変更する
gen_imgs = 0.5 * gen_imgs + 0.5
fig, axs = plt.subplots(r, c)
cnt = 0
for i in range(r):
for j in range(c):
axs[i,j].imshow(gen_imgs[cnt, :,:,:])
axs[i,j].axis('off')
cnt += 1
fig.savefig("{}/{}.png".format(log_path, epoch))
plt.show()
plt.close()
def train(epochs, batch_size=64, save_interval=1):
half_batch = int(batch_size / 2)
num_batches = int(X_train.shape[0] / half_batch)
print("Number of Batches : ", num_batches)
log_path = 'log/{}/images'.format(dt.datetime.now().strftime("%Y-%m-%d_%H%M%S"))
os.makedirs(log_path, exist_ok=True)
for epoch in range(epochs):
start_time = time.time()
for iteration in range(num_batches):
# NoiseからGeneratorで生成
noise = np.random.normal(0, 1, (half_batch, z_size))
gen_imgs = generator.predict(noise)
# データセットから画像をピックアップ
idx = np.random.randint(0, X_train.shape[0], half_batch)
imgs = X_train[idx]
# それぞれのデータでDiscriminatorを学習
d_loss_real = discriminator.train_on_batch(imgs, np.ones((half_batch, 1)))
d_loss_fake = discriminator.train_on_batch(gen_imgs, np.zeros((half_batch, 1)))
# DiscriminatorのLossを算出
d_loss = np.add(d_loss_real, d_loss_fake) / 2
# ノイズ生成
noise = np.random.normal(0, 1, (batch_size, z_size))
# 騙すことが正解になる目的変数
valid_y = np.array([1] * batch_size)
#Generatorを学習
g_loss = combined_model.train_on_batch(noise, valid_y)
if epoch % save_interval == 0:
# 生成画像の表示と保存
per_1_time = time.time() - start_time
print ("epoch:%d, iter:%d, Time:%2f seconds, End: %2f min, [D loss: %f, acc.: %.2f%%] [G loss: %f]" % (epoch, iteration,per_1_time, (epochs- epoch)*per_1_time/60, d_loss[0], 100*d_loss[1], g_loss))
save_imgs(log_path, epoch)
train(epochs = 500)
save_img関数ですが
gen_imgs = generator.predict(noise)
gen_imgs = 0.5 * gen_imgs + 0.5
generatorから生成された画像は先ほども伝えたようにtanhで活性化しているため[-1,1]スケールになってます。またdtypeはfloat32です。
これをそのままmatplotlibのimshowメソッドで表示するとうまく表示されません。もしそのまま突っ込むと
Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers).
おそらくこうしたエラーが出ます。
要は**「float型でRGBデータ（カラー画像）をちゃんと表示したいんやったら[0,1]スケール、もしくはデータ型をintegers(整数値)にして[0,255]スケールにしてくれや」**ってことです。
今回はまず半分にして0.5を足すことで[0,1]スケールにしています。
(gen_imgs * 127.5 + 127.5).astype(uint8)
とかでもいいと思います。
今回はtrain_on_batchでバッチごとに学習と損失の計算をしています。
tf.GradientTapeで勾配を計算するやり方も記事にできたらと思います。
また学習においては、コチラを参考にミニバッチをさらに半分に分けて、本物と偽物を混ぜて学習しないようにしてます。
そうして500エポックで学習した記録を載せます。
・1エポック
・10エポック
・30エポック
失敗か...?
・50エポック
・・・・・・・。
・100エポック
なんかそれっぽいやつが現れた
・200エポック
ポケモンのラフ画に見えなくもない
・300エポック
同じような形のポケモンが目立つ
・400エポック
ドラゴンっぽい
・500エポック
学習時間は大体1時間半ほど。
D-lossとG-lossも安定していたので、エポック数増やせばもっと鮮明なポケモンが生まれるかも。
LightweightGANでやってみた
今年登場した、高性能ながら軽量（でも数時間の学習は必要）のモデルです。
pipで簡単に落とせるので気軽に試せます。
# lightweight_ganをインストール
!pip install lightweight-gan
!lightweight_gan \
--data ./data \
--name 'pokemon' \
--batch-size 16 \
--gradient-accumulate-every 4 \
--num-train-steps 15000
使い方は公式Githubにあります。
これで生成できた画像を載せます
・1000エポック
それっぽいのがもう出現してる
・2000エポック
もうゲームとかに普通にいそう笑
・5000エポック
草タイプ多いな
・10000エポック
似た形のポケモンも多いですね。色違いかな？
・15000エポック
赤の着色が目立つポケモン
ここまでで学習時間が7時間くらいだった気がします。
当たり前ですけど、自前のDCGANモデルより高画質ですね（ポケモンに見えるかは置いといて）
手軽に高解像度の画像が生成できるのは興味が惹かれますね。
このモデルのソースコードがPytorchで書かれているので、Pytorchの勉強もして読み砕いで、TensorFlowで書き直してみたいです。
まとめ
自分で（ほぼ参考通り）実装したGANモデルと、すでに完成されている高性能なGANモデルの2つで新しいポケモンを生成してみました。
やはり言われているように学習はかなり難しいですね。層の構成、Optimizerの学習率など調整するところが多くて大変な印象でした。
ただその分、ただのノイズから画像が生成されるのは非常に興味深いです。
次はCGAN（ラベル情報もノイズに加えて入力することで生成する画像のコントロールができる）モデルを構築してみたいと思います。