[Keras] GAN으로 MNIST 이미지 생성하기

※ 이 글은 '코딩셰프의 3분 딥러닝 케라스맛'이라는 책을 보고 실습한걸 기록한 글입니다.

GAN으로 MNIST 이미지를 생성하는 예제입니다.

책의 코드에서 뺀 부분이 많습니다.

코드:

from keras import layers, models, optimizers
from keras import datasets
from keras import backend as K
import matplotlib.pyplot as plt
import numpy as np
 
K.set_image_data_format('channels_first')
print(K.image_data_format())
 
class GAN(models.Sequential):
  def __init__(self, input_dim):
    super().__init__()
    
    self.input_dim = input_dim
    
    self.generator = self.make_G()
    self.discriminator = self.make_D()
    
    self.add(self.generator)
    self.discriminator.trainable = False
    self.add(self.discriminator)
    
    self.compile_all()
    
  def make_G(self):
    input_dim = self.input_dim
    
    model = models.Sequential()
    model.add(layers.Dense(1024, activation='tanh', input_dim=input_dim))
    model.add(layers.Dense(128 * 7 * 7, activation='tanh'))
    model.add(layers.BatchNormalization())
    model.add(layers.Reshape((128, 7, 7), input_shape=(128 * 7 * 7,)))
    model.add(layers.UpSampling2D(size=(2, 2)))
    model.add(layers.Conv2D(64, (5, 5), padding='same', activation='tanh'))
    model.add(layers.UpSampling2D(size=(2, 2)))
    model.add(layers.Conv2D(1, (5, 5), padding='same', activation='tanh'))
    return model
  
  def make_D(self):
    model = models.Sequential()
    model.add(layers.Conv2D(64, (5, 5), padding='same', activation='tanh',
                           input_shape=(1, 28, 28)))
    model.add(layers.MaxPooling2D(pool_size=(2, 2)))
    model.add(layers.Conv2D(128, (5, 5), padding='same', activation='tanh'))
    model.add(layers.MaxPooling2D(pool_size=(2, 2)))
    model.add(layers.Flatten())
    model.add(layers.Dense(1024, activation='tanh'))
    model.add(layers.Dense(1, activation='sigmoid'))
    return model
  
  def compile_all(self):
    opt_D = optimizers.SGD(lr=0.0005, momentum=0.9, nesterov=True)
    opt_G = optimizers.SGD(lr=0.0005, momentum=0.9, nesterov=True)
    
    self.compile(loss='binary_crossentropy', optimizer=opt_G)
    
    self.discriminator.trainable = True
    self.discriminator.compile(loss='binary_crossentropy', optimizer=opt_D)
    
  def get_z(self, ln):
    return np.random.uniform(-1.0, 1.0, (ln, self.input_dim))
  
  def train_once(self, x):
    ln = x.shape[0]
    
    z = self.get_z(ln)
    gen = self.generator.predict(z, verbose=0)
    input_D = np.concatenate((x, gen))
    y_D = [1] * ln + [0] * ln
    loss_D = self.discriminator.train_on_batch(input_D, y_D)
    
    z = self.get_z(ln)
    self.discriminator.trainable = False
    loss_G = self.train_on_batch(z, [1] * ln)
    self.discriminator.trainable = True
    
    return loss_D, loss_G
  
def get_x(x_train, index, batch_size):
  return x_train[index * batch_size:(index + 1) * batch_size]
 
class MnistData():
    def __init__(self):
        (x_train, y_train), (x_test, y_test) = datasets.mnist.load_data()
 
        img_rows, img_cols = x_train.shape[1:]
 
        x_train = x_train.astype('float32') - 127.5
        x_test = x_test.astype('float32') - 127.5
        x_train /= 127.5
        x_test /= 127.5
 
        self.num_classes = 10
        self.x_train, self.y_train = x_train, y_train
        self.x_test, self.y_test = x_test, y_test
  
def main():
  batch_size = 100
  epochs = 30
  input_dim = 100
  sample_size = 6
  
  data = MnistData()
  x_train = data.x_train
  x_train = x_train.reshape((x_train.shape[0], 1) + x_train.shape[1:])
  
  gan = GAN(input_dim)
  
  for epoch in range(epochs):
    print("Epoch", epoch)
    
    for index in range(int(x_train.shape[0] / batch_size)):
      x = get_x(x_train, index, batch_size)
      loss_D, loss_G = gan.train_once(x)
      
    print('Loss D:', loss_D)
    print('Loss G:', loss_G)
      
    if epoch % 2 == 0 or epoch == epochs - 1:
      z = gan.get_z(sample_size)
      gen = gan.generator.predict(z, verbose=0)
      
      plt.figure(figsize=(20, 2))
  
      for i in range(sample_size):
        ax = plt.subplot(1, sample_size, i + 1)
        plt.imshow(gen[i].reshape((28, 28)))
        ax.get_xaxis().set_visible(False)
        ax.get_yaxis().set_visible(False)
        
      plt.show()
 
if __name__ == '__main__':
  main()

결과:

Using TensorFlow backend.
channels_first

Epoch 0
Loss D: 0.7338002
Loss G: 0.93040544

(중략)

Epoch 9
Loss D: 0.35202882
Loss G: 2.2948298
Epoch 10
Loss D: 0.46802875
Loss G: 1.4171509

(중략)

Epoch 29
Loss D: 0.2630578
Loss G: 1.8819554

생성망을 만드는데 새로운 레이어가 등장했습니다.

BatchNormalization 레이어입니다.

각 배치의 입력마다 평균을 0, 편차를 1로 변환해서 출력해주는 레이어입니다.

오버피팅을 방지하는 등 드롭아웃과 비슷한 효과를 낸다고 합니다.

생성망의 전체적인 구조는 별다를게 없습니다.

노이즈 입력을 변환해서 업스케일링과 합성곱을 거쳐 이미지를 생성합니다.

판별기는 전형적인 CNN 모습이니 넘어갑니다.

다음은 컴파일인데 책에서는 생성망도 compile 메소드를 호출해줬지만 생성망 자체는 train 작업이 없습니다.

생성망과 판별망을 합쳐서 train하기 때문이죠.

그래서 그 부분은 빼고 했는데 잘 되는거 같습니다.

이전 글에서 trainable 속성 설명은 했으니 넘어가고

최적화기는 SGD(Stochastic Gradient Descent)를 썼네요.

Loss를 계산할때 전체 Batch가 아니라 Mini-Batch를 가지고 계산한다고 합니다.

get_z 메소드는 생성망의 입력으로 쓸 노이즈 데이터를 생성합니다.

train_once 메소드에서 판별망과 생성망+판별망을 차례로 한번씩 학습해주고 있습니다.

이때 판별망의 학습 데이터는 (Input: 실제 이미지*n + 생성된 가짜 이미지*n, Target: 진짜*n+가짜*n) 입니다.

생성망은 판별망이 진짜(1)라고 출력하게 해야하므로 학습 데이터는 (Input: 노이즈 데이터*n, Target: 진짜*n) 입니다.

그 다음은 별거 없습니다.

끝!