(Pytorch) 01.Adversarial Auto Encoder

Adversarial Auto Encoder

( 이론 참고 : https://seunghan96.github.io/bnn/37.Adversarial-Autoencoders-(2016)/ )

( 코드 참고 : https://github.com/eriklindernoren/PyTorch-GAN/blob/master/implementations/aae/aae.py )

1. Import Packages

AutoEncoder구현을 위한 주요 패키지들을 불러온다. ( Pytorch, Numpy … )

import numpy as np
import math
import os
import itertools

import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from torchvision import datasets
from torch.autograd import Variable

import torch.nn as nn
import torch.nn.functional as F
import torch

GPU 사용 여부

• torch.tensor( ).to(device)를 계속 사용하는 것을 피하기 위해

  Tensor = torch.cuda.FloatTensor를 사용한다.

cuda = True if torch.cuda.is_available() else False
Tensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor

2. Set hyperparameters

pytorch 공식 깃헙에서 사용한 default 값으로 hyperparameter를 설정한다.

n_epochs = 100
batch_size = 64
lr = 0.0002
b1 = 0.5
b2 = 0.999
n_cpu = 8
latent_dim = 10
img_size = 32
channels = 1
sample_interval = 400

img_shape = (channels, img_size, img_size)

3. Adversarial Auto Encoder 구현

Adversarial Auto Encoder(AAE)를 구현하기 위해서, 다음과 같은 함수 및 클래스들을 구현해줘야 한다.

• 1) reparam : reparameterization trick을 사용하기 위한 함수

• 2) Encoder : \(X \rightarrow Z\) 로 차원 축소해주는 인코더

• 3) Decoder : \(Z \rightarrow X'\) 로 복원하려는 디코더

• 4) Discriminator : AAE가 AE와 다른 점으로, discriminator가 존재한다.

  • positive sample \(p(\mathbf{z})\) vs negative sample \(q(\mathbf{z})\)

  • regularizer의 역할을 한다

    ( 차원축소되어 생성된 \(q(\mathbf{z})\) 로 하여금, prior인 \(p(\mathbf{z})\)와 비슷하도록 유도한다 )

