VAE (Variational autoencoder): 다변수 오토인코더

import os
import numpy as np
import matplotlib.pyplot as plt
import random
import tensorflow as tf


# GPU 디바이스 목록 가져오기

gpus = tf.config.experimental.list_physical_devices('GPU')

if gpus:
    print("사용 가능한 GPU 디바이스:")
    for gpu in gpus:
        print("디바이스 이름:", gpu.name)
else:
    print("사용 가능한 GPU 디바이스가 없습니다.")

# 이미지 전처리 코드
image_dir = "C:\\Users\\crazy\\Downloads\\archive\\img_align_celeba" #이미지 데이터셋 경로입니다.
resize_width, resize_height = 128, 128  # 변경할 크기

# 이미지 목록 가져오기
image_files = os.listdir(image_dir)

# 이미지를 NumPy 배열로 로드하고 크기 변경하기
images = []
for filename in image_files:
    img = tf.keras.preprocessing.image.load_img(os.path.join(image_dir, filename), target_size=(resize_width, resize_height))
    img = tf.keras.preprocessing.image.img_to_array(img)
    img = img / 255.0  # 이미지 정규화
    images.append(img)
images = np.array(images)


# 이미지 중에서 무작위로 하나 선택
random_index = random.randint(0, len(images) - 1)
selected_image = images[random_index]

keras = tf.keras
K = tf.keras.backend

