[Generative-AI] 오토인코더(AE)에서 VAE까지

0. 신경망이 학습을 한다는 것

- 구두, 후드, 청바지, 면바지, 원피스, 운동화 등 총 10가지로 분류할수 있는 이미지가 10,000장 있다고 가정

- classification학습: CNN신경망 + 최종 노드가 10인 DNN + softmax로 학습 -> 이미지 입력되면 10가지중 1개로 알려줌 

- (2차원으로)임베딩 학습: CNN신경망 + 최종노드가 2인 DNN 학습 -> 2차원 공간에 비슷한것끼리 뭉치는 개념

무작위 데이터가 임베딩(학습)을 통해 2개/3개/4개로 분류되는 모습(필자 아이패드 그림ㅋ)

 

1. AE(오토인코더) 아이디어

- 그림을 생성하는 모델을 어떻게 만들수 있을까?

- 이미지를 신경망에 입력하고, output역시 입력된 이미지가 나오는 신경망을 학습 가능하지 않은가?

입력된 이미지가 그대로 나오는 신경망

- 이미지를 입력받아 CNN을 통과하고, 최종적으로 2차원으로 줄이는 DNN을 붙여서 왼쪽 신경망(인코더)

- 2차원입력을 받아서 DNN을 통과해서, 역CNN을 통과하여 이미지로 복원(디코더)

- input: 이미지, label: 이미지, 이렇게 학습을 하면,

- 입력된 이미지가 2개의 좌표로 표현되고, 2개의 좌표를 입력하면 이미지가 생성되는 신경망이 학습된다.

- 여기서 디코더만 때서, 임의의 숫자 2개를 입력하면 이미지가 생성되지 않을까?

 

2. AE 코드로 확인하기

2-1. 데이터 셋팅

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

# 데이터를 로드합니다.
(x_train, y_train), (x_test, y_test) = datasets.fashion_mnist.load_data()

# 데이터 전처리
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)

# 훈련 세트에 있는 의류 아이템 일부를 출력합니다.
display(x_train)

- MNIST fashion데이터 셋을 불러옵니다

- 이미지 전처리를 진행합니다.

- 아래와 같이 데이터 확인 가능합니다

MNIST fashion데이터

 

2-2. AE 만들기(인코더 신경망)

# 인코더
encoder_input = layers.Input(
    shape=(IMAGE_SIZE, IMAGE_SIZE, CHANNELS), 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)
# K.int_shape(x)함수는 해당 텐서를 받아서 형태를 반환 (배치, 높이, 너비, 채널)
# [1:] > (배치, 높이, 너비, 채널)에서 (높이,너비,채널)가져오는 거
shape_before_flattening = K.int_shape(x)[1:]  # 디코더에 필요합니다!

x = layers.Flatten()(x)
encoder_output = layers.Dense(EMBEDDING_DIM, name="encoder_output")(x)

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

- 입력으로 이미지를 받아서, CNN 레이어 3개 통과

- flatten한후 - 2차원으로 임베딩 합니다(EMBEDDING_DIM=2)

- 이미지를 2차원으로 임베딩하는 신경망이 학습됩니다.

 

2-3. AE 만들기(디코더 신경망)

# 디코더
decoder_input = layers.Input(shape=(EMBEDDING_DIM,), name="decoder_input")
# np.prod(shape_before_flattening) = 높이*너비*채널
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(
    CHANNELS,
    (3, 3),
    strides=1,
    activation="sigmoid",
    padding="same",
    name="decoder_output",
)(x)

decoder = models.Model(decoder_input, decoder_output)
decoder.summary()

- 2개 노드(2차원입력)에서 시작해서, CNN을 3개 통과

- 2차원 데이터를 입력받아, 이미지를 생성하는 부분입니다.

 

2-4. 오토인코더 만들기(인코더 + 디코더)

# 오토인코더
autoencoder = models.Model(
    encoder_input, decoder(encoder_output)
)
autoencoder.summary()