• 5) 그 밖의 loss function들과 optimizer들

  • adversarial_loss : discriminator의 판별 값과, 정답값(0 or 1) 사이의 loss
  • pixelwise_loss : 실제 \(X\)와 복원된 \(X'\)간의 reconstruction loss
  • opt_G,opt_D : 각각 Generator와 Discriminator의 optimizer
    • Generator : Encoder & Decoder
    • Discriminator : Discrimnator

3-1. Reparameterization trick

reparameterization trick을 사용하기 위한 함수

• input : \(\mu\), \(\text{log}(\sigma^2)\)
• output : \(\mu + \sigma \odot \epsilon\) where \(\epsilon \sim N(0,1)\)

def reparam(mu,logvar):
  std = torch.exp(logvar/2)
  z =  Variable(Tensor(np.random.normal(0, 1, (mu.size(0), latent_dim))))
  return mu + z*std

Reparameterization trick을 구현한 코드들을 보면, 대부분 \(\sigma^2\)말고 \(\text{log}\sigma^2\)를 사용한 경우를 종종보게 된다. 이는 r.v의 variance가 positive해야한다는 조건 때문에 그런 것은 이해할 수 있다. 하지만 왜 \(\log{\sigma}\)를 사용하지 않고 왜 \(\text{log}\sigma^2\)를 사용할까?

( 출처 : https://stats.stackexchange.com/questions/486203/why-we-learn-log-sigma2-in-vae-reparameterization-trick-instead-of-standar )

1. It doesn’t make any real difference; since \(\text{log}\sigma^2=2\text{log}\sigma\); learning one is as easy as learning the other
2. It’s traditional in statistics to think of \(\sigma^2\) as the second parameter of a Normal distribution (rather than \(\sigma\) ).
3. There’s a simple unbiased estimator for \(\sigma^2\) but not for \(\sigma\)
4. The math for representing the Normal as a two-parameter exponential family is slightly simpler as \((\mu,\sigma^2)\) than \((\mu,\sigma)\)

3-2. Encoder

\(X \rightarrow Z\) 로 차원 축소해주는 인코더

• Hidden layer의 unit 개수 = 512

• Activation function : LeakyReLU

  ( nn.LeakyReLU( .... ,inplace=True)에서, inplace=True 의미 ?

  들어오는 값에 직접적으로 해당 함수를 적용한다! ( 별도의 output을 저장 X )

class Encoder(nn.Module):
  def __init__(self):
    super(Encoder,self).__init__()
    self.model = nn.Sequential(    
        nn.Linear(int(np.prod(img_shape)), 512),
        nn.LeakyReLU(0.2, inplace=True),
        nn.Linear(512, 512),
        nn.BatchNorm1d(512),
        nn.LeakyReLU(0.2, inplace=True),
      )
    self.mu = nn.Linear(512, latent_dim)
    self.logvar = nn.Linear(512, latent_dim)

  def forward(self,x):
    x_flatten = x.view(x.shape[0], -1)
    out = self.model(x_flatten)
    mu = self.mu(out)
    logvar = self.logvar(out)
    z = reparam(mu,logvar)
    return z

3-3. Decoder

\(Z \rightarrow X'\) 로 복원하는 디코더

• Hidden layer & Activation Function은 Encoder와 동일

class Decoder(nn.Module):
  def __init__(self):
    super(Decoder,self).__init__()
    self.model = nn.Sequential(
        nn.Linear(latent_dim, 512),
        nn.LeakyReLU(0.2, inplace=True),
        nn.Linear(512, 512),
        nn.BatchNorm1d(512),
        nn.LeakyReLU(0.2, inplace=True),
        nn.Linear(512, int(np.prod(img_shape))),
        nn.Tanh()
        )
  
  def forward(self,z):
    x_flatten = self.model(z)
    x = x_flatten.view(x_flatten.shape[0],*img_shape)
    return x

3-4. Discriminator

기존의 AutoEncoder와의 차이점이 바로 이 “Discriminator”의 존재에 있다.

이 Discriminator는, (1) 랜덤하게 샘플된 \(z\)와 (2) 실제 input \(x\)가 차원축소된 \(z\)를 잘 구분하도록 만드는 함수이다. 마지막 activation function은 sigmoid로, 0~1사이 값을 반환하게끔 한다.

2가지 Loss

• real loss : 임의로 샘플된 \(z\)가 인코딩된 것 vs 1로 채워진 벡터

• fake loss : \(x\)가 인코딩된 \(z\) vs 0으로 채워진 벡터

  ( 임의로 샘플된 \(z\)는 1에 가깝도록, \(x\)가 인코딩된 \(z\)는 0에 가깝도록 유도한다 )

  이 2가지 loss를 minimize하도록 Discriminator는 학습된다.

class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()

        self.model = nn.Sequential(
            nn.Linear(latent_dim, 512),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(512, 256),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(256, 1),
            nn.Sigmoid(),
        )

    def forward(self, z):
        validity = self.model(z)
        return validity

이에 반해, Generator ( Encoder & Decoder )는 loss에서는 다음의 term

adversarial_loss(discriminator(encoded_imgs), ONES)이 들어있어서, 자기가 (\(x\)를 \(z\)로 차원축소 시켜서) 생성해낸 것을 진짜가 1인 컷처럼 ( 즉, 임의로 샘플링된 \(z\)처럼 보이도록 ) 유도함으로써 regularizer의 효과를 가진다! ( 즉, 차원축소되어 생성된 \(q(\mathbf{z})\) 로 하여금, prior인 \(p(\mathbf{z})\)와 비슷하도록 유도한다 )

3-5. Loss Function

Loss function으로는,

• Discriminator가 판별하는데에 사용하는 함수인 Binary Cross Entropy와
• Reconstruction error를 측정하는 mean absolute error (MAE) 를 사용한다.

adversarial_loss = torch.nn.BCELoss()
pixelwise_loss = torch.nn.L1Loss()

지금까지 위에서 구현한 (1) 모델들과 (2) loss function을 생성하고, (if GPU 사용) gpu로 보낸다.

encoder = Encoder()
decoder = Decoder()
discriminator = Discriminator()

if cuda:
    encoder.cuda()
    decoder.cuda()
    discriminator.cuda()
    adversarial_loss.cuda()
    pixelwise_loss.cuda()

3-6. Optimizer

Optimizer로는 Adam을 사용한다.

( lr는 익숙한 파라미터지만, beta1,2는 생소하다. 이 둘은, weight decay rate에 영향을 주는 파라미터로써, 자세한 것은 찾아보길 바란다 :) )

opt_G = torch.optim.Adam(itertools.chain(encoder.parameters(), decoder.parameters()), lr=lr, betas=(b1, b2))

opt_D = torch.optim.Adam(discriminator.parameters(), lr=lr, betas=(b1, b2))

4. Train

for epoch in range(1,n_epochs+1):
  for i, (imgs, _) in enumerate(dataloader):
      ONES = Variable(Tensor(imgs.shape[0], 1).fill_(1.0), requires_grad=False)
      ZEROS = Variable(Tensor(imgs.shape[0], 1).fill_(0.0), requires_grad=False)

      # Configure input
      real_imgs = Variable(imgs.type(Tensor))

      # -----------------
      #  Train Generator
      # -----------------
      opt_G.zero_grad()
      encoded_imgs = encoder(real_imgs)
      decoded_imgs = decoder(encoded_imgs)

      # Generator Loss
      g_loss = 0.001 * adversarial_loss(discriminator(encoded_imgs), ONES) + 0.999 * pixelwise_loss(decoded_imgs, real_imgs)
      g_loss.backward()
      opt_G.step()

      # ---------------------
      #  Train Discriminator
      # ---------------------
      opt_D.zero_grad()
      z = Variable(Tensor(np.random.normal(0, 1, (imgs.shape[0], latent_dim))))

      # Discriminator Loss
      real_loss = adversarial_loss(discriminator(z), ONES)
      fake_loss = adversarial_loss(discriminator(encoded_imgs.detach()), ZEROS)
      d_loss = 0.5 * (real_loss + fake_loss)
      d_loss.backward()
      opt_D.step()

      if epoch%5==0 & i==len(dataloader):
        print(
          "[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]"
          % (epoch, n_epochs, i, len(dataloader), d_loss.item(), g_loss.item())
          )