[Machine Learning] PyTorch로 ACGAN 구현하기 (Auxiliary Conditional GANs using PyTorch)

GAN(Generative Adversarial Networks)은 실제 데이터와 매우 유사한 가짜 데이터를 생성하는 데 사용되는 딥러닝 알고리즘이다.

GAN은 두 개의 네트워크 즉, Generator(생성기)와 Discriminator(판별기)으로 구성된다.

Generator는 가짜 데이터를 생성하고 Discriminator는 데이터가 진짜인지 가짜인지 구분한다.

제네레이터는 디스크리미네이터를 속일 수 있는 데이터를 생성하도록 훈련되고 디스크리미네이터는 실제 데이터와 가짜 데이터를 구별하도록 트레이닝된다.

 

GAN은 사실적인 이미지나 비디오 등 학습한 데이터와 유사한 데이터를 생성하는 데 효과적이다.

하지만 이런 GAN도 부족한 점이 있다.

GAN의 한계는 생성된 데이터의 속성을 제어하는 능력이 부족하다는 것이다.

즉, 생성된 정보가 훈련된 데이터의 확률 분포를 따라갈 뿐 내가 원하는 값의 정보를 생성하지는 못하는 것이다.

 

이러한 단점을 해결해주는 GAN을 알아보자.

---

GAN, Conditional GAN, Auxiliary Conditional GAN 차이

먼저 기존 GAN의 스트럭쳐를 살펴보자.

GAN Structure, 출처: https://github.com/hwalsuklee/tensorflow-generative-model-collections

GAN에서 제네레이터는 노이즈를 입력받아 페이크 데이터를 생성한다.

디스크리미네이터는 데이터를 입력받아 진짜와 가짜를 구분한다.

이때 중요한 점은 여러 클래스를 학습한 이상적인 GAN의 출력은 학습한 데이터의 확률 분포를 따를뿐, 제네레이터로 원하는 클래스의 출력을 얻기는 쉽지 않다는 것이다.

 

이런 이유로 조건을 주는 GAN이 만들어지게 된다.

Conditional GAN, 출처: https://github.com/hwalsuklee/tensorflow-generative-model-collections

CGAN은 각 네트워크가 데이터의 클래스나 카테고리에 해당하는 정보를 추가적으로 입력받는다.

이 정보에 따라 CGAN은 입력받은 클래스에 대한 특징을 가진 가짜 데이터를 생성할 수 있다.

 

오늘 살펴볼 GAN은 CGAN에 기능이 좀더 추가된 버전이다.

Auxiliary Conditional GAN, 출처: https://github.com/hwalsuklee/tensorflow-generative-model-collections

ACGAN(Auxiliary Conditional GAN)은 CGAN에서 디스크리미네이터에 기능을 추가한 버전이다.

ACGAN의 디스크리미네이터에는 기존 GAN에서 하는 역할인 실제 이미지와 가짜 이미지를 구별하는 역할에 더해 추가적으로 이미지를 각각의 클래스로 분류하는 두 가지 목표가 있다.

 

얼핏 보면 결국 제네레이터의 역할은 비슷해보이지만 ACGAN은 CGAN보다 좋은 점이 몇 가지 있다.

ACGAN은 CGAN보다 클래스 측면에서 더 나은 결과를 생성하는 것으로 알려져있다.

이는 ACGAN의 제네레이터가 디스크리미네이터를 속이는 것에 더해 올바른 클래스에 속하는 샘플을 생성하도록 훈련되기 때문이다.

이 외에도 ACGAN은 CGAN보다 수렴 속도가 빠르고 이는 ACGAN이 지도학습과 비지도학습 모두에 사용 가능하기 때문에 ACGAN이 CGAN보다 더 유연하다고 한다.(나도 ACGAN을 사용한 연구로 논문을 써서 졸업하기는 했지만 이 점은 잘 모르겠다.)

 

ACGAN에 대한 더 자세한 내용은 다음 논문을 참고하자.

Conditional GANs with Auxiliary Discriminative Classifier

Conditional GANs with Auxiliary Discriminative Classifier

 

아무튼 이제 직접 PyTorch로 ACGAN을 구현해보자.

---

PyTorch로 ACGAN 구현하기

