시작하기 앞서 틀린 부분이 있을 수 있으니, 틀린 부분이 있다면 지적해주시면 감사하겠습니다.

이전글에서는 variational autoencoder (VAE)에 대해 설명하였습니다. 이번글에서는 linear layer로 이루어진 vanilla VAE의 구현에 대해 설명하도록 하겠습니다. 학습에 사용한 데이터는 MNIST 데이터를 사용하였으며, 구현은 python의 PyTorch를 이용하였습니다. 그리고 VAE 구현 뿐 아니라, t-SNE를 통한 잠재 변수(latent variable) 가시화 및 잠재 변수들이 변함에 따라 생성되는 데이터가 어떠한지 살펴보도록 하겠습니다.

그리고 VAE에 대한 글은 여기, t-SNE, UMAP에 관한 글은 여기를 참고하시기 바랍니다. 이렇게 구현한 VAE의 코드는 GitHub에 올려놓았으니 아래 링크를 참고하시기 바랍니다(본 글에서는 모델의 구현과 잠재 변수 가시화에 초점을 맞추고 있기 때문에, 데이터 전처리 및 학습을 위한 전체 코드는 아래 GitHub 링크를 참고하시기 바랍니다).

오늘의 컨텐츠입니다.

1. Vanilla VAE 구현
2. VAE Loss
3. VAE 결과 및 t-SNE를 통한 잠재 변수 가시화

VAE 구현 및 잠재 변수(latent variable) 가시화

여기서는 기본적인 vanilla VAE의 구현 코드를 살펴보겠습니다. 코드는 PyTorch로 작성 되었으며, vanilla VAE의 전체적인 구조는 단순히 linear layer을 여러개 쌓은 모습입니다. 그리고 loss function과 reparameterization trick 함수를 추가적으로 구현을 해주어야 합니다. 한 줄씩 자세한 설명은 코드 아래쪽에 설명을 참고하시기 바랍니다.