- 두개의 신경망을 하나로 이어줍니다.

 

2-5. AE훈련하기

# 오토인코더 컴파일
autoencoder.compile(optimizer="adam", loss="binary_crossentropy")

# 모델 저장 체크포인트 생성
model_checkpoint_callback = callbacks.ModelCheckpoint(
    filepath="./checkpoint",
    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=EPOCHS,
    batch_size=BATCH_SIZE,
    shuffle=True,
    validation_data=(x_test, x_test),
    callbacks=[model_checkpoint_callback, tensorboard_callback],
)

# 최종 모델을 저장합니다.
autoencoder.save("./models/autoencoder")
encoder.save("./models/encoder")
decoder.save("./models/decoder")

- input도 x_train, label도 x_train = 이미지 입력하고, 동일한 이미지 출력하는 신경망(오토인코더) 학습하기

 

2-6. 디코더만 때서 이미지 생성하기

# 잠재 공간에서 포인트를 샘플링합니다.
grid_width, grid_height = (6, 3)
sample = np.random.uniform(
    mins, maxs, size=(grid_width * grid_height, EMBEDDING_DIM)
)

# 샘플링된 포인트를 디코딩합니다.
reconstructions = decoder.predict(sample)

imshow(reconstructions[i, :, :], cmap="Greys")

- 임의의 2차원 숫자 가져와서, 디코더에 넣습니다.

- 이미지를 그려봅니다.

2차원 숫자입력시 출력되는 이미지

- 그럴듯한 이미지도 존재하고, 전혀 말도 안되는 느낌의 이미지도 존재합니다.

 

3. AE의 한계

- 인코더가 2차원으로 임베딩하는 학습을 하였지만, 학습한 이미지의 임베딩값이 입력되지 않으면, 뜨금없는 이미지가 등장

- 이는 학습이 각각의 임베딩 포인트로 되었기 때문

- 결국 배우지 않은 입력값에 대해서는 (학습하지 않은 값이 디코더에 입력되기 때문에) 그럴듯한 이미지를 못 생성해냄.

- VAE 아이디어가 이 문제를 해결

 

4. 변이형 오토인코더(VAE) 아이디어

- 기존에 AE는 2차원 잠재공간의 1점에 맵핑

- 하지만 VAE는 분포로 맵핑(기존의 2차원 값 > (평균, 분산))

- 디코더에는 기존과 동일하게 2차원 값이 입력되지만, 이는 분포에서 나온 2차원 값(동일한 이미지를 입력해도, 디코더에 입력되는 값은 다르다)

- AE는 동일한 이미지가 입력되면, 디코더에 입력되는 입력값 동일

- VAE에서는 동일한 이미지 입력되어도, 분포를 기반으로 디코더에 서로다른 값이 입력됨

(예: 이미지1 입력 > Z(z1,z2) = 평균 + 표준편차*랜덤, 분포에 따라 매 입력때마다 다른 Z(z1,z2) 출력)

 

5. VAE 코드로 확인하기

5-1. 데이터 셋팅

IMAGE_SIZE = 32
BATCH_SIZE = 100
VALIDATION_SPLIT = 0.2
EMBEDDING_DIM = 2
EPOCHS = 5
BETA = 500

# 데이터 로드
(x_train, y_train), (x_test, y_test) = datasets.fashion_mnist.load_data()

# 데이터 전처리
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)

- 동일함

 

5-2. VAE만들기(인코더)

class Sampling(layers.Layer):
    def call(self, inputs):
        z_mean, z_log_var = inputs
        batch = tf.shape(z_mean)[0]
        dim = tf.shape(z_mean)[1]
        epsilon = K.random_normal(shape=(batch, dim))
        return z_mean + tf.exp(0.5 * z_log_var) * epsilon
        
 