ACGAN이 뭔가 복잡해보이지만 GAN과 달라지는 점은 별로 없다.

1. 각 네트워크의 입력에 클래스에 대한 정보를 넣는다.
   넣는 방법에 대한 정답은 없으므로 여러 코드를 보고 직접 돌리면서 결과를 확인해보자.
   (정답이 없어서 이것 저것 시도해봐야한다.)
2. 디스크리미네이터의 출력을 두 가지로 만든다.
   진짜 데이터와 제네레이터가 만든 데이터를 구분하는 기존 GAN의 기능과 소프트맥스 함수 등을 사용해 클래스를 구분하는 출력을 만든다.
3. 트레이닝 루프에서 기존 GAN에 있는 진짜 가짜에 대한 비용함수에 더해 클래스에 대한 비용함수를 추가한다.

결국 GAN에 다중 로지스틱 회귀를 추가한 버전이라고 생각하면 쉽다.

다중 로지스틱 회귀는 다음 글을 참고하자.

[Machine Learning] PyTorch로 다중 로지스틱 회귀 구현하기 (Multinomial Logistic Regression)

 
이제 직접 코드를 써보자.

필요한 라이브러리와 하이퍼 파라미터, 및 훈련에 사용할 MNIST 데이터를 준비하자.

import os
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
from torchvision.utils import save_image

if not os.path.exists("imgs"):
    os.makedirs("imgs")

# Device configuration
device = torch.device('cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu')
# Hyperparameters
latent_size = 100
num_epochs = 200
batch_size = 64
image_size = 28
num_classes = 10
lr = 0.0002
beta1 = 0.5

# MNIST dataset
transform = transforms.Compose([
    transforms.Resize(image_size),
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

train_dataset = datasets.MNIST(root='data/', train=True, transform=transform, download=True)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)

이제 네트워크를 만들자.
네트워크는 DCGAN을 기반으로 각 네트워크의 입력에 클래스에 대한 입력을 받을 수 있게 하고, 디스크리미네이터의 출력을 두 개로 만들었다.

# ACGAN Generator network
class Generator(nn.Module):
    def __init__(self, latent_size=100, num_classes=10):
        super(Generator, self).__init__()

        self.label_emb = nn.Embedding(num_classes, num_classes)
        self.fc = nn.Sequential(
            nn.Linear(latent_size+num_classes, 7*7*128),
            nn.BatchNorm1d(7*7*128),
            nn.ReLU(True)
        )
        self.upconv1 = nn.Sequential(
            nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(True)
        )
        self.upconv2 = nn.Sequential(
            nn.ConvTranspose2d(64, 1, kernel_size=4, stride=2, padding=1),
            nn.Tanh()
        )

    def forward(self, z, labels):
        inputs = torch.cat([z, self.label_emb(labels)], dim=1)
        x = self.fc(inputs)
        x = x.view(-1, 128, 7, 7)
        x = self.upconv1(x)
        x = self.upconv2(x)
        return x

class Discriminator(nn.Module):
    def __init__(self, num_classes=10):
        super(Discriminator, self).__init__()

        self.label_emb = nn.Embedding(num_classes, num_classes)
        self.conv1 = nn.Sequential(
            nn.Conv2d(1+num_classes, 64, kernel_size=4, stride=2, padding=1),
            nn.LeakyReLU(0.2, inplace=True)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.2, inplace=True)
        )
        self.fc_real_fake = nn.Linear(128*7*7, 1)
        self.fc_class = nn.Linear(128*7*7, num_classes)
        self.sigmoid = nn.Sigmoid()
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x, labels):
        c = self.label_emb(labels)
        c = c.unsqueeze(-1).unsqueeze(-1)
        c = c.repeat(1, 1, x.size(2), x.size(3))
        x = torch.cat([x, c], dim=1)
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(-1, 128*7*7)
        real_fake = self.sigmoid(self.fc_real_fake(x))
        class_pred = self.softmax(self.fc_class(x))
        return real_fake, class_pred

맨 밑 줄을 보면 디스크리미네이터가 진짜와 가짜를 구분하는 0에서 1사이의 출력을 가지는 real_fake와 클래스를 구분하는 class_pred가 있음을 확인할 수 있다.

 
비용함수와 옵티마이저, 그리고 학습 과정을 모니터링할 노이즈와 라벨(클래스)도 만들자.

