파이토치를 활용한 GAN 딥러닝, VAE 만들기

1. 서론

인공지능의 발전과 함께 생성 모델(Generative Models)의 중요성이 커지고 있습니다. 생성 모델은 구조적으로 서로 다른 데이터를 생성하는 역할을 하며, 특히 GAN(Generative Adversarial Networks)과 VAE(Variational Autoencoder)가 널리 사용됩니다. 이 글에서는 파이토치(PyTorch)를 활용하여 GAN과 VAE를 구현하는 방법을 자세히 설명하겠습니다.

2. GAN(Generative Adversarial Networks)

GAN은 Ian Goodfellow가 2014년에 제안한 모델로, 두 개의 신경망(생성자와 판별자)이 서로 경쟁하며 학습합니다. 생성자는 가짜 데이터를 만들어내고, 판별자는 진짜 데이터와 가짜 데이터를 구분하는 역할을 합니다.

2.1 GAN의 구조

GAN은 다음과 같은 구조로 구성됩니다:

  • 생성자(Generator): 랜덤 노이즈를 입력으로 받아 진짜 데이터처럼 높은 품질의 가짜 데이터를 생성합니다.
  • 판별자(Discriminator): 입력 데이터를 보고 이 데이터가 진짜인지 가짜인지 판단합니다.

2.2 GAN 학습 과정

GAN의 학습 과정은 다음과 같은 단계를 포함합니다.

  1. 생성자는 랜덤 노이즈를 생성하여 가짜 데이터를 생성합니다.
  2. 판별자는 생성된 가짜 데이터와 진짜 데이터를 입력받아 각 클래스의 확률을 출력합니다.
  3. 생성자는 판별자가 가짜 데이터를 진짜라고 판단하도록 손실(loss)을 최소화하려고 합니다.
  4. 판별자는 진짜 데이터에 대해 높은 확률을, 가짜 데이터에 대해서는 낮은 확률을 출력하도록 손실을 최소화합니다.

2.3 GAN 구현 코드

다음은 간단한 GAN을 구현하는 파이썬 코드입니다:


import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms

# 생성자 클래스 정의
class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()
        self.model = nn.Sequential(
            nn.Linear(100, 256),
            nn.ReLU(inplace=True),
            nn.Linear(256, 512),
            nn.ReLU(inplace=True),
            nn.Linear(512, 1024),
            nn.ReLU(inplace=True),
            nn.Linear(1024, 784),
            nn.Tanh()
        )

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

# 판별자 클래스 정의
class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.model = nn.Sequential(
            nn.Linear(784, 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, img):
        return self.model(img)

# 데이터 로딩
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize([0.5], [0.5])
])
mnist = datasets.MNIST('data', train=True, download=True, transform=transform)
dataloader = torch.utils.data.DataLoader(mnist, batch_size=64, shuffle=True)

# 모델 초기화
generator = Generator()
discriminator = Discriminator()

# 손실 함수 및 최적화 알고리즘 설정
criterion = nn.BCELoss()
optimizer_G = optim.Adam(generator.parameters(), lr=0.0002, betas=(0.5, 0.999))
optimizer_D = optim.Adam(discriminator.parameters(), lr=0.0002, betas=(0.5, 0.999))

# GAN 학습
num_epochs = 100
for epoch in range(num_epochs):
    for i, (imgs, _) in enumerate(dataloader):
        # 실제 데이터 레이블 및 가짜 데이터 레이블 설정
        real_labels = torch.ones(imgs.size(0), 1)
        fake_labels = torch.zeros(imgs.size(0), 1)

        # 판별자 학습
        optimizer_D.zero_grad()
        outputs = discriminator(imgs.view(imgs.size(0), -1))
        d_loss_real = criterion(outputs, real_labels)
        d_loss_real.backward()

        z = torch.randn(imgs.size(0), 100)
        fake_imgs = generator(z)
        outputs = discriminator(fake_imgs.detach())
        d_loss_fake = criterion(outputs, fake_labels)
        d_loss_fake.backward()
        optimizer_D.step()

        # 생성자 학습
        optimizer_G.zero_grad()
        outputs = discriminator(fake_imgs)
        g_loss = criterion(outputs, real_labels)
        g_loss.backward()
        optimizer_G.step()

    print(f'Epoch [{epoch+1}/{num_epochs}], d_loss: {d_loss_real.item() + d_loss_fake.item()}, g_loss: {g_loss.item()}')
        