class VAE(nn.Module):
    def __init__(self, config:Config, color_channel:int):
        super(VAE, self).__init__()
        self.height = config.height
        self.width = config.width
        self.hidden_dim = config.hidden_dim
        self.latent_dim = config.latent_dim
        self.dropout = config.dropout
        self.color_channel = color_channel
        
        self.encoder = nn.Sequential(
            nn.Linear(self.height*self.width*self.color_channel, self.hidden_dim),
            nn.ReLU(),
            nn.Dropout(self.dropout),

            nn.Linear(self.hidden_dim, self.hidden_dim//2),
            nn.ReLU(),
            nn.Dropout(self.dropout),
        )
        self.decoder = nn.Sequential(
            nn.Linear(self.latent_dim, self.hidden_dim//2),
            nn.ReLU(),
            nn.Dropout(self.dropout),

            nn.Linear(self.hidden_dim//2, self.hidden_dim),
            nn.ReLU(),
            nn.Dropout(self.dropout),

            nn.Linear(self.hidden_dim, self.height*self.width*self.color_channel),
            nn.Sigmoid()
        )
        self.fc_mu = nn.Linear(self.hidden_dim//2, self.latent_dim)
        self.fc_log_var = nn.Linear(self.hidden_dim//2, self.latent_dim)

    
    def reparameterization_trick(self, encoded):
        mu = self.fc_mu(encoded)
        log_var = self.fc_log_var(encoded)
        std = torch.exp(0.5*log_var)
        eps = torch.randn_like(std)

        return mu + std*eps, mu, log_var


    def forward(self, x):
        batch_size = x.size(0)
        x = x.view(batch_size, -1)
        output = self.encoder(x)
        z, mu, log_var = self.reparameterization_trick(output)
        output = self.decoder(z)

        return output, mu, log_var

모델 초기화
먼저 모델의 고정된 값을 초기화하여 hidden layer까지 초기화하는 부분입니다. GitHub 코드에 보면 config.json이라는 파일에 존재하는 변수 값들을 모델에 적용하여 초기화 하는 것입니다.

• 4, 5번째 줄: 학습 이미지를 모두 같은 크기로 전처리 하였을 때의 세로 가로 크기.
• 6, 7번째 줄: hidden layer의 차원 및 sampling 할 잠재 변수(latent variable)의 차원.
• 8번째 줄: overfitting (과적합)을 방지하기 위한 dropout 비율.
• 9번째 줄: 이미지 전처리를 하였을 때, color channel 수(흑백으로 처리를 했다면 1, 칼라로 처리 했다면 3).
• 11 ~ 19번째 줄: encoder를 정의, 첫 번째 hidden layer는 (data size * hidden dim), 두 번째 hidden layer는 (hidden dim * hidden dim//2)의 크기를 가짐.
• 20 ~ 31번째 줄: decoder를 정의, 첫 번째 hidden layer는 (latent variable dim * hidden dim//2), 두 번째 hidden layer는 (hidden dim//2 * hidden dim)의 크기, 세 번째 hidden layer는 (hidden dim z * data size)의 크기를 가짐(sampling 한 잠재 변수를 가지고 실제 데이터 생성하는 부분.)
• 32 ~ 33번째 줄: encoder를 통해 나온 결과를 넣는 부분, 각각의 레이어를 통해 나온 최종 결과는 평균, log 분산을 의미하며 레이어는 (hidden//2 * latent dim)의 크기를 가짐(잠재 변수가 존재하는 분포를 가정하고 학습하는 부분).

Reparameterization Trick
다음은 reparameterization trick 부분의 코드입니다. 이 부분에서는 encoder를 거쳐서 나온 결과가 잠재 변수가 존재하는 분포를 가정하기 위해 평균과 log 분산을 구하는 레이어를 통과합니다. 그리고 reparameterization trick을 통해 잠재 변수 z를 sampling 합니다.

• 37번째 줄: 잠재 변수가 존재하는 분포의 평균.
• 38번째 줄: 잠재 변수가 존재하는 분포의 log 분산.
• 39번째 줄: log 분산을 표준편차로 변환.
• 40번째 줄: 평균이 0, 분산이 1인 표준정규분포에서 noise 하나 sampling.
• 42번째 줄: reparameterization trick을 통해 sampling된 잠재 변수, 평균, log 분산 반환.

학습될 때 거치는 부분
다음은 forward 부분의 코드입니다. 이 부분은 모델을 정의 하고나서 실제로 데이터가 학습할 때 직접적으로 거치게 되는 부분입니다. Pytorch의 모델을 정의할 때 nn.Module을 상속 받기 때문에 자동으로 데이터가 forward라는 method를 거치게 됩니다. 그리고 PyTorch에서는 forward를 거치면 자동으로 데이터 backpropagation이 가능하게 구현이 되어있습니다.

• 46번째 줄: 데이터가 들어왔을 때 batch size를 저장.
• 47번째 줄: 데이터의 크기를 변화.
  (e.g. 데이터가 크기가 칼라일 경우 128 * 3 * 10 * 10 (batch * channel * height * width)이고, 이를 128 * 300으로 변환).
• 48번째 줄: encoder를 거쳐 latent variable 추출.
  (e.g. 위 예시의 128 * 300 데이터가 인코더를 거쳐 128 * 32 크기를 갖는 잠재 변수 추출).
• 49번째 줄: encoder를 거쳐 나온 결과를 reparameterization_trick 함수로 넣어서 잠재 변수가 존재하는 분포의 평균, log 분산을 구하여 잠재 변수 샘플링.
• 50번째 줄: 잠재 변수(latent variable)를 decoder를 거쳐 원레의 데이터와 유사하게 복구.
  (e.g. 128 * 32 크기의 잠재 변수를 원래 크기인 128 * 300으로 복구).
• 52번째 줄: decoder의 복구 결과와 잠재 변수를 t-SNE로 가시화하기 위해 학습된 잠재 변수 분포의 평균, 표준편차와 생성된 데이터를 반환.

VAE의 모델을 구현하였으니 학습을 하기 위한 loss function을 구현해보겠습니다. 이전글에서도 설명했듯이, VAE의 loss function은 크게 encoder와 관련된 regularization term, decoder와 관련된 reconstruction term으로 이루어져있습니다. 자세한 코드는 아래에서 확인해보겠습니다.

decoder_loss = nn.BCELoss(reduction='sum')

def VAE_loss(x, output, mu, log_var, decoder_loss):
    batch_size = x.size(0)
    x = x.view(batch_size, -1)

    BCE_loss = decoder_loss(output, x)
    KLD_loss = 0.5 * torch.sum((torch.square(mu) + torch.exp(log_var) - log_var - 1))

    return BCE_loss + KLD_loss

• 1번째 줄: deocoder를 Bernoulli distribution (베르누이 분포)으로 가정하였기 때문에 binary cross entropy (BCE) loss를 사용(redunction sum을 하여 평균값을 내지 않고 loss를 더하도록 설정).
• 4, 5번째 줄: 데이터가 들어왔을 때 batch size를 저장하고, 데이터 크기를 변환.
• 7번째 줄: decoder와 관련된 reconstruction loss 부분.
• 8번째 줄: encoder와 관련된 regularization loss 부분(Kullback-Leibler divergence 계산 식), loss를 mini batch에 따라 평균내지 않고 합하여 loss 계산.
• 10번째 줄: 최종 전체 loss 반환.

VAE 결과
위에서 구현한 VAE 모델이 잘 학습이 되었다면, 모델은 sampling한 잠재 변수 z에 대해서 데이터를 원래 가지고 있는 데이터처럼 비슷하게 생성을 해낼 것입니다. 그렇다면 VAE를 통해 생성된 이미지의 결과와 실제 인풋으로 들어간 이미지를 비교해보겠습니다.

아래 결과는 학습에 전혀 사용되지 않은 test 데이터의 결과입니다. 아래 결과에서 왼쪽은 VAE에 들어간 인풋 데이터, 오른쪽은 sampling 된 잠재 변수 z로부터 생성된 데이터입니다. Autoencoder 구현 글에서 나온 결과와 비교했을 때 VAE의 결과가 더 blur한 것을 확인할 수 있습니다. 이는 어쩔 수 없이 평균값의 형태로 나오는 vanilla VAE의 한계라고 볼 수 있습니다.

잠재 변수 z의 분포 가시화
실제 GitHub에 구현되어있는 코드에서 잠재 변수(latent variable)의 차원이 10차원으로 구현되어있습니다. 우리는 10차원의 잠재 변수를 가시화할 수 없기에 t-SNE를 이용하여 2차원으로 축소하여 가시화 한 결과입니다. VAE는 autoencoder와 다르게 데이터 생성을 하기 위해 잠재 변수 z의 분포를 가정하고 학습하는 것이기에, 잠재 변수들의 분포가 좀 더 명확히 구분되어있습니다.

위에서 가시화한 잠재 변수는 test 데이터를 모델의 encoder에 통과시켜 나온 평균과 표준편차를 이용하여 z를 sampling한 결과입니다. 이때 sampling 하기 위해 추출해야하는 표준정규분포를 따르는 노이즈는 각 데이터마다 다르게 random 추출을 통해 나온 값을 사용하였습니다. 자세한 코드는 GitHub의 train.py의 test 함수를 살펴보시기 바랍니다.

Walking in the latent space (from DCGAN)
이제 잠재 변수가 변화함에 따라 어떻게 생성되는 데이터가 바뀌는지 살펴보겠습니다. 이 부분은 DCGAN 논문에서 walking in the latent space라고 표현된 부분입니다(DCGAN과 관련된 글은 여기를 참고하시기 바랍니다). 즉 latent space (z space)가 부드럽게 변한다는 것을 의미합니다. 먼저 test data를 통해 나오는 잠재 변수 z를 숫자별로 다 평균을 내어줍니다. 따라서 최종적으로 0의 잠재 변수 평균값, 1의 잠재 변수 평균값 이렇게 하여 숫자 9까지 총 10개의 잠재 변수 값을 가지고 두 개의 잠재 변수를 선택하여 interpolate 해보겠습니다. 아래는 3의 잠재 변수 평균값과 7의 잠재 변수 평균값을 interpolate 한 값들에 대해 생성되는 데이터의 변화를 살펴본 결과입니다.

3과 7의 잠재 변수들 사이를 interpolate 하여 새롭게 만든 잠재 변수를 이용하여 데이터를 생성한 결과, 데이터의 모습이 3에서 7로 서서히 바뀌는 것을 볼 수 있습니다. 자세한 코드는 GitHub의 train.py의 test 함수를 살펴보시기 바랍니다.

지금까지 vanilla VAE 구현 코드와 잠재 변수(latent variable) 가시화를 해보았습니다. 학습 과정에 대한 전체 코드는 GitHub에 있으니 참고하시면 될 것 같습니다. 다음에는 VAE와 동일한 생성모델이자, 가장 많이 응용되고 있는 모델 중 하나인 generative adversarial network (GAN)에 관하 이야기 해보도록 하겠습니다.