#샘플링 클래스 정의
#keras의 layer를 상속받아 만들어진 클래스 입니다.
#call 메서드는 z_mean(평균) 과 z_log_var(로그 분산)을 받습니다.
#z_mean 과 z_log_var를 사용해 잠재변수(z)를 return 합니다.
class Sampling(keras.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 = keras.layers.Input(shape=(resize_width, resize_height, 3), name="encoder_input")
x = keras.layers.Conv2D(32, (3, 3), strides=2, activation="relu", padding="same")(encoder_input)
x = keras.layers.Conv2D(64, (3, 3), strides=2, activation="relu", padding="same")(x)
x = keras.layers.Conv2D(128, (3, 3), strides=2, activation="relu", padding="same")(x)
shape_before_flattening = K.int_shape(x)[1:]
x = keras.layers.Flatten()(x)
z_mean = keras.layers.Dense(200, name="z_mean")(x)
z_log_var = keras.layers.Dense(200, name="z_log_var")(x)
z = Sampling()([z_mean, z_log_var])
encoder = keras.models.Model(encoder_input, [z_mean, z_log_var, z], name="encoder")

#디코더 부분
decoder_input = keras.layers.Input(shape=(200,), name="decoder_input")
x = keras.layers.Dense(np.prod(shape_before_flattening))(decoder_input)
x = keras.layers.Reshape(shape_before_flattening)(x)
x = keras.layers.Conv2DTranspose(128, (3, 3), strides=2, activation = 'relu', padding="same")(x)
x = keras.layers.Conv2DTranspose(64, (3, 3), strides=2, activation = 'relu', padding="same")(x)
x = keras.layers.Conv2DTranspose(32, (3, 3), strides=2, activation = 'relu', padding="same")(x)
decoder_output = keras.layers.Conv2D(3,(3, 3),strides = 1, activation="relu", padding="same", name="decoder_output")(x)
decoder = keras.models.Model(decoder_input, decoder_output)

#VAE 클래스는 Keras의 Model 클래스를 상속받아서 만들어진 클래스입니다. 이 클래스는 VAE 모델을 정의하고 학습하기 위한 기능을 포함합니다.
#__init__ 메서드: 클래스의 생성자로, encoder와 decoder 모델을 인수로 받습니다. 이 두 모델은 VAE의 핵심 구성 요소인 인코더와 디코더입니다.
#또한 **kwargs를 사용하여 추가 인수를 받을 수 있습니다.
class VAE(keras.models.Model):
    def __init__(self, encoder, decoder, **kwargs):
        super(VAE, self).__init__(**kwargs)
        self.encoder = encoder
        self.decoder = decoder

        #loss 값을 모니터링 하는 요소
        self.total_loss_tracker = keras.metrics.Mean(name='total_loss')
        self.reconstruction_loss_tracker = keras.metrics.Mean(name="reconstruction_loss")
        self.kl_loss_tracker = keras.metrics.Mean(name="kl_loss")

    #metrics 프로퍼티: VAE 모델의 성능을 모니터링하는 메트릭을 정의합니다. 이 모델은 총 손실 (total_loss), 재구성 손실 (reconstruction_loss), 및 KL-다이버전스 손실 (kl_loss)을 모니터링합니다.
    @property
    def metrics(self):
        return [
            self.total_loss_tracker,
            self.reconstruction_loss_tracker,
            self.kl_loss_tracker,
        ]

    #call 메서드는 모델을 호출할 때 실행되는 메서드입니다. 입력 이미지를 받아서 잠재 변수를 인코더로부터 추출하고, 해당 잠재 변수를 사용하여 이미지를 디코더로부터 재구성합니다.
    #input은 입력 이미지 데이터를 나타냅니다.
    #z_mean, z_log_var, z는 인코더로부터 얻은 잠재 변수의 평균, 로그 분산 및 샘플링된 잠재 변수를 나타냅니다.
    #reconstruction은 디코더로부터 재구성된 이미지를 나타냅니다.
    def call(self, input):
        z_mean, z_log_var, z = encoder(input)
        reconstruction = decoder(z)
        return z_mean, z_log_var, reconstruction
    
    #train_step 메서드는 VAE 모델의 학습 단계를 정의합니다.
    #입력으로 데이터 (data)를 받아서 학습을 진행합니다.
    #tf.GradientTape() 내에서 재구성 손실과 KL-다이버전스 손실을 계산합니다.
    #reconstruction_loss는 재구성 손실을 계산하는 부분입니다. tf.losses.binary_crossentropy를 사용하여 재구성 손실을 계산하며, 이 값을 500으로 스케일링합니다.
    #kl_loss는 KL-다이버전스 손실을 계산하는 부분입니다. 이것은 VAE의 정규화 항으로, 잠재 변수의 분포와 정규 분포 간의 차이를 측정합니다.
    #total_loss는 재구성 손실과 KL-다이버전스 손실을 더한 총 손실입니다.
    #경사 하강법을 사용하여 총 손실을 최소화하도록 모델을 업데이트하고, 메트릭 값을 업데이트합니다.
    #최종적으로 업데이트된 메트릭 값을 딕셔너리로 반환합니다.
    def train_step(self, data):
        with tf.GradientTape() as tape:
            z_mean, z_log_var, reconstruction = self(data)
            reconstruction_loss = tf.reduce_mean(500 * tf.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}

#모델 새로 정의할때
vae = VAE(encoder, decoder)

#훈련시키기 (훈련후에 결과가 맘에들지않으면 다시 훈련을 시켜줍니다.)
vae.compile(optimizer="adam")
vae.fit(images, epochs=100, batch_size=320)

# 이미지 중에서 무작위로 하나 선택
random_index1 = random.randint(0, len(images) - 1)
selected_image1 = images[random_index1]

# 또 다른 이미지를 무작위로 선택
random_index2 = random.randint(0, len(images) - 1)
selected_image2 = images[random_index2]

# 선택한 이미지 표시
plt.figure(figsize=(10, 5))  # 그림 크기 조정
plt.subplot(1, 2, 1)  # 1x2 그리드의 첫 번째 위치
plt.imshow(selected_image1)
plt.title("Original Image 1")

# 인코더에 넣어서 임베딩 벡터 확인
embedded1 = np.array(encoder.predict(np.array([selected_image1])))[2]

# 임베딩 벡터를 디코더에 넣어서 결과값 예측
predicted_img1 = decoder.predict(embedded1)
plt.subplot(1, 2, 2)  # 1x2 그리드의 두 번째 위치
plt.imshow(predicted_img1.reshape(resize_width, resize_height, 3))
plt.title("Reconstructed Image 1")

plt.show()

# 선택한 다른 이미지 표시
plt.figure(figsize=(10, 5))  # 그림 크기 조정
plt.subplot(1, 2, 1)  # 1x2 그리드의 첫 번째 위치
plt.imshow(selected_image2)
plt.title("Original Image 2")

# 인코더에 넣어서 임베딩 벡터 확인
embedded2 = np.array(encoder.predict(np.array([selected_image2])))[2]

# 임베딩 벡터를 디코더에 넣어서 결과값 예측
predicted_img2 = decoder.predict(embedded2)
plt.subplot(1, 2, 2)  # 1x2 그리드의 두 번째 위치
plt.imshow(predicted_img2.reshape(resize_width, resize_height, 3))
plt.title("Reconstructed Image 2")

plt.show()

# 두 이미지의 임베딩 벡터의 중간값을 이용해 결과값 예측
predicted_img_combined = decoder.predict((embedded1 + embedded2) / 2)
plt.figure(figsize=(5, 5))
plt.imshow(predicted_img_combined.reshape(resize_width, resize_height, 3))
plt.title("Reconstructed Combined Image")

plt.show()