파이토치를 활용한 GAN 딥러닝, VAE를 사용하여 얼굴 이미지 생성

최근 인공지능 분야에서 생성적 적대 신경망(GAN)과 변형 오토인코더(VAE)는 이미지 생성의 효율과 품질을 크게 향상시키는 중요한 기술로 자리 잡았습니다. 이 글에서는 GAN과 VAE의 기본 개념과 함께, 파이토치를 사용하여 얼굴 이미지를 생성하는 과정을 자세히 살펴보겠습니다.

1. GAN(Generative Adversarial Networks) 개요

Generative Adversarial Networks(GAN)는 두 개의 신경망인 생성자(Generator)와 판별자(Discriminator)가 서로 경쟁하며 학습하는 구조를 가지고 있습니다. 생성자는 실제와 유사한 이미지를 생성하려고 하고, 판별자는 생성된 이미지가 진짜인지 가짜인지를 판별합니다. 이 과정은 생성자가 판별자를 속이도록 학습하면서 점점 더 현실적인 이미지를 생성하는 데 도움을 주게 됩니다.

1.1 GAN의 동작 원리

GAN은 다음과 같은 두 가지의 네트워크로 구성됩니다:

  • 생성자(Generator): 랜덤 노이즈를 입력으로 받아 실제와 유사한 이미지를 생성합니다.
  • 판별자(Discriminator): 입력된 이미지가 진짜인지 가짜인지 분류합니다.

학습이 진행됨에 따라 생성자는 점차 품질 높은 이미지를 생성하고, 판별자는 더욱 정확하게 이미지를 분석합니다. 이 과정은 제로섬 게임의 형태로 진행되며, GAN 모델의 목표는 이 두 네트워크의 성능을 동시에 향상시키는 것입니다.

2. VAE(Variational Autoencoder) 개요

변형 오토인코더(VAE)는 이미지나 데이터의 잠재 공간(latent space)을 학습하여 새로운 데이터를 생성할 수 있도록 하는 모델입니다. VAE는 입력 데이터를 인코더를 통해 저차원 잠재 공간으로 변환한 후, 이 잠재 공간에서 샘플링하여 디코더를 사용하여 이미지를 재구성합니다. VAE는 확률적 모델로, 입력 데이터의 분포를 학습하고 이를 기반으로 새로운 샘플을 생성합니다.

2.1 VAE의 구조

VAE는 다음과 같은 세 가지 주요 구성 요소로 이루어집니다:

  • 인코더(Encoder): 입력 데이터를 잠재 변수로 변환합니다.
  • 샘플링(Sampling): 잠재 변수에서 샘플을 추출합니다.
  • 디코더(Decoder): 샘플된 잠재 변수를 사용하여 새로운 이미지를 생성합니다.

3. 프로젝트 목표와 데이터셋

이번 프로젝트의 목표는 GAN 및 VAE를 사용하여 실제와 유사한 얼굴 이미지를 생성하는 것입니다. 이를 위해 CelebA 데이터셋을 사용하겠습니다. CelebA 데이터셋에는 다양한 얼굴 이미지가 포함되어 있으며, GAN과 VAE의 성능을 측정하는 데 적합한 데이터셋입니다.

4. 환경 설정

이 프로젝트를 진행하기 위해 Python과 PyTorch 프레임워크가 필요합니다. 다음은 필요한 패키지 목록입니다:

pip install torch torchvision matplotlib

5. 파이토치로 GAN 구현하기

먼저 GAN 모델을 구현해 보겠습니다. GAN의 구조는 다음과 같은 단계로 구성됩니다:

  • 데이터셋 로딩
  • 생성자와 판별자 정의
  • 훈련 루프 설정
  • 결과 시각화

5.1 데이터셋 로딩

우선, CelebA 데이터셋을 다운로드하고 준비합니다.

import torchvision.transforms as transforms
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader

transform = transforms.Compose([
    transforms.Resize((64, 64)),
    transforms.ToTensor(),
])

dataset = ImageFolder(root='path_to_celeba', transform=transform)
dataloader = DataLoader(dataset, batch_size=64, shuffle=True)

5.2 생성자 및 판별자 정의

GAN의 생성자 및 판별자를 정의합니다.

import torch
import torch.nn as nn

class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()
        self.model = nn.Sequential(
            nn.Linear(100, 256),
            nn.ReLU(),
            nn.Linear(256, 512),
            nn.ReLU(),
            nn.Linear(512, 1024),
            nn.ReLU(),
            nn.Linear(1024, 3 * 64 * 64),
            nn.Tanh(),
        )

    def forward(self, z):
        z = self.model(z)
        return z.view(-1, 3, 64, 64)

class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.model = nn.Sequential(
            nn.Linear(3 * 64 * 64, 512),
            nn.LeakyReLU(0.2),
            nn.Linear(512, 256),
            nn.LeakyReLU(0.2),
            nn.Linear(256, 1),
            nn.Sigmoid(),
        )

    def forward(self, img):
        img_flat = img.view(img.size(0), -1)
        return self.model(img_flat)

5.3 훈련 루프 설정

이제 GAN의 훈련 과정을 구현합니다.

import torch.optim as optim

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
generator = Generator().to(device)
discriminator = Discriminator().to(device)

criterion = nn.BCELoss()
g_optimizer = optim.Adam(generator.parameters(), lr=0.0002, betas=(0.5, 0.999))
d_optimizer = optim.Adam(discriminator.parameters(), lr=0.0002, betas=(0.5, 0.999))