# 인코더
encoder_input = layers.Input(
    shape=(IMAGE_SIZE, IMAGE_SIZE, 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)
z_mean = layers.Dense(EMBEDDING_DIM, name="z_mean")(x)
z_log_var = layers.Dense(EMBEDDING_DIM, name="z_log_var")(x)
z = Sampling()([z_mean, z_log_var])

encoder = models.Model(encoder_input, [z_mean, z_log_var, z], name="encoder")
encoder.summary()

- AE의 인코더와 다르게 2차원의 데이터가 아닌, z_mean, z_log_var (평균, 분산)을 학습

- 학습해서 나온 평균과 분산으로 디코더에 입력할 임베딩 데이터 완성(Sampling함수)

 

5-3. VAE만들기(디코더 신경망)

# 디코더
decoder_input = layers.Input(shape=(EMBEDDING_DIM,), 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)
decoder.summary()

- AE 디코더 동일

 

5-4. VAE만들기(인코더+디코더 합치기)

class VAE(models.Model):
    def __init__(self, encoder, decoder, **kwargs):
        super(VAE, self).__init__(**kwargs)
        self.encoder = encoder
        self.decoder = decoder
        self.total_loss_tracker = metrics.Mean(name="total_loss")
        self.reconstruction_loss_tracker = metrics.Mean(
            name="reconstruction_loss"
        )
        self.kl_loss_tracker = metrics.Mean(name="kl_loss")

    @property
    def metrics(self):
        return [
            self.total_loss_tracker,
            self.reconstruction_loss_tracker,
            self.kl_loss_tracker,
        ]

    def call(self, inputs):
        """특정 입력에서 모델을 호출합니다."""
        z_mean, z_log_var, z = encoder(inputs)
        reconstruction = decoder(z)
        return z_mean, z_log_var, reconstruction

    def train_step(self, data):
        """훈련 스텝을 실행합니다."""
        with tf.GradientTape() as tape:
            z_mean, z_log_var, reconstruction = self(data)
            reconstruction_loss = tf.reduce_mean(
                BETA
                * losses.binary_crossentropy(
                    data, reconstruction, axis=(1, 2, 3)
                )
            )
            kl_loss = tf.reduce_mean(
                tf.reduce_sum(
                    -0.5
                    * (1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var)),
                    axis=1,
                )
            )
            total_loss = reconstruction_loss + kl_loss

        grads = tape.gradient(total_loss, self.trainable_weights)
        self.optimizer.apply_gradients(zip(grads, self.trainable_weights))

        self.total_loss_tracker.update_state(total_loss)
        self.reconstruction_loss_tracker.update_state(reconstruction_loss)
        self.kl_loss_tracker.update_state(kl_loss)

        return {m.name: m.result() for m in self.metrics}

    def test_step(self, data):
        """Step run during validation."""
        if isinstance(data, tuple):
            data = data[0]

        z_mean, z_log_var, reconstruction = self(data)
        reconstruction_loss = tf.reduce_mean(
            BETA
            * losses.binary_crossentropy(data, reconstruction, axis=(1, 2, 3))
        )
        kl_loss = tf.reduce_mean(
            tf.reduce_sum(
                -0.5 * (1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var)),
                axis=1,
            )
        )
        total_loss = reconstruction_loss + kl_loss

        return {
            "loss": total_loss,
            "reconstruction_loss": reconstruction_loss,
            "kl_loss": kl_loss,
        }

    def get_config(self):
        return {}
        

# 변이형 오토인코더 생성
vae = VAE(encoder, decoder)

- VAE는 KL발산은 확률분포간 차이를 측정하는 도구

- 위의 kl_loss는 샘플의 분포가 표준정규분포에 크게 벗어난 인코딩 값에 대해서는 벌칙을 가합니다

- 이는 인코더가 학습시 결과값으로 내보내는 분포가 될수 있으면 정규분포에 가까운 분포를 내보내도록 합니다. 이는 동일한 이미지에 대해서 안정적인 디코더 입력값이 나올수 있음(만약 정규분포가 아닐경우, 동일한 이미지인데 굉장히 격차가 큰 2차원 값이 인코더에서 나오고 이 값이 디코더에 입력...이는 안정적인 학습에 방해가 됨)

