[파이썬] PyTorch 확률론적 모델 및 VAE

PyTorch는 딥러닝과 확률론적 모델링을 모두 지원하는 강력한 프레임워크입니다. 이번 포스트에서는 PyTorch를 사용하여 확률론적 모델링과 변이형 오토인코더(Variational Autoencoder, VAE)를 구현하는 방법에 대해 알아보겠습니다.

확률론적 모델 (Probabilistic Models)

확률론적 모델은 데이터와 관련된 확률 분포를 모델링하는 데 사용됩니다. 이러한 모델은 데이터 생성과 변수 간의 관계를 표현할 수 있어 유용합니다. PyTorch는 확률론적 모델을 정의하고 학습하는 데 필요한 다양한 기능 및 클래스를 제공합니다.

변이형 오토인코더 (Variational Autoencoder, VAE)

변이형 오토인코더는 생성 모델링에 사용되는 확률론적 모델입니다. VAE는 데이터의 잠재 표현을 학습하는 동시에 이를 사용하여 새로운 샘플을 생성할 수 있습니다. VAE는 오토인코더의 변형으로, 잠재 변수에 대한 사전 분포와 재구성 오차를 최소화하는 방식으로 학습됩니다.

import torch
import torch.nn as nn

class VariationalAutoencoder(nn.Module):
    def __init__(self, input_dim, hidden_dim, latent_dim):
        super(VariationalAutoencoder, self).__init__()
        
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, latent_dim * 2)  # mean and variance
        )
        
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, input_dim),
            nn.Sigmoid()  # output activation
        )
        
    def reparameterize(self, mu, log_var):
        std = torch.exp(0.5 * log_var)
        eps = torch.randn_like(std)
        return mu + eps * std
    
    def forward(self, x):
        z_params = self.encoder(x)
        mu, log_var = torch.split(z_params, z_params.shape[1] // 2, dim=1)
        z = self.reparameterize(mu, log_var)
        x_hat = self.decoder(z)
        return x_hat, mu, log_var

위의 코드는 PyTorch를 사용하여 VAE를 구현하는 예시입니다. VariationalAutoencoder 클래스는 인코더와 디코더를 포함하고 있으며, forward 메서드를 통해 잠재 변수 z를 추론하고 입력 데이터를 재구성합니다. reparameterize 메서드는 잠재 변수를 샘플링하는 역할을 합니다.

손실 함수 및 학습

VAE의 손실 함수는 재구성 오차와 잠재 변수 z의 KL 다이버전스로 구성됩니다. 이러한 손실 함수는 nn.BCELoss와 KL 다이버전스를 사용하여 정의할 수 있습니다. 학습에는 주어진 데이터셋에 대해 해당 손실 함수를 최소화하는 방향으로 모델 파라미터를 업데이트하는 것이 일반적입니다.

def loss_function(x, x_hat, mu, log_var):
    reconstruction_loss = nn.BCELoss(reduction='sum')(x_hat, x)
    kl_divergence = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())
    return reconstruction_loss + kl_divergence

model = VariationalAutoencoder(input_dim, hidden_dim, latent_dim)
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

for epoch in range(num_epochs):
    for batch in data_loader:
        # Forward pass
        x, _ = batch
        x_hat, mu, log_var = model(x)
        
        # Compute loss
        loss = loss_function(x, x_hat, mu, log_var)
        
        # Backward pass and optimize
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

위의 코드는 VAE를 학습하는 간단한 예시입니다. 데이터셋을 로드한 후, 데이터 배치를 모델에 전달하여 재구성 및 KL 다이버전스 손실을 계산합니다. 그런 다음 손실에 대한 역전파와 옵티마이저를 사용하여 모델 파라미터를 업데이트합니다.

결론

PyTorch를 사용하여 확률론적 모델링과 변이형 오토인코더(VAE)를 구현하는 방법에 대해 살펴보았습니다. PyTorch의 다양한 기능과 클래스를 활용하여 복잡한 모델을 구현할 수 있으며, 데이터 생성 및 변형에 유용한 기술을 구현할 수 있습니다. 이를 토대로 새로운 모델 및 알고리즘을 개발하는 데 활용해보세요!