3. VAE(Variational Autoencoder)

VAE는 D. P. Kingma와 M. Welling이 2013년에 제안한 모델로, 확률적 방식으로 데이터를 생성합니다. VAE는 인코더와 디코더로 구성되어 있으며, 인코더는 데이터를 잠재 공간(latent space)으로 압축하고, 디코더는 이 latent space로부터 데이터를 재구성합니다.

3.1 VAE의 구조

VAE의 주요 구성 요소는 다음과 같습니다:

  • 인코더(Encoder): 입력 데이터를 잠재 벡터(latent vector)로 변환하며, 이 벡터는 정규분포를 따르도록 학습됩니다.
  • 디코더(Decoder): 잠재 벡터를 입력 받아 원본 데이터와 비슷한 출력을 생성합니다.

3.2 VAE 학습 과정

VAE의 학습 과정은 다음과 같습니다.

  1. 데이터를 인코더에 통과시켜 평균과 분산을 얻습니다.
  2. 재파라미터화 기법을 사용하여 샘플링합니다.
  3. 샘플링한 잠재 벡터를 디코더에 통과시켜 데이터를 재구성합니다.
  4. 재구성된 데이터와 원본 데이터 간의 손실을 계산합니다.

3.3 VAE 구현 코드

다음은 간단한 VAE를 구현하는 파이썬 코드입니다:


class VAE(nn.Module):
    def __init__(self):
        super(VAE, self).__init__()
        self.encoder = nn.Sequential(
            nn.Linear(784, 400),
            nn.ReLU()
        )
        self.fc_mu = nn.Linear(400, 20)
        self.fc_logvar = nn.Linear(400, 20)
        self.decoder = nn.Sequential(
            nn.Linear(20, 400),
            nn.ReLU(),
            nn.Linear(400, 784),
            nn.Sigmoid()
        )

    def reparametrize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std

    def forward(self, x):
        h1 = self.encoder(x.view(-1, 784))
        mu = self.fc_mu(h1)
        logvar = self.fc_logvar(h1)
        z = self.reparametrize(mu, logvar)
        return self.decoder(z), mu, logvar

# VAE 학습
vae = VAE()
optimizer = optim.Adam(vae.parameters(), lr=0.001)
criterion = nn.BCELoss(reduction='sum')

num_epochs = 10
for epoch in range(num_epochs):
    for imgs, _ in dataloader:
        optimizer.zero_grad()
        recon_batch, mu, logvar = vae(imgs)
        recon_loss = criterion(recon_batch, imgs.view(-1, 784))
        # Kullback-Leibler divergence
        kld = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
        loss = recon_loss + kld
        loss.backward()
        optimizer.step()
    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item()}')
        

4. 결론

GAN과 VAE는 각각 고유한 장점이 있으며, 다양한 생성적 작업에서 사용될 수 있습니다. 본 글에서는 파이토치를 활용하여 GAN과 VAE를 구현하는 방법에 대해 설명하였으며, 각 모델의 동작 원리를 이해하고 실제로 코드로 구현해보는 기회를 제공하였습니다. 생성 모델인 GAN과 VAE는 이미지 생성, 스타일 변환, 데이터 증강 등 다양한 분야에서 활용되고 있습니다. 이러한 모델들은 앞으로 더욱 발전할 가능성이 있으며, 인공지능 분야에서 중요한 역할을 하게 될 것입니다.