BEGAN을 이용한 이미지생성

라이브러리 읽어들이기

In [0]:
#https://drive.google.com/uc?export=download&id=1URxZOJTO38qb3v1hOT3usuAE7nTGQDnL
import os
import numpy as np
from tensorflow.python import keras
from tensorflow.python.keras import backend as K
from tensorflow.python.keras import losses
from tensorflow.python.keras.optimizers import Adam
from tensorflow.python.keras.models import Sequential, Model
from tensorflow.python.keras.layers import Conv2D, Flatten, Dense, UpSampling2D, Reshape, Lambda, Input, Dropout
from tensorflow.python.keras.preprocessing.image import ImageDataGenerator
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from keras.datasets import mnist

구글 드라이브 연동하기

In [2]:
from google.colab import drive
drive.mount('/content/drive')
Go to this URL in a browser: https://accounts.google.com/o/oauth2/auth?client_id=947318989803-6bn6qk8qdgf4n4g3pfee6491hc0brc4i.apps.googleusercontent.com&redirect_uri=urn%3aietf%3awg%3aoauth%3a2.0%3aoob&response_type=code&scope=email%20https%3a%2f%2fwww.googleapis.com%2fauth%2fdocs.test%20https%3a%2f%2fwww.googleapis.com%2fauth%2fdrive%20https%3a%2f%2fwww.googleapis.com%2fauth%2fdrive.photos.readonly%20https%3a%2f%2fwww.googleapis.com%2fauth%2fpeopleapi.readonly

Enter your authorization code:
··········
Mounted at /content/drive

생성된 이미지 보여주기

In [0]:
# 생성된 MNIST 이미지를 rowsxcols Grid로 보여주는 plot 함수 정의
def plot_imgs(path, imgs, rows, cols):

    fig = plt.figure(figsize=(rows, cols))
    fig.set_figheight(3)
    fig.set_figwidth(3)
    gs = gridspec.GridSpec(rows, cols)
    gs.update(wspace=0.05, hspace=0.05)
   
    for i, img in enumerate(imgs):
        ax = plt.subplot(gs[i])
        plt.axis('off')
        ax.set_xticklabels([])
        ax.set_yticklabels([])
        ax.set_aspect('equal')  
        
        plt.imshow(img.reshape((28, 28)), cmap='gray') 

    return fig

이미지 데이터 읽어 들이기

In [0]:
BATCH_SIZE = 64
IMG_SHAPE = (28, 28, 1)

# mnist 데이터 셋 불러옴 
(X_train, __), (__, __) = mnist.load_data()

X_train = X_train.astype('float32')
X_train = np.expand_dims(X_train, axis=3)

data_gen = ImageDataGenerator(rescale=1/255.)  # 이미지 전처리(Rescale 0 to 1)
train_data_generator = data_gen.flow(X_train, batch_size=BATCH_SIZE)
In [6]:
plt.imshow(X_train[0,:,:,0], cmap='gray')
Out[6]:
<matplotlib.image.AxesImage at 0x7fe1f2fe8e80>

Encoder 정의

In [0]:
def build_encoder(input_shape, z_size, n_filters, n_layers):
    """Encoder구축
    
    Arguments:
        input_shape (int): 이미지의 shape
        z_size (int): 특징 공간의 차원 수
        n_filters (int): 필터 수
 
    """
    model = Sequential()
    model.add(Conv2D(n_filters, 3, activation='elu',input_shape=input_shape, padding='same'))
    model.add(Conv2D(n_filters, 3, padding='same'))
    for i in range(2, n_layers + 1):
        model.add(Conv2D(i*n_filters, 3, activation='elu', padding='same'))
        model.add(Conv2D(i*n_filters, 3, activation='elu', padding='same'))
  
    model.add(Conv2D(n_layers*n_filters, 3, padding='same'))
    model.add(Flatten())
    model.add(Dense(z_size))
    model.summary()
    return model

생성자(Generator)/Decoder 정의

