[Python] Keras로 DCGAN 구현하고 MNIST 이미지 생성하기

사실 정석의 DCGAN과는 차이(드롭아웃을 썼다거나)가 있지만 신경망 구성이야 쉽게 바꿀 수 있으니 틀만 보신다는 느낌으로 봐주세요.

참고로 제가 처음부터 짠 코드는 아니고 오픈소스를 좀 수정했습니다.

import numpy as np
from keras.models import *
from keras.layers import *
from keras.optimizers import *
from keras.datasets import mnist
import keras.backend as K
import matplotlib.pyplot as plt

K.set_image_data_format('channels_last')

class Gan:
  def __init__(self, img_data):
    img_size = img_data.shape[1]
    channel = img_data.shape[3] if len(img_data.shape) >= 4 else 1
    
    self.img_data = img_data
    self.input_shape = (img_size, img_size, channel)
    
    self.img_rows = img_size
    self.img_cols = img_size
    self.channel = channel
    self.noise_size = 100
    
    # Create D and G.
    self.create_d()
    self.create_g()
    
    # Build model to train D.
    optimizer = Adam(lr=0.0008)
    self.D.compile(loss='binary_crossentropy', optimizer=optimizer)
    
    # Build model to train G.
    optimizer = Adam(lr=0.0004)
    self.D.trainable = False
    self.AM = Sequential()
    self.AM.add(self.G)
    self.AM.add(self.D)
    self.AM.compile(loss='binary_crossentropy', optimizer=optimizer)
  
  def create_d(self):
    self.D = Sequential()
    depth = 64
    dropout = 0.4
    self.D.add(Conv2D(depth*1, 5, strides=2, input_shape=self.input_shape,
                      padding='same'))
    self.D.add(LeakyReLU(alpha=0.2))
    self.D.add(Dropout(dropout))
    self.D.add(Conv2D(depth*2, 5, strides=2, padding='same'))
    self.D.add(LeakyReLU(alpha=0.2))
    self.D.add(Dropout(dropout))
    self.D.add(Conv2D(depth*4, 5, strides=2, padding='same'))
    self.D.add(LeakyReLU(alpha=0.2))
    self.D.add(Dropout(dropout))
    self.D.add(Conv2D(depth*8, 5, strides=1, padding='same'))
    self.D.add(LeakyReLU(alpha=0.2))
    self.D.add(Dropout(dropout))
    self.D.add(Flatten())
    self.D.add(Dense(1))
    self.D.add(Activation('sigmoid'))
    self.D.summary()
    return self.D
  
  def create_g(self):
    self.G = Sequential()
    dropout = 0.4
    depth = 64+64+64+64
    dim = 7
    self.G.add(Dense(dim*dim*depth, input_dim=self.noise_size))
    self.G.add(BatchNormalization(momentum=0.9))
    self.G.add(Activation('relu'))
    self.G.add(Reshape((dim, dim, depth)))
    self.G.add(Dropout(dropout))
    self.G.add(UpSampling2D())
    self.G.add(Conv2DTranspose(int(depth/2), 5, padding='same'))
    self.G.add(BatchNormalization(momentum=0.9))
    self.G.add(Activation('relu'))
    self.G.add(UpSampling2D())
    self.G.add(Conv2DTranspose(int(depth/4), 5, padding='same'))
    self.G.add(BatchNormalization(momentum=0.9))
    self.G.add(Activation('relu'))
    self.G.add(Conv2DTranspose(int(depth/8), 5, padding='same'))
    self.G.add(BatchNormalization(momentum=0.9))
    self.G.add(Activation('relu'))
    self.G.add(Conv2DTranspose(1, 5, padding='same'))
    self.G.add(Activation('sigmoid'))
    self.G.summary()
    return self.G
  
  def train(self, batch_size=100):
    # Pick image data randomly.
    images_train = self.img_data[np.random.randint(0, self.img_data.shape[0], size=batch_size), :, :, :]
    
    # Generate images from noise.
    noise = np.random.uniform(-1.0, 1.0, size=[batch_size, self.noise_size])
    images_fake = self.G.predict(noise)
    
    # Train D.
    x = np.concatenate((images_train, images_fake))
    y = np.ones([2*batch_size, 1])
    y[batch_size:, :] = 0
    self.D.trainable = True
    d_loss = self.D.train_on_batch(x, y)
    
    # Train G.
    y = np.ones([batch_size, 1])
    noise = np.random.uniform(-1.0, 1.0, size=[batch_size, self.noise_size])
    self.D.trainable = False
    a_loss = self.AM.train_on_batch(noise, y)
    
    return d_loss, a_loss, images_fake

# Load MNIST dataset.
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train = x_train.astype('float32') / 255.0
x_train = x_train.reshape((x_train.shape[0],) + (28, 28, 1))
print(x_train.shape)

# Init network
gan = Gan(x_train)

# Some parameters.
epochs = 30
sample_size = 10
batch_size = 100
train_per_epoch = x_train.shape[0] // batch_size

for epoch in range(0, epochs):
  print("Epoch:", epoch + 1)
  
  total_d_loss = 0.0
  total_a_loss = 0.0
  
  for batch in range(0, train_per_epoch):
    d_loss, a_loss, imgs = gan.train(batch_size)
    total_d_loss += d_loss
    total_a_loss += a_loss
  
  total_d_loss /= train_per_epoch
  total_a_loss /= train_per_epoch
  
  print("D Loss: {}, AM Loss: {}".format(total_d_loss, total_a_loss))

  # Show generated images.
  fig, ax = plt.subplots(1, sample_size, figsize=(sample_size, 1))
  for i in range(0, sample_size):
    ax[i].set_axis_off()
    ax[i].imshow(imgs[i].reshape((28, 28)));
  plt.show()
  plt.close(fig);

결과:

뭐 그럭저럭 괜찮네요.

사실 원래 목적은 포켓몬 이미지를 생성하는 거였는데 계속 안돼서 MNIST로 회귀했습니다 ㅠㅠ

GAN이 딥러닝 모델? 중에서 가장 재미있는 것 같습니다 ㅎㅎ

두번째로는 Seq2Seq가 재밌는 것 같고요.

LAPGAN이란 것도 있다는데 고해상도 이미지 생성을 위한 GAN이라고 합니다.

주말에 공부해서 돌려볼 수 있으면 좋겠네요 ㅎ.. (근데 이해는 해도 컴퓨터 성능이 안좋은데 고화질까지 가면... ㅠ)