Machine Learning Diary :: 03 - 오토인코더로 이미지 압축, 복원2 (Building Simple AutoEncoder with MNIST)

개인 프로젝트/(DL) 생성 모델링2020. 2. 26. 10:59

03 - 오토인코더로 이미지 압축, 복원2 (Building Simple AutoEncoder with MNIST)

세상에서 가장 간단한 오토인코더를 만들었었다. 이 오토인코더에 몇 개의 레이어를 추가해 성능을 조금 더 향상시켜보자.

원래 코드를 맥북에서 작성하는게 더 편해서, 간단한 딥러닝 모델 정도는 맥북으로 그냥 돌리고 마는데, 여기부터는 컴퓨터가 힘들어하기 시작해서 게이밍 노트북으로 옮겨서 실행했다. 역시 GPU는 짱짱이다.

그냥 뭐 설명을 적어놓을건 딱히 더 없으니 코드만 대충 약간 올려놔야겠다.

import pandas as pd
import numpy as np
import keras
from keras.layers import Input, Conv2D, Flatten, Dense, Conv2DTranspose, Reshape, Activation, LeakyReLU, Dropout, BatchNormalization, MaxPooling2D, Lambda
from keras.models import Model
from keras import backend as K
from keras.optimizers import Adam
from keras.utils import to_categorical
from sklearn.metrics import mean_squared_error as mse

import matplotlib.pyplot as plt
%matplotlib inline
import seaborn as sns

(Xtrain, Ytrain), (Xtest, Ytest) = keras.datasets.mnist.load_data()
Xtrain=Xtrain/255
Xtest=Xtest/255

Xtrain=Xtrain.reshape(len(Xtrain), 28, 28, 1)
Xtest=Xtest.reshape(len(Xtest), 28, 28, 1)

Ytrain=to_categorical(Ytrain)
Ytest=to_categorical(Ytest)

#Encoder
encoder_input=Input(shape=(28,28,1))

x=Conv2D(filters=32, kernel_size=3, strides=1, padding='same')(encoder_input)
x=LeakyReLU()(x)

x=Conv2D(filters=64, kernel_size=3, strides=2, padding='same')(x)
x=LeakyReLU()(x)

x=Conv2D(filters=64, kernel_size=3, strides=2, padding='same')(x)
x=LeakyReLU()(x)

x=Conv2D(filters=64, kernel_size=3, strides=1, padding='same')(x)
x=LeakyReLU()(x)

shape_before_flatten=K.int_shape(x)[1:]
x=Flatten()(x)
encoder_output=Dense(units=2)(x)
model_encoder=Model(encoder_input, encoder_output)
print(model_encoder.summary())

#Decoder
decoder_input=Input(shape=(2,))
x=Dense(units=np.prod(shape_before_flatten))(decoder_input)
x=Reshape(shape_before_flatten)(x)

x=Conv2DTranspose(filters=64, kernel_size=3, strides=1, padding='same')(x)
x=LeakyReLU()(x)

x=Conv2DTranspose(filters=64, kernel_size=3, strides=2, padding='same')(x)
x=LeakyReLU()(x)

x=Conv2DTranspose(filters=32, kernel_size=3, strides=2, padding='same')(x)
x=LeakyReLU()(x)

x=Conv2DTranspose(filters=1, kernel_size=3, strides=1, padding='same')(x)
x=Activation('sigmoid')(x)

decoder_output=x
model_decoder=Model(decoder_input, decoder_output)
print(model_decoder.summary())


#Connect Two Models
model_input = encoder_input
model_output = model_decoder(encoder_output)

AutoEncoder=Model(model_input, model_output)

인코더의 Input shape와, 디코더의 Output shape 가 같으므로 돌아갈 수 있는 모델이다.

optimizer=Adam(lr=0.0005)
AutoEncoder.compile(optimizer=optimizer, loss='mean_squared_error')


for i in range(10):
    print(f"Trial {i}...")
    AutoEncoder.fit(Xtrain, Xtrain, batch_size=32, shuffle=True, epochs=20)
    result=AutoEncoder.predict(Xtrain)
    
    random=np.random.randint(0, len(Xtrain))
    
    fig = plt.figure()
    rows = 1
    cols = 2
    img1 = Xtrain[random].reshape(28,28)
    img2 = result[random].reshape(28,28)
    ax1 = fig.add_subplot(rows, cols, 1)
    ax1.imshow(img1)
    ax1.set_title('Correct')
    ax1.axis("off")
    ax2 = fig.add_subplot(rows, cols, 2)
    ax2.imshow(img2)
    ax2.set_title('Generated')
    ax2.axis("off")
    plt.show()

오토인코더 학습을 반복문으로 10회 반복 학습시키고, 학습당 20 에포크를 준다. 시간 여유가 좀 있으니까 (+GPU를 사용하면 빠르니까) 총 200 에포크만큼 반복 학습을 시킨다. 이때 20 에포크당 Xtrain에서 랜덤하게 데이터 1개를 가져와서 결과물 이미지랑 비교한다.