In [0]:
def build_decoder(output_shape, z_size, n_filters, n_layers):
    """Decoder 구축
    
    Arguments:
        output_shape (np.array): 이미지 shape
        z_size (int): 특징 공간의 차원 수
        n_filters (int): 필터 수
        n_layers (int): 레이어 수

    """
    # UpSampling2D로 몇 배로 확대할지 계산
    scale = 2**(n_layers - 1)
    # 합성곱층의 처음 입력 사이즈를 scale로부터 역산
    fc_shape = (output_shape[0]//scale, output_shape[1]//scale, n_filters )
    # 완전연결 계층에서 필요한 사이즈를 역산
    fc_size = fc_shape[0]*fc_shape[1]*fc_shape[2]
    
    model = Sequential()
    # 완전연결 계층
    model.add(Dense(fc_size, input_shape=(z_size,)))
    model.add(Reshape(fc_shape))
    
    # 합성곱층 반복
    for i in range(n_layers - 1):
        model.add(Conv2D(n_filters, 3, activation='elu', padding='same'))
        model.add(Conv2D(n_filters, 3, activation='elu', padding='same'))
        model.add(UpSampling2D())
        
    # 마지막 층은 UpSampling2D가 불필요 
    model.add(Conv2D(n_filters, 3, activation='elu', padding='same'))
    model.add(Conv2D(n_filters, 3, activation='elu', padding='same'))
    # 출력층에서는 1채널로
    model.add(Conv2D(1, 3, padding='same'))
    
    return model

생성자(Generator) 정의

In [0]:
def build_generator(img_shape, z_size, n_filters, n_layers):
    decoder = build_decoder(img_shape, z_size, n_filters, n_layers)
    return decoder

구분자(Discriminator) 정의

In [0]:
def build_discriminator(img_shape, z_size, n_filters, n_layers):
    encoder = build_encoder(img_shape, z_size, n_filters, n_layers)
    decoder = build_decoder(img_shape, z_size, n_filters, n_layers)
    return keras.models.Sequential((encoder, decoder))

구분자(Discriminator) 학습용 네트워크

In [0]:
def build_discriminator_trainer(discriminator):
    img_shape = discriminator.input_shape[1:]
    real_inputs = Input(img_shape)
    fake_inputs = Input(img_shape)
    real_outputs = discriminator(real_inputs)
    fake_outputs = discriminator(fake_inputs)

    return Model(
        inputs=[real_inputs, fake_inputs],
        outputs=[real_outputs, fake_outputs]
    )

네트워크 구축

In [87]:
n_filters = 64 #  필터 수
n_layers = 3 # 레이어 수
z_size = 32  #  특징 공간의 차원

generator = build_generator(IMG_SHAPE, z_size, n_filters, n_layers)
discriminator = build_discriminator(IMG_SHAPE, z_size, n_filters, n_layers)
discriminator_trainer = build_discriminator_trainer(discriminator)

generator.summary()

# discriminator.layers[1]은 디코더를 나타냄
discriminator.layers[1].summary()
Model: "sequential_57"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
conv2d_253 (Conv2D)          (None, 28, 28, 64)        640       
_________________________________________________________________
conv2d_254 (Conv2D)          (None, 28, 28, 64)        36928     
_________________________________________________________________
conv2d_255 (Conv2D)          (None, 28, 28, 128)       73856     
_________________________________________________________________
conv2d_256 (Conv2D)          (None, 28, 28, 128)       147584    
_________________________________________________________________
conv2d_257 (Conv2D)          (None, 28, 28, 192)       221376    
_________________________________________________________________
conv2d_258 (Conv2D)          (None, 28, 28, 192)       331968    
_________________________________________________________________
conv2d_259 (Conv2D)          (None, 28, 28, 192)       331968    
_________________________________________________________________
flatten_14 (Flatten)         (None, 150528)            0         
_________________________________________________________________
dense_43 (Dense)             (None, 32)                4816928   
=================================================================
Total params: 5,961,248
Trainable params: 5,961,248
Non-trainable params: 0
_________________________________________________________________
Model: "sequential_56"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
dense_42 (Dense)             (None, 3136)              103488    
_________________________________________________________________
reshape_28 (Reshape)         (None, 7, 7, 64)          0         
_________________________________________________________________
conv2d_246 (Conv2D)          (None, 7, 7, 64)          36928     
_________________________________________________________________
conv2d_247 (Conv2D)          (None, 7, 7, 64)          36928     
_________________________________________________________________
up_sampling2d_40 (UpSampling (None, 14, 14, 64)        0         
_________________________________________________________________
conv2d_248 (Conv2D)          (None, 14, 14, 64)        36928     
_________________________________________________________________
conv2d_249 (Conv2D)          (None, 14, 14, 64)        36928     
_________________________________________________________________
up_sampling2d_41 (UpSampling (None, 28, 28, 64)        0         
_________________________________________________________________
conv2d_250 (Conv2D)          (None, 28, 28, 64)        36928     
_________________________________________________________________
conv2d_251 (Conv2D)          (None, 28, 28, 64)        36928     
_________________________________________________________________
conv2d_252 (Conv2D)          (None, 28, 28, 1)         577       
=================================================================
Total params: 325,633
Trainable params: 325,633
Non-trainable params: 0
_________________________________________________________________
Model: "sequential_58"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
dense_44 (Dense)             (None, 3136)              103488    
_________________________________________________________________
reshape_29 (Reshape)         (None, 7, 7, 64)          0         
_________________________________________________________________
conv2d_260 (Conv2D)          (None, 7, 7, 64)          36928     
_________________________________________________________________
conv2d_261 (Conv2D)          (None, 7, 7, 64)          36928     
_________________________________________________________________
up_sampling2d_42 (UpSampling (None, 14, 14, 64)        0         
_________________________________________________________________
conv2d_262 (Conv2D)          (None, 14, 14, 64)        36928     
_________________________________________________________________
conv2d_263 (Conv2D)          (None, 14, 14, 64)        36928     
_________________________________________________________________
up_sampling2d_43 (UpSampling (None, 28, 28, 64)        0         
_________________________________________________________________
conv2d_264 (Conv2D)          (None, 28, 28, 64)        36928     
_________________________________________________________________
conv2d_265 (Conv2D)          (None, 28, 28, 64)        36928     
_________________________________________________________________
conv2d_266 (Conv2D)          (None, 28, 28, 1)         577       
=================================================================
Total params: 325,633
Trainable params: 325,633
Non-trainable params: 0
_________________________________________________________________

손실(loss) 함수 정의

In [0]:
from tensorflow.python.keras.losses import mean_absolute_error

def build_generator_loss(discriminator):
    # discriminator를 사용해서 손실 함수 정의
    def loss(y_true, y_pred):
        # y_true는 더미
        reconst = discriminator(y_pred)
        return mean_absolute_error(reconst, y_pred)
    return loss

generator 컴파일

In [0]:
# 초기 학습률(Generator)
g_lr = 0.0001
generator_loss = build_generator_loss(discriminator)
generator.compile(loss=generator_loss, optimizer=Adam(g_lr))

discriminator 컴파일

In [0]:
# 초기 학습률(Discriminator)
# k_var는 수치(일반 변수)
k_var = 0.0
# k : Keras(TensorFlow) Variable
k = K.variable(k_var)

d_lr = 0.0001

discriminator_trainer.compile(loss=[ mean_absolute_error, mean_absolute_error],loss_weights=[1., -k], optimizer=Adam(d_lr))

수렴 판정용 함수 정의

In [0]:
def measure(real_loss, fake_loss, gamma):
    return real_loss + np.abs(gamma*real_loss - fake_loss)

학습 코드

In [0]:
# k의 갱신에 이용할 파라미터
GAMMA =  1
Lambda = 0.001

# 반복 수. 100000~1000000 정도로 지정
TOTAL_STEPS = 1000

# 모델과 확인용 생성 이미지를 저장할 폴더
IMG_SAVE_DIR = '/content/drive/My Drive/data/imgs'
# 확인용으로 4x4 개의 이미지를 생성
IMG_SAMPLE_SHAPE = (4, 4)
N_IMG_SAMPLES = np.prod(IMG_SAMPLE_SHAPE)

# 저장할 폴더가 없다면 생성
os.makedirs(IMG_SAVE_DIR, exist_ok=True)

# 샘플이미지용 랜덤 시드
sample_seeds = np.random.uniform(-1, 1, (N_IMG_SAMPLES, z_size))

history = []
logs = []

for epoch, batch in enumerate(train_data_generator): 
    # 학습 종료 
    if epoch > TOTAL_STEPS:
      break
    #임의의 값(noise) 생성, 잠재변수의 input으로 사용할 noise를 균등분포(Uniform Distribution)에서 BATCH_SIZE만큼 샘플링
    z_g = np.random.uniform(-1, 1, (BATCH_SIZE, z_size))  # 균등 분포 -1과 1사이에 랜덤값 추출
    z_d = np.random.uniform(-1, 1, (BATCH_SIZE, z_size))
    # 생성 이미지(구분자의 학습에 이용), z_g 입력받아 가짜 이미지 생성
    g_pred = generator.predict(z_d)
    
    # 생성자를 1스텝 학습시킨다
    generator.train_on_batch(z_g, batch)
    # discriminator 1스텝 학습시킨다
    _, real_loss, fake_loss = discriminator_trainer.train_on_batch([batch, g_pred],[batch, g_pred]) 

    # k 를 갱신, generator & discriminator loss 균형맞춤. discriminator가 얼마나  fake images에 집중할 것인지 컨트롤. 매 batch마다 업데이트.
    k_var += Lambda*(GAMMA*real_loss - fake_loss)
    K.set_value(k, k_var)
    

    # g_measure 을 계산하기 위한 loss 저장
    history.append({'real_loss': real_loss,'fake_loss': fake_loss })

    # 100번에 1번씩 로그 표시
    if epoch%100 == 0:
        # 과거 100 번의 measure 의 평균
        measurement = np.mean([measure(loss['real_loss'],loss['fake_loss'],GAMMA) for loss in history[-100:]])
        logs.append({'Epoch:':epoch, 'k': K.get_value(k),'measure': measurement,'real_loss': real_loss,'fake_loss': fake_loss })
        print(logs[-1])

        # 생성된 이미지 저장 및 보이기
        img_path = '{}/generated_{}.png'.format(IMG_SAVE_DIR, epoch)
        fig  = plot_imgs(img_path, generator.predict(sample_seeds), rows=IMG_SAMPLE_SHAPE[0],cols=IMG_SAMPLE_SHAPE[1])
        plt.savefig(img_path, bbox_inches='tight')     
        plt.axis('off')
        plt.show()    
{'Epoch:': 0, 'k': 0.00011048602, 'measure': 0.24255565404891968, 'real_loss': 0.13206963, 'fake_loss': 0.0612045}
{'Epoch:': 100, 'k': 0.01051935, 'measure': 0.21105107249319555, 'real_loss': 0.090453535, 'fake_loss': 0.024226436}
{'Epoch:': 200, 'k': 0.019882021, 'measure': 0.17836286436021326, 'real_loss': 0.08530679, 'fake_loss': 0.01522289}
{'Epoch:': 300, 'k': 0.028015548, 'measure': 0.15981527408026155, 'real_loss': 0.07372745, 'fake_loss': 0.027405646}
{'Epoch:': 400, 'k': 0.034613445, 'measure': 0.139005820132792, 'real_loss': 0.067190014, 'fake_loss': 0.028980296}
{'Epoch:': 500, 'k': 0.03945609, 'measure': 0.11123871646821498, 'real_loss': 0.055348106, 'fake_loss': 0.032984372}
{'Epoch:': 600, 'k': 0.04329482, 'measure': 0.09449506510794164, 'real_loss': 0.05426988, 'fake_loss': 0.04434004}
{'Epoch:': 700, 'k': 0.046639696, 'measure': 0.08501894175633785, 'real_loss': 0.047984272, 'fake_loss': 0.035520133}