PyTorch로 GAN 구현하기: 기초부터 MNIST 이미지 생성까지 완벽 가이드

최근 생성형 AI가 큰 주목을 받으면서 많은 개발자가 이미지 생성 모델의 근간이 되는 GAN에 관심을 가지고 있습니다. 이번 포스팅에서는 복잡한 이론보다는 PyTorch를 사용하여 어떻게 GAN을 직접 구현할 수 있는지, 그 실무적인 흐름을 정리해 보았습니다. 임베디드 환경에서의 경량화나 고도화된 모델 연구에 앞서, GAN의 기본 구조를 이해하는 데 도움이 되기를 바랍니다.

Generated by Gemini AI.

핵심 요약

• GAN은 생성자(Generator)와 판별자(Discriminator)가 서로 경쟁하며 점진적으로 데이터 품질을 높이는 구조입니다.
• 학습의 핵심은 판별자는 진짜와 가짜를 정확히 구분하고, 생성자는 판별자를 속일 수 있을 만큼 정교한 데이터를 만드는 것입니다.
• 손실 함수와 최적화 과정에서 생성자와 판별자의 균형을 맞추는 것이 GAN 학습의 성패를 좌우합니다.

GAN의 기본 개념

GAN은 두 개의 신경망이 서로의 성능을 끌어올리는 제로섬 게임(Zero-sum game) 형태의 학습 방식을 취합니다.

구성 요소	역할	목표
생성자 (Generator)	랜덤 노이즈를 입력받아 가짜 이미지 생성	판별자가 진짜로 착각하게 만들기
판별자 (Discriminator)	이미지(진짜/가짜)를 판별	진짜와 가짜를 정확히 구분하기

GAN의 목적 함수는 다음과 같습니다.

$$\min_G \max_D V(D, G) = E_{x \sim P_{data}}[\log D(x)] + E_{z \sim P_z}[\log(1 - D(G(z)))]$$

판별자는 실제 데이터($x$)에는 높은 점수를, 생성된 데이터($G(z)$)에는 낮은 점수를 주려고 하며, 생성자는 반대로 판별자를 속이기 위해 노력합니다.

PyTorch를 활용한 GAN 구현

1. 라이브러리 및 데이터셋 준비

먼저 필요한 라이브러리를 불러오고, MNIST 데이터를 정규화하여 학습 준비를 마칩니다.

Python

 

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.datasets as datasets
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt

# 데이터 준비
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])
dataloader = DataLoader(datasets.MNIST(root='data', train=True, transform=transform, download=True), 
                        batch_size=64, shuffle=True)

2. 생성자와 판별자 설계

생성자는 노이즈 차원에서 픽셀 데이터 차원으로 확장하고, 판별자는 이를 다시 0과 1 사이의 확률값으로 압축합니다.

레이어 단계	생성자 구조 (Generator)	판별자 구조 (Discriminator)
입력	노이즈 벡터 (100)	이미지 (28x28)
중간층	Linear + ReLU (256, 512, 1024)	Linear + LeakyReLU (1024, 512, 256)
출력	28x28 (Tanh 활성화)	1 (Sigmoid 활성화)

Python

 

class Generator(nn.Module):
    def __init__(self, noise_dim=100):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(noise_dim, 256), nn.ReLU(),
            nn.Linear(256, 512), nn.ReLU(),
            nn.Linear(512, 1024), nn.ReLU(),
            nn.Linear(1024, 28*28), nn.Tanh()
        )
    def forward(self, z): return self.model(z).view(-1, 1, 28, 28)

class Discriminator(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(28*28, 1024), nn.LeakyReLU(0.2),
            nn.Linear(1024, 512), nn.LeakyReLU(0.2),
            nn.Linear(512, 256), nn.LeakyReLU(0.2),
            nn.Linear(256, 1), nn.Sigmoid()
        )
    def forward(self, x): return self.model(x.view(-1, 28*28))

3. 학습 루프 구현

학습 시 판별자는 진짜 데이터와 가짜 데이터를 각각 1과 0으로 학습시키고, 생성자는 가짜 데이터를 1로 학습시킵니다.

Python

 

# 모델 설정
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
generator = Generator().to(device)
discriminator = Discriminator().to(device)
g_optimizer = optim.Adam(generator.parameters(), lr=0.0002)
d_optimizer = optim.Adam(discriminator.parameters(), lr=0.0002)
criterion = nn.BCELoss()

# 학습 루프 생략 (위의 코드를 활용)

개발을 위한 팁

• 학습 균형 맞추기: 판별자가 너무 빨리 학습되면 생성자가 학습할 수 있는 그라디언트(Gradient)가 사라집니다. 판별자의 학습 속도를 생성자보다 조금 낮추는 것이 좋습니다.
• 활성화 함수 선택: 판별자에서는 ReLU 대신 LeakyReLU를 사용해야 학습 과정에서 '죽은 뉴런' 문제를 방지할 수 있습니다.
• 시각화 활용: 매 5~10 에포크마다 생성된 이미지를 직접 확인하며 모델이 데이터를 잘 생성하고 있는지 모니터링하세요.

흔히 하는 실수

• detach() 누락: 생성된 가짜 이미지를 판별자에 넣을 때 .detach()를 하지 않으면, 판별자 학습 시 생성자 네트워크까지 그라디언트가 흘러가 학습이 꼬이게 됩니다.
• 정규화 범위 착각: Tanh 함수를 사용했다면 데이터의 범위가 -1에서 1 사이입니다. 이 부분을 맞추지 않으면 이미지가 제대로 출력되지 않습니다.
• 학습률 설정: 너무 높은 학습률은 GAN의 학습을 불안정하게 만듭니다. 0.0002 정도의 낮은 학습률로 시작하는 것이 안정적입니다.

결론

지금까지 PyTorch를 이용한 기본적인 GAN 구현 방법을 살펴보았습니다. GAN은 기초적인 형태 이후에도 DCGAN, WGAN 등 더 강력한 모델로 확장될 수 있습니다. 이번 코드를 기반으로 나만의 데이터를 넣거나 구조를 변경하며 학습의 차이를 직접 경험해 보시기 바랍니다. 다음 포스팅에서는 조금 더 복잡한 구조인 DCGAN에 대해 자세히 다뤄보겠습니다.

도움이 되셨나요? 블로그의 다른 기술 포스팅들도 확인해 보세요.