IT/AI\ML

[python/Tensorflow2.0] AutoEncoder(오토인코더) ; Variational autoencoder(VAE)

개발자 두더지 2020. 4. 22. 11:46
728x90

1. VAE(Variational AutoEncoder)란?

VAE(Variational AutoEncoder)는 2014년 D.Kingma와 M.Welling이 Auto-Encoding Variational Bayes논문에서 제안한 오토인코더의 한 종류이다.

VAE는 이전에 살펴 본 오터인코더와 다음과 같은 차이점이 있다.

- VAE는 확률적 오터인코더(probabilistic autoencoder)다. 즉, 학습이 끝난 후에도 출력이 부분적으로 우연에 의해 결정된다.

- VAE는 생성 오토인코더(generatic autoencoder)이며, 학습 데이터셋에서 샘플링된 것과 같은 새로운 샘플을 생성할 수 있다.

VAE의 구조는 다음의 그림과 같다.

 VAE의 코딩층은 다른 오토인코더와는 다른 부분이 있는데 주어진 입력에 대해 바로 코딩을 만드는 것이 아니라, 인코더(encoder)는 평균 코딩μ와 표준편차 코딩δ을 만든다. 실제 코딩은 평균이μ이고 표준편차가δ인 가우시안 분포(gaussian distribution)에서 랜덤하게 샘플링되며, 이렇게 샘플링된 코딩을 디코더(decoder)가 원본 입력으로 재구성하게 된다.

 VAE는 마치 가우시안 분포에서 샘플링된 것처럼 보이는 코딩을 만드는 경향이 있는데, 학습하는 동안 손실함수가 코딩(coding)을 가우시안 샘플들의 집합처럼 보이는 형태를 가진 코딩 공간(coding space) 또는 잠재 변수 공간(latent space)로 이동시키기 때문이다. 

 이러한 이유로 VAE는 학습이 끝난 후에 새로운 샘플을 가우시안 분포로 부터 랜덤한 코딩을 샘플링해 디코딩해서 생성할 수 있다.


2. 생성 모델로서의 Variational AutoEncoder

판별모델(discriminative model)과 생성모델(generative model)

일반적인 머신러닝은 판별 모델로 각각을 나누기 위한 선을 긋는다고 표현할 수 있다. 생성모델은 판별보다는 범위를 고려한다고 볼 수 있다.

판별모델과 생성모델

 VAE는 Generative Model 중 하나로, 확률분포( P(x) )를 학습함으로써, 데이터를 생성하는 것이 목적이다. Encoder 네트워크는 학습용 데이터(이하  x )를 입력으로 받고 잠재 변수(이하  z )의 확률분포에 대한 파라미터를 출력한다(Gaussian 정규 분포의 경우,  μ,σ2 ). 그리고 Decoder는 잠재변수에 대한 확률 분포  p(z) 에서 샘플링한 벡터를 입력받아, 이를 이용해 원본 이미지를 복원한다.

VAE는 최적화를 통해 아래의 두가지 문제를 푼다.

①주어진 데이터를 잘 설명하는 잠재 변수의 분포를 찾는 것이고, (Encoder의 역할)
②잠재변수로 부터 원본 이미지와 같은 이미지를 잘 복원하는 것이다. (Decoder의 역할)


1) Encoder

 Encoder의 역할은 데이터가 주어졌을 때 Decoder가 원래의 데이터로 잘 복원할 수 있는  z 를 샘플링 할 수 있는 이상적인 확률분포  p(z|x) 를 찾는 것이다. 하지만 어떤 것이 이상적인 확률분포  p(z|x)  인지는 아무도 모른다. VAE 방법론에서는 이 문제를 해결하기 위해 Variational inference를 사용한다.

Variational inference

우리가 이상적인 확률분포를 모르지만, 이를 추정하기 위해서 다루기 쉬운 분포(approximation class, 대표적으로 Gaussian distribution)를 가정하고 이 확률분포의 모수를 바꿔가며, 이상적인 확률분포에 근사하게 만들어 그 확률분포를 대신 사용하는 것이다. 이 다루기 쉬운 분포를 qϕ 라고 한다면, Encoder는 ϕ 라는 파라미터들을 바꾸어가며, qϕ(z|x) 확률 분포를 이상적인 확률 분포 p(z|x) 에 근사시키는 역할을 수행한다. 보통 qϕ(⋅) 은 Gaussian 정규 분포라고 가정한다. 이유는 뒤에서 설명 할 것이다. 그리고 이때, z 의 marginal distribution은 평균이 0이고 분산이 1인 표준 정규분포로 가정한다. 아래의 코드가 encoder부분을 구현 한 것이다.

Variational inference

class Encoder(layers.Layer):
  """Maps MNIST digits to a triplet (z_mean, z_log_var, z)."""

  def __init__(self,
               latent_dim=32,
               intermediate_dim=64,
               name='encoder',
               **kwargs):
    super(Encoder, self).__init__(name=name, **kwargs)
    self.dense_proj = layers.Dense(intermediate_dim, activation='relu')
    self.dense_mean = layers.Dense(latent_dim)
    self.dense_log_var = layers.Dense(latent_dim)
    self.sampling = Sampling()

  def call(self, inputs):
    x = self.dense_proj(inputs)
    # 가정한 확률분포의 모수
    z_mean = self.dense_mean(x)
    z_log_var = self.dense_log_var(x)
    z = self.sampling((z_mean, z_log_var))
    return z_mean, z_log_var, z
The reparameterization trick