# Initialize generator and discriminator
generator = Generator(latent_size, num_classes).to(device)
discriminator = Discriminator(num_classes).to(device)

# Initialize loss functions
adversarial_loss = nn.BCELoss()
auxiliary_loss = nn.CrossEntropyLoss()

# Initialize optimizers
optimizer_G = optim.Adam(generator.parameters(), lr=lr, betas=(beta1, 0.999))
optimizer_D = optim.Adam(discriminator.parameters(), lr=lr, betas=(beta1, 0.999))

# Fixed noise settings to monitor the training process
fixed_noise = torch.randn(100, 100).to(device)
fixed_gen_label = torch.tensor([ int(i/10) for i in range(100) ], dtype=torch.int64).to(device)

 
이제 트레이닝 루프만 만들면 된다. (간단하다.)

# Training loop
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        batch_size = images.size(0)
        images = images.to(device)
        labels = labels.to(device)

        # Adversarial ground truths
        valid = torch.ones(batch_size, 1).to(device)
        fake = torch.zeros(batch_size, 1).to(device)

        # Configure input
        z = torch.randn(batch_size, latent_size).to(device)

        # Generate fake images
        gen_labels = torch.randint(0, num_classes, (batch_size,)).to(device)
        fake_images = generator(z, gen_labels)

        # Train discriminator with real images
        real_pred, real_aux = discriminator(images, labels)
        d_loss_real = adversarial_loss(real_pred, valid)
        d_aux_real = auxiliary_loss(real_aux, labels)
        d_real_loss = d_loss_real + d_aux_real

        # Train discriminator with fake images
        fake_pred, fake_aux = discriminator(fake_images.detach(), gen_labels)
        d_loss_fake = adversarial_loss(fake_pred, fake)
        d_aux_fake = auxiliary_loss(fake_aux, gen_labels)
        d_fake_loss = d_loss_fake + d_aux_fake

        # Total discriminator loss
        d_loss = 0.5 * (d_real_loss + d_fake_loss)

        # Backward pass for discriminator
        optimizer_D.zero_grad()
        d_loss.backward()
        optimizer_D.step()

        # Train generator
        gen_labels = torch.randint(0, num_classes, (batch_size,)).to(device)
        fake_images = generator(z, gen_labels)
        fake_pred, fake_aux = discriminator(fake_images, gen_labels)
        g_loss_adv = adversarial_loss(fake_pred, valid)
        g_loss_aux = auxiliary_loss(fake_aux, gen_labels)
        g_loss = g_loss_adv + g_loss_aux

        # Backward pass for generator
        optimizer_G.zero_grad()
        g_loss.backward()
        optimizer_G.step()

        # Print progress
        if i % 100 == 0:
            print(f"Epoch [{epoch}/{num_epochs}], Batch [{i}/{len(train_loader)}], "
                  f"D_loss: {d_loss.item():.4f}, G_loss: {g_loss.item():.4f}")

    # Save generated images
    with torch.no_grad():
        save_image(generator(fixed_noise, fixed_gen_label), "imgs/%03d.png"%(epoch+1), nrow=10, normalize=True)

직접 돌려보고 싶으면 위에서부터 코드를 다 복사해서 돌리기만 하면 된다.

---

결과 확인

fixed_gen_label을 직접 프린트 해보면 알겠지만 0이 10개, 1이 10개, 2가 10개, ..., 9가 10개가 있는 길이가 100인 텐서이다.

이거를 노이즈와 같이 제네레이터에 넣어 트레이닝 루프의 맨 마지막에서 볼 수 있는 것처럼, 매 에포크마다 이미지를 뽑은 다음 저장을 한다.
이미지 폴더에 이미지가 굉장히 많이 쌓이게 되는데 이거를 다 보여주기는 힘들어서 gif 파일로 만들었다.

ACGAN output

저번에 만든 GAN의 출력은 중구난방이었지만 이번에 만든 ACGAN은 클래스를 줄 수 있었기 때문에 레이블을 입력한 대로 출력이 예쁘게 잘 정렬되어있다.