AI - Tensorflow로 AutoEncoder 구현하기

 

이전 글에서 AutoEncoder의 개념과 구조를 배웠다면,

이번에는 직접 Tensorflow와 Keras를 사용해 Fashion-MNIST 이미지에 AutoEncoder를 구현해보겠습니다.

2025.07.28 - [AI/AI] - AI - AutoEncoder 직관적으로 이해하기

AI - AutoEncoder 직관적으로 이해하기

 

라이브러리 임포트

import numpy as np
import matplotlib.pyplot as plt
from tensorflow.keras import layers, models, datasets, callbacks
import tensorflow.keras.backend as K

 

데이터 로드

(X_train, y_train), (X_test, y_test) = datasets.fashion_mnist.load_data()

• datasets.fashion_mnist는 흑백 옷 이미지(28x28)를 제공하는 내장 데이터셋입니다.

 

데이터 전처리 함수

def preprocess(imgs):
    imgs = imgs.astype('float32') / 255.0
    imgs = np.pad(imgs, ((0,0), (2,2), (2,2)), constant_values=0.0)
    imgs = np.expand_dims(imgs, -1)
    return imgs

X_train = preprocess(X_train)
X_test = preprocess(X_test)

 

• imgs.astype / 255.0 : 정규화, 데이터 값을 0과 1 사이로 변환
• pad : 28 x 28 이미지를 Conv Layer와의 호환성이 맞는 32 x 32 크기로 맞추기 위해 패딩 연산
• expand_dims : (batch, 32, 32, 1) 형태로 reshape

 

하이퍼파라미터 설정

IMAGE_SIZE = 32
CHANNELS = 1
BATCH_SIZE = 100
BUFFER_SIZE = 1000
VALIDATION_SPLIT = 0.2
EMBEDDING_DIM = 2
EPOCHS = 3

 

• Embedding_dim = 2 : Latent Vector의 차원을 2로 설정하여 시각화 가능케 하기 위함

Encoder 정의

Encoder는 고차원 데이터를 점점 줄여서 가장 핵심적인 특징만 남긴 Latent Vector를 생성합니다.

encoder_input = layers.Input(shape=(32, 32, 1), name='encoder_input')

x = layers.Conv2D(32, (3,3), strides=2, activation='relu', padding='same')(encoder_input)
x = layers.Conv2D(64, (3,3), strides=2, activation='relu', padding='same')(x)
x = layers.Conv2D(128, (3,3), strides=2, activation='relu', padding='same')(x)

shape_before_flattening = K.int_shape(x)[1:]
x = layers.Flatten()(x)
encoder_output = layers.Dense(2, name='encoder_output')(x)

encoder = models.Model(encoder_input, encoder_output)
encoder.summary()

• Conv2D 3 Layer : 점점 이미지를 축소하면서 정보만 추출
• ReLU : 비선형성을 위함
• Flatten : (4, 4, 128) => (2048, )
• Dense(2) : 최종적으로 2개의 숫자로 압축(Latent Vector), 2차원 이미지 시각화 위함

 

Decoder 정의

Decoder는 Encoder가 만든 Latent Vector를 입력으로 받아 다시 이미지 형태로 복원하는 과정입니다.

decoder_input = layers.Input(shape=(2,), name='decoder_input')

x = layers.Dense(np.prod(shape_before_flattening))(decoder_input)
x = layers.Reshape(shape_before_flattening)(x)
x = layers.Conv2DTranspose(128, (3,3), strides=2, activation='relu', padding='same')(x)
x = layers.Conv2DTranspose(64, (3,3), strides=2, activation='relu', padding='same')(x)
x = layers.Conv2DTranspose(32, (3,3), strides=2, activation='relu', padding='same')(x)

decoder_output = layers.Conv2D(1, (3,3), strides=1, activation='sigmoid', padding='same', name='decoder_output')(x)

decoder = models.Model(decoder_input, decoder_output)

• Dense => Reshape : 2D Latent Vector => (4, 4, 128) 형태로 다시 변환
• Conv2D Transpose : Upsampling
• Sigmoid : 픽셀 값을 0 ~ 1 범위로 제한

Model 연결

autoencoder = models.Model(encoder_input, decoder(encoder_output))

 

Encoder와 Decoder를 결합합니다.

 

AutoEncoder 학습 

autoencoder.compile(optimizer="adam", loss="binary_crossentropy")

model_checkpoint_callback = callbacks.ModelCheckpoint(
    filepath="./checkpoint.keras",
    save_weights_only=False,
    save_freq="epoch",
    monitor="loss",
    mode="min",
    save_best_only=True,
    verbose=0,
)
tensorboard_callback = callbacks.TensorBoard(log_dir="./logs")

autoencoder.fit(
    X_train,
    X_train,
    epochs=10,
    batch_size=BATCH_SIZE,
    shuffle=True,
    validation_data=(X_test, X_test),
    callbacks=[model_checkpoint_callback, tensorboard_callback],
)

 

• optimizer : adam 사용
• loss : binary crossentropy
• epoch : 10

모델 저장

autoencoder.save("./models/autoencoder.keras")
encoder.save("./models/encoder.keras")
decoder.save("./models/decoder.keras")

 

AutoEncoder 복원 시각화

n_to_predict = 5000
example_images = X_test[:n_to_predict]
example_labels = y_test[:n_to_predict]

predictions = autoencoder.predict(example_images)

n = min(10, len(example_images))
plt.figure(figsize=(20, 4))

for i in range(n):
    ax = plt.subplot(2, n, i + 1)
    plt.imshow(example_images[i].squeeze(), cmap='gray')
    plt.title("Input")
    plt.axis("off")

 

자, 이제 AutoEncoder가 복원한 이미지를 시각화 해보겠습니다.

 

Latent Vector 시각화

embeddings = encoder.predict(example_images)
example_labels = y_test[:n_to_predict]

figsize = 8
plt.figure(figsize=(figsize, figsize))
plt.scatter(
    embeddings[:, 0],
    embeddings[:, 1],
    cmap="rainbow",
    c=example_labels,
    alpha=0.8,
    s=3,
)
plt.colorbar()
plt.show()

 

 

---

이번 실습에서는 TensorFlow/Keras를 활용해 Fashion-MNIST 이미지 데이터를 기반으로
Convolutional AutoEncoder를 처음부터 구현하고, 시각적으로 분석해보았습니다.

 

다음 글에서는 Pytorch를 이용하여 똑같이 구현해보겠습니다.