Encoder가 출력하는 것은 qϕ(z|x) 확률 분포의 모수이다. 우리는 qϕ(z|x) 를 정규 분포라고 가정했기 때문에, 이 경우에는 평균과 분산이다. 다음 단계는 이 확률분포로 부터 샘플링을 하는 것이다. 이 때, 그냥 샘플링을 한다면, Back propagation이 불가능하다. Back propagation은 편미분을 구함으로써 Gradient를 구하는 것인데, z 를 확률분포에서 그냥 샘플링 한다면 체인룰이 중간에 끊기게 된다. 여기서는 이를 극복하기 위해서 Reparameterization trick을 사용했다. Reparameterization trick이란, 가우시안 정규 분포의 샘플을 추출하고 싶을 때, 아래의 식과 같이 샘플링을 하는 것을 말한다. 이렇게 샘플을 추출 하더라도 원래의 확률적 특성을 보존 한다. 그리고 이렇게 샘플링 하면 z 는 확률분포의 모수인 분산과 평균이 더해진 형태이므로 Back propagation 또한 가능하다.

zi,l~ N(μi,σ2i)→zi,l=μi+σ2i⊙ϵ 

ϵ~ N(0,1)

class Sampling(layers.Layer):
  """Uses (z_mean, z_log_var) to sample z, the vector encoding a digit."""

  def call(self, inputs):
    z_mean, z_log_var = inputs
    batch = tf.shape(z_mean)[0]
    dim = tf.shape(z_mean)[1]
    epsilon = tf.keras.backend.random_normal(shape=(batch, dim))
    return z_mean + tf.exp(0.5 * z_log_var) * epsilon

2) Decoder

 Decoder는 추출한 샘플을 입력으로 받아, 다시 원본으로 재구축하는 역할을 수행한다. 추후에 설명을 간편하기 위해서 이를  θ  라는 파라미터들을 가지는  gθ(⋅)  이라고 하겠다.

class Decoder(layers.Layer):
  """Converts z, the encoded digit vector, back into a readable digit."""

  def __init__(self,
               original_dim,
               intermediate_dim=64,
               name='decoder',
               **kwargs):
    super(Decoder, self).__init__(name=name, **kwargs)
    self.dense_proj = layers.Dense(intermediate_dim, activation='relu')
    self.dense_output = layers.Dense(original_dim, activation='sigmoid')

  def call(self, inputs):
    x = self.dense_proj(inputs)
    return self.dense_output(x)

 우리가 궁극적으로 알고 싶은 것은  p(x)  즉 실제 데이터의 분포를 알고 싶은 것이다. 이상적인 확률 분포 p(x|z)p(x|z)를 찾는다는 것은 ELBO(ϕ)를 최대화하는 것과 같다. 수식에 관련된 설명은 링크를 참고하라.

 VAE의 손실함수는 두 부분으로 구성되어 있다. 첫 번째는 오토인코더가 입력을 재구성하도록 만드는 일반적인 재구성 손실(reconstruction loss)이고, 두 번째는 가우시안 분포에서 샘플된 것 샅은 코딩을 가지도록 오토인코더를 제어하는 latent loss이다.

class VariationalAutoEncoder(tf.keras.Model):
  """Combines the encoder and decoder into an end-to-end model for training."""

  def __init__(self,
               original_dim,
               intermediate_dim=64,
               latent_dim=32,
               name='autoencoder',
               **kwargs):
    super(VariationalAutoEncoder, self).__init__(name=name, **kwargs)
    self.original_dim = original_dim
    self.encoder = Encoder(latent_dim=latent_dim,
                           intermediate_dim=intermediate_dim)
    self.decoder = Decoder(original_dim, intermediate_dim=intermediate_dim)

  def call(self, inputs):
    # self._set_inputs(inputs)
    z_mean, z_log_var, z = self.encoder(inputs)
    reconstructed = self.decoder(z)
    # Add KL divergence regularization loss.
    kl_loss = - 0.5 * tf.reduce_mean(
        z_log_var - tf.square(z_mean) - tf.exp(z_log_var) + 1)
    self.add_loss(kl_loss)
    return reconstructed
original_dim = 784
vae = VariationalAutoEncoder(original_dim, 64, 32)  #, input_shape=(784,)

optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)
mse_loss_fn = tf.keras.losses.MeanSquaredError()

loss_metric = tf.keras.metrics.Mean()

3) Dataset

(x_train, _), _ = tf.keras.datasets.mnist.load_data()
x_train = x_train.reshape(60000, 784).astype('float32') / 255

train_dataset = tf.data.Dataset.from_tensor_slices(x_train)
train_dataset = train_dataset.shuffle(buffer_size=1024).batch(64)

4) Training

# Iterate over epochs.
for epoch in range(1):
  print('Start of epoch %d' % (epoch,))

  # Iterate over the batches of the dataset.
  for step, x_batch_train in enumerate(train_dataset):
    with tf.GradientTape() as tape:
      # !!! uncomment the following two lines to use workaround and skip !!!
      # if step == 0 and epoch == 0:
      #   vae._set_inputs(x_batch_train)
      reconstructed = vae(x_batch_train)
      # Compute reconstruction loss
      loss = mse_loss_fn(x_batch_train, reconstructed)
      loss += sum(vae.losses)  # Add KLD regularization loss

    grads = tape.gradient(loss, vae.trainable_weights)
    optimizer.apply_gradients(zip(grads, vae.trainable_weights))

    loss_metric(loss)

    if step % 100 == 0:
      print('step %s: mean loss = %s' % (step, loss_metric.result()))

vae.save('vae')

참고자료

https://gist.github.com/RomanSteinberg/c4a47470ab1c06b0c45fa92d07afe2e3

https://datascienceschool.net/view-notebook/c5248de280a64ae2a96c1d4e690fdf79/

https://www.slideshare.net/ssuser06e0c5/variational-autoencoder-76552518

https://excelsior-cjh.tistory.com/187

728x90