num_epochs = 50
for epoch in range(num_epochs):
    for i, (imgs, _) in enumerate(dataloader):
        imgs = imgs.to(device)
        batch_size = imgs.size(0)

        # 레이블 설정
        real_labels = torch.ones(batch_size, 1).to(device)
        fake_labels = torch.zeros(batch_size, 1).to(device)

        # 판별자 훈련
        d_optimizer.zero_grad()
        outputs = discriminator(imgs)
        d_loss_real = criterion(outputs, real_labels)
        d_loss_real.backward()

        z = torch.randn(batch_size, 100).to(device)
        fake_imgs = generator(z)
        outputs = discriminator(fake_imgs.detach())
        d_loss_fake = criterion(outputs, fake_labels)
        d_loss_fake.backward()
        
        d_loss = d_loss_real + d_loss_fake
        d_optimizer.step()

        # 생성자 훈련
        g_optimizer.zero_grad()
        outputs = discriminator(fake_imgs)
        g_loss = criterion(outputs, real_labels)
        g_loss.backward()
        g_optimizer.step()

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

5.4 결과 시각화

훈련된 생성자가 생성한 이미지를 시각화합니다.

import matplotlib.pyplot as plt

z = torch.randn(64, 100).to(device)
fake_images = generator(z).detach().cpu()

plt.figure(figsize=(8, 8))
for i in range(64):
    plt.subplot(8, 8, i + 1)
    plt.imshow(fake_images[i].permute(1, 2, 0).numpy() * 0.5 + 0.5)
    plt.axis('off')
plt.show()

6. 파이토치로 VAE 구현하기

이제 VAE를 구현해 보겠습니다. VAE의 구조는 GAN과 비슷하지만, 확률적인 접근 방식을 사용합니다. VAE의 구현 단계는 아래와 같습니다:

  • 데이터셋 준비
  • 인코더와 디코더 정의
  • 훈련 루프 설정
  • 결과 시각화

6.1 데이터셋 준비

데이터셋은 GAN을 사용할 때와 동일하게 로드합니다.

6.2 인코더와 디코더 정의

VAE의 인코더와 디코더를 정의합니다.

class VAE(nn.Module):
    def __init__(self):
        super(VAE, self).__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(3, 16, 4, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(16, 32, 4, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(32, 64, 4, stride=2, padding=1),
            nn.ReLU(),
        )
        self.fc_mu = nn.Linear(64 * 8 * 8, 128)
        self.fc_logvar = nn.Linear(64 * 8 * 8, 128)
        self.fc_decode = nn.Linear(128, 64 * 8 * 8)
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(64, 32, 4, stride=2, padding=1),
            nn.ReLU(),
            nn.ConvTranspose2d(32, 16, 4, stride=2, padding=1),
            nn.ReLU(),
            nn.ConvTranspose2d(16, 3, 4, stride=2, padding=1),
            nn.Sigmoid(),
        )

    def encode(self, x):
        h = self.encoder(x)
        h = h.view(h.size(0), -1)
        return self.fc_mu(h), self.fc_logvar(h)

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

    def decode(self, z):
        z = self.fc_decode(z).view(-1, 64, 8, 8)
        return self.decoder(z)

    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        return self.decode(z), mu, logvar

6.3 훈련 루프 설정

VAE의 훈련 과정을 구현합니다. VAE는 두 가지 손실을 사용하여 훈련되며, 원본 이미지와 복원된 이미지 간의 차이(재구성 손실)와 잠재 공간의 분포와 정규 분포 간의 차이(쿨백-라이블러 발산 손실)로 구성됩니다.

vae = VAE().to(device)
optimizer = optim.Adam(vae.parameters(), lr=0.0002)

num_epochs = 50
for epoch in range(num_epochs):
    for imgs, _ in dataloader:
        imgs = imgs.to(device)

        optimizer.zero_grad()
        reconstructed, mu, logvar = vae(imgs)

        re_loss = nn.functional.binary_cross_entropy(reconstructed.view(-1, 3 * 64 * 64), imgs.view(-1, 3 * 64 * 64), reduction='sum')
        kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
        loss = re_loss + kl_loss

        loss.backward()
        optimizer.step()

    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item()}')

6.4 결과 시각화

훈련된 VAE를 사용하여 이미지를 복원하고 시각화합니다.

with torch.no_grad():
    z = torch.randn(64, 128).to(device)
    generated_images = vae.decode(z).cpu()

plt.figure(figsize=(8, 8))
for i in range(64):
    plt.subplot(8, 8, i + 1)
    plt.imshow(generated_images[i].permute(1, 2, 0).numpy())
    plt.axis('off')
plt.show()

7. 결론

이 글에서는 파이토치를 활용하여 GAN과 VAE를 사용하여 얼굴 이미지를 생성하는 방법을 살펴보았습니다. GAN은 생성자와 판별자가 서로 경쟁하면서 점점 더 현실적인 이미지를 생성하도록 학습하는 반면, VAE는 잠재 공간의 분포를 학습하여 새로운 이미지를 생성합니다. 두 기술 모두 이미지 생성 분야에서 중요한 역할을 하고 있으며, 각기 다른 방법으로 놀라운 결과를 만들어낼 수 있습니다.

8. 추가 참고 자료