- KL발산항에 너무 큰 가중치를 주면, 역으로 KL발산 손실에 치중되어 재구성 이미지 품질이 나빠질수 있음... 잘 조절해야 함

 

5-5. VAE 훈련

# 변이형 오토인코더 컴파일
optimizer = optimizers.Adam(learning_rate=0.0005)
vae.compile(optimizer=optimizer)

# 모델 저장 체크포인트 생성
model_checkpoint_callback = callbacks.ModelCheckpoint(
    filepath="./checkpoint",
    save_weights_only=False,
    save_freq="epoch",
    monitor="loss",
    mode="min",
    save_best_only=True,
    verbose=0,
)
tensorboard_callback = callbacks.TensorBoard(log_dir="./logs")

vae.fit(
    x_train,
    epochs=EPOCHS,
    batch_size=BATCH_SIZE,
    shuffle=True,
    validation_data=(x_test, x_test),
    callbacks=[model_checkpoint_callback, tensorboard_callback],
)

# 최종 모델 저장
vae.save("./models/vae")
encoder.save("./models/encoder")
decoder.save("./models/decoder")

 

5-6. 디코더로 이미지 생성하기

# 예제 이미지 인코딩
z_mean, z_var, z = encoder.predict(example_images)

# 2D 공간에서 인코딩된 포인트 표시
figsize = 8

plt.figure(figsize=(figsize, figsize))
plt.scatter(z[:, 0], z[:, 1], c="black", alpha=0.5, s=3)
plt.show()

# 표준 정규 분포에서 잠재 공간의 일부 포인트를 샘플링합니다.
grid_width, grid_height = (6, 3)
z_sample = np.random.normal(size=(grid_width * grid_height, 2))

# 샘플링된 포인트 디코딩
reconstructions = decoder.predict(z_sample)

# 원본 임베딩과 샘플링된 임베딩을 p값으로 변환하기
p = norm.cdf(z)
p_sample = norm.cdf(z_sample)

# 그래프를 그립니다....
figsize = 8
plt.figure(figsize=(figsize, figsize))

# ... 원본 임베딩 ...
plt.scatter(z[:, 0], z[:, 1], c="black", alpha=0.5, s=2)

# ... 잠재 공간에 새로 생성된 포인트
plt.scatter(z_sample[:, 0], z_sample[:, 1], c="#00B0F0", alpha=1, s=40)
plt.show()

# 디코딩된 이미지 그리드를 추가합니다.
fig = plt.figure(figsize=(figsize, grid_height * 2))
fig.subplots_adjust(hspace=0.4, wspace=0.4)

for i in range(grid_width * grid_height):
    ax = fig.add_subplot(grid_height, grid_width, i + 1)
    ax.axis("off")
    ax.text(
        0.5,
        -0.35,
        str(np.round(z_sample[i, :], 1)),
        fontsize=10,
        ha="center",
        transform=ax.transAxes,
    )
    ax.imshow(reconstructions[i, :, :], cmap="Greys")

- AE와 다르게 분포를 입력해서 이미지를 출력합니다.

AE대비 상대적으로 괜찮은 이미지를 출력하는 VAE

 

6. AE, VAE정리

- AE는 인코더, 디코더 아이디어를 통해서, 다양한 이미지를 저차원의 잠재공간으로 매핑할수 있었고, 이 매핑값만 가지고, 이미지를 생성할수 있었다. 하지만, 포인트 단위 학습의 한계가 존재.

- VAE의 인코더는 잠재공간의 한개의 포인트가 아닌, 분포로 학습. 이는 잠재공간상의 한개 점이 아닌, 영역을 학습하는 개념. 당연히, AE보다 다양한 입력값에 대해서, 디코더가 안정적으로 이미지를 생성해냄

 

AE, VAE는 생성형 모델의 기초로... 이후, GAN, 디퓨전, 에너지모델 등이 나온다.