[tensorflow GAN]tensorflow GAN(심층 합성곱 생성적 적대 신경망) 기본예제

tensorflow사이트에서 제공하는 GAN예제를 따라해본것입니다.

더욱 자세한 내용은 밑의 주소에서 확인하시면 됩니다.

www.tensorflow.org/tutorials/generative/dcgan?hl=ko

심층 합성곱 생성적 적대 신경망  |  TensorFlow Core

 

 

생성적 적대 신경망(GAN)은 두개의 모델이 적대적인 과정을 통해 동시에 훈련이 이루어집니다.

생성자는 진짜처럼 보이는 이미지를 생성하도록 배우는 와중에, 감별자가 가짜의 이미지로부터 진짜를 구별하게 되는 것을 배우게 됩니다.

훈련과정을 반복하다보면 생성자는 점차 실제같은 이미지를 더 잘 생성하게 되고, 감별자의 경우는 점차 진짜와 가짜를 더 잘 구별하게 됩니다. 이 과정은 감별자가 가짜 이미지에서 진짜 이미지를 더 이상 구별하지 못하게 될때 , 평형상태에 도달하게 됩니다.

 

 

이 기본예제를 끝까지 진행하시면 다음과 같은 결과를 얻으 실 수 있습니다. (50epoch만큼 학습을 진행합니다.)

이 기본예제를 진행하기에 앞서 텐서플로우와 imageio를 설치하셔야 합니다.

텐서플로우에서 제공하는 버전 중 이 예제는 2.0.0-rc1버전을 사용합니다.

imageio의 경우는 pip를 이용하여 설치하시면 됩니다.

 

pip install -q tensorflow==2.0.0-rc1
pip install -q imageio

 

 

먼저 필요한 라이브러리들을 import 시켜줍니다. 저의경우는 주피터 노트북에서 진행하였습니다.

 

import glob
import imageio
import matplotlib.pyplot as plt
import numpy as np
import os
import PIL
from tensorflow.keras import layers
import time

from IPython import display

 

라이브러리들이 import되었으면, tensorflow에서 제공하는 데이터셋을 로딩과 준비를 합니다.

 

(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 # 이미지를 [-1, 1]로 정규화합니다.

BUFFER_SIZE = 60000
BATCH_SIZE = 256

# 데이터 배치를 만들고 섞습니다.
train_dataset = tf.data.Dataset.from_tensor_slices(train_images).shuffle(BUFFER_SIZE).batch(BATCH_SIZE)

이제 준비가 끝나면 모델을 만들어줍니다. 생성자와 감별자는 keras Sequential API를 이용해서 정의가 됩니다.

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) # 주목: 배치사이즈로 None이 주어집니다.

    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

generator = make_generator_model()

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

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

위 이미지는 아직 훈련이 되지않은 생성자를 이용해서 만든 이미지입니다.

밑의 코드는 감별자입니다. 감별자는 CNN기반의 이미지 분류기를 사용합니다.

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

discriminator = make_discriminator_model()
decision = discriminator(generated_image)
print (decision)

# 이 메서드는 크로스 엔트로피 손실함수 (cross entropy loss)를 계산하기 위해 헬퍼 (helper) 함수를 반환합니다.
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)

체크포인트 저장

 

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

# 이 시드를 시간이 지나도 재활용하겠습니다. 
# (GIF 애니메이션에서 진전 내용을 시각화하는데 쉽기 때문입니다.) 
seed = tf.random.normal([num_examples_to_generate, noise_dim])

# `tf.function`이 어떻게 사용되는지 주목해 주세요.
# 이 데코레이터는 함수를 "컴파일"합니다.
@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)

    # GIF를 위한 이미지를 바로 생성합니다.
    display.clear_output(wait=True)
    generate_and_save_images(generator,
                             epoch + 1,
                             seed)

    # 15 에포크가 지날 때마다 모델을 저장합니다.
    if (epoch + 1) % 15 == 0:
      checkpoint.save(file_prefix = checkpoint_prefix)

    # print (' 에포크 {} 에서 걸린 시간은 {} 초 입니다'.format(epoch +1, time.time()-start))
    print ('Time for epoch {} 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):
  # `training`이 False로 맞춰진 것을 주목하세요.
  # 이렇게 하면 (배치정규화를 포함하여) 모든 층들이 추론 모드로 실행됩니다. 
  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))
  plt.show()

 

이제 훈련을 시켜주면 됩니다. 

%%time
train(train_dataset, EPOCHS)

훈련이 끝나게 되면 위의 이미지처럼 결과가 나오게 됩니다.

시간이 오래걸리므로 다른 작업을 하시거나 휴식을 하면 될 것 같습니다.

 

훈련이 끝나 위의 이미지처럼 결과까지 얻으셨다면 마지막 체크포인트를 복구하고, gif를 생성하면서 마무리하시면 됩니다.

 

checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))

# 에포크 숫자를 사용하여 하나의 이미지를 보여줍니다.
def display_image(epoch_no):
  return PIL.Image.open('image_at_epoch_{:04d}.png'.format(epoch_no))

anim_file = 'dcgan.gif'

with imageio.get_writer(anim_file, mode='I') as writer:
  filenames = glob.glob('image*.png')
  filenames = sorted(filenames)
  last = -1
  for i,filename in enumerate(filenames):
    frame = 2*(i**0.5)
    if round(frame) > round(last):
      last = frame
    else:
      continue
    image = imageio.imread(filename)
    writer.append_data(image)
  image = imageio.imread(filename)
  writer.append_data(image)

import IPython
if IPython.version_info > (6,2,0,''):
  display.Image(filename=anim_file)

만약에 colab에서 작업하고있으시다면 아래의 코드로 이미지를 다운받으시면 됩니다.

try:
  from google.colab import files
except ImportError:
  pass
else:
  files.download(anim_file)