아래 예시는 PyTorch를 활용하여 MNIST 데이터를 대상으로 **VAE(Variational Autoencoder)**를 간단히 구현하는 코드 예시입니다. 실제 프로젝트에서는 하이퍼파라미터나 네트워크 구조를 더 세부 조절하거나, 다양한 최적화 기법을 적용할 수 있지만, 여기서는 “ELBO를 어떻게 코드로 구현해서 학습시키는지”에 초점을 맞춘 기본 뼈대를 보여드립니다.

부록 A. PyTorch로 간단한 VAE 구현 예시

A.1 라이브러리 임포트 및 데이터셋 준비

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# 1) 하이퍼파라미터 설정
batch_size = 128
learning_rate = 1e-3
num_epochs = 10
z_dim = 20  # 잠재변수 차원

# 2) MNIST 데이터셋 로드 & 전처리
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))  # MNIST 평균, 표준편차
])

train_dataset = datasets.MNIST(
    root='data', train=True, download=True, transform=transform
)
test_dataset = datasets.MNIST(
    root='data', train=False, download=True, transform=transform
)

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")

A.2 VAE 모델 구현

A.2.1 인코더(Encoder)와 디코더(Decoder)

class Encoder(nn.Module):
    def __init__(self, input_dim=784, hidden_dim=400, z_dim=20):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.fc_mu = nn.Linear(hidden_dim, z_dim)
        self.fc_logvar = nn.Linear(hidden_dim, z_dim)

    def forward(self, x):
        """
        x: (batch_size, 784) 형태 (MNIST 이미지를 28x28 -> 784로 flatten)
        return: (mu, log_var)
        """
        h = torch.relu(self.fc1(x))
        mu = self.fc_mu(h)
        logvar = self.fc_logvar(h)
        return mu, logvar

class Decoder(nn.Module):
    def __init__(self, z_dim=20, hidden_dim=400, output_dim=784):
        super().__init__()
        self.fc1 = nn.Linear(z_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, output_dim)

    def forward(self, z):
        """
        z: (batch_size, z_dim) 형태
        return: (batch_size, 784)
        """
        h = torch.relu(self.fc1(z))
        x_recon = torch.sigmoid(self.fc2(h))  # 픽셀값 0~1 범위로
        return x_recon

A.2.2 전체 VAE 클래스

class VAE(nn.Module):
    def __init__(self, input_dim=784, hidden_dim=400, z_dim=20):
        super().__init__()
        self.encoder = Encoder(input_dim, hidden_dim, z_dim)
        self.decoder = Decoder(z_dim, hidden_dim, input_dim)

    def reparameterize(self, mu, logvar):
        """
        Reparameterization Trick:
        z = mu + sigma * eps,
        where eps ~ N(0, I)
        """
        std = torch.exp(0.5 * logvar)   # logvar = log(sigma^2)
        eps = torch.randn_like(std)     # std와 같은 shape의 표준정규 샘플
        return mu + eps * std

    def forward(self, x):
        """
        x: (batch_size, 784)
        return: x_recon, mu, logvar
        """
        mu, logvar = self.encoder(x)
        z = self.reparameterize(mu, logvar)
        x_recon = self.decoder(z)
        return x_recon, mu, logvar

A.3 ELBO 계산을 위한 Loss 함수

VAE를 학습할 때는 ELBO를 최대화(동등하게 Loss를 최소화)합니다.