파이토치를 활용한 GAN 딥러닝, 사과와 오렌지

Generative Adversarial Networks (GANs)는 생성 모델의 일종으로, 두 개의 신경망(생성자와 판별자) 간의 경쟁을 통해 실제와 유사한 데이터를 생성하는 기술입니다.
이번 글에서는 GAN을 이용하여 사과와 오렌지 이미지를 생성하는 방법을 탐구하겠습니다. 이 과정에서 PyTorch 프레임워크를 사용하여 GAN을 구현할 것이며,
실습을 위한 파이썬 코드도 제공하겠습니다.

1. GAN이란?

GAN은 Ian Goodfellow에 의해 2014년에 제안된 모델로, 두 개의 인공신경망 구조가 서로 경쟁하여 학습합니다.
이러한 구조는 다음과 같은 두 부분으로 나눌 수 있습니다:

생성자(Generator): 무작위 노이즈를 입력으로 받아 실제와 유사한 데이터를 생성하는 역할을 합니다.
판별자(Discriminator): 입력된 데이터가 실제 데이터인지 생성자가 만든 가짜 데이터인지를 판별합니다.

GAN의 훈련 과정은 다음과 같습니다:

생성자가 무작위 노이즈를 통해 데이터를 생성합니다.

판별자는 실제 데이터와 생성된 데이터를 비교하여 진짜인지 가짜인지 판단합니다.

생성자는 판별자의 판단을 바탕으로 더 진짜 같은 데이터를 생성하기 위해 업데이트됩니다.

판별자는 더 정확하게 진짜와 가짜를 구분하기 위해 업데이트됩니다.

2. 데이터셋 준비

GAN을 훈련하기 위해서 사과와 오렌지 이미지를 포함하는 데이터셋을 준비해야 합니다. 이번 예제에서는 Kaggle이나 다른 오픈 데이터셋에서 사과와 오렌지 데이터를 수집할 예정입니다.
이미지 데이터는 같은 크기로 변경 및 정규화되며, 그 후 텐서로 변환되어야 합니다. 일반적으로 이미지를 (64, 64) 크기로 조정하고,
[-1, 1] 범위로 정규화하는 것이 일반적입니다.

2.1. 이미지 전처리

다음은 이미지 전처리 과정을 구현하는 파이썬 코드입니다:


import os
import numpy as np
import cv2
from torchvision import transforms
from PIL import Image
import torch

def load_images_from_folder(folder):
    images = []
    for filename in os.listdir(folder):
        img = cv2.imread(os.path.join(folder, filename))
        if img is not None:
            img = cv2.resize(img, (64, 64))
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            images.append(img)
    return np.array(images)

folder = 'path_to_your_dataset'
dataset = load_images_from_folder(folder)

transform = transforms.Compose([
    transforms.ToPILImage(),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])
tensor_images = [transform(Image.fromarray(img)).unsqueeze(0) for img in dataset]

images_tensor = torch.cat(tensor_images)

3. GAN 구조 구현

GAN을 구현하기 위해 먼저 생성자와 판별자를 정의해야 합니다.
생성자는 일반적으로 Fully Connected Layer와 Convolutional Layer를 사용하여 이미지를 생성합니다.
판별자는 Convolutional Layer를 사용하여 이미지의 진짜 여부를 판단합니다.
아래는 파이토치로 작성된 간단한 GAN 모델입니다.

3.1. 생성자 모델


import torch.nn as nn

class Generator(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(Generator, self).__init__()
        self.model = nn.Sequential(
            nn.Linear(input_dim, 128),
            nn.ReLU(),
            nn.Linear(128, 256),
            nn.ReLU(),
            nn.Linear(256, output_dim),
            nn.Tanh()  # Output range to [-1, 1]
        )
    
    def forward(self, z):
        img = self.model(z)
        return img.view(img.size(0), 3, 64, 64)  # Reshape for image output

3.2. 판별자 모델


class Discriminator(nn.Module):
    def __init__(self, input_dim):
        super(Discriminator, self).__init__()
        self.model = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=4, stride=2, padding=1),
            nn.LeakyReLU(0.2),
            nn.Conv2d(32, 64, kernel_size=4, stride=2, padding=1),
            nn.LeakyReLU(0.2),
            nn.Flatten(),
            nn.Linear(64 * 16 * 16, 1),
            nn.Sigmoid()  # Output range to [0, 1]
        )

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

4. GAN 훈련

GAN을 훈련시키기 위해 다음의 과정을 반복합니다.
생성자는 무작위 노이즈를 사용하여 이미지를 생성하고,
판별자는 생성된 이미지와 실제 이미지를 구분합니다.
이후 손실 함수를 바탕으로 두 모델을 각각 업데이트합니다.


import torch.optim as optim

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Hyperparameters
input_dim = 100
output_dim = 3 * 64 * 64
lr = 0.0002
num_epochs = 200

# Models and optimizers
generator = Generator(input_dim, output_dim).to(device)
discriminator = Discriminator(output_dim).to(device)
criterion = nn.BCELoss()
optimizer_G = optim.Adam(generator.parameters(), lr=lr)
optimizer_D = optim.Adam(discriminator.parameters(), lr=lr)

# Label for real and fake images
real_labels = torch.ones(batch_size, 1).to(device)
fake_labels = torch.zeros(batch_size, 1).to(device)

for epoch in range(num_epochs):
    for i, imgs in enumerate(dataloader):
        # Train Discriminator
        optimizer_D.zero_grad()
        real_imgs = imgs.to(device)
        real_loss = criterion(discriminator(real_imgs), real_labels)
        
        z = torch.randn(batch_size, input_dim).to(device)
        fake_imgs = generator(z)
        fake_loss = criterion(discriminator(fake_imgs.detach()), fake_labels)
        
        d_loss = real_loss + fake_loss
        d_loss.backward()
        optimizer_D.step()
        
        # Train Generator
        optimizer_G.zero_grad()
        g_loss = criterion(discriminator(fake_imgs), real_labels)
        g_loss.backward()
        optimizer_G.step()

        if (i + 1) % 100 == 0:
            print(f'Epoch [{epoch + 1}/{num_epochs}], Step [{i + 1}/{len(dataloader)}], '
                  f'D Loss: {d_loss.item():.4f}, G Loss: {g_loss.item():.4f}')

5. 결과 및 시각화

훈련이 완료되면 생성된 이미지를 시각화하여 성능을 평가할 수 있습니다.
아래는 생성된 이미지를 그리드 형태로 표시하는 파이썬 코드입니다.


import matplotlib.pyplot as plt

def show_generated_images(generator, num_images):
    z = torch.randn(num_images, input_dim).to(device)
    generated_images = generator(z)

    grid = torchvision.utils.make_grid(generated_images.cpu().detach(), nrow=5, normalize=True)
    
    plt.imshow(grid.permute(1, 2, 0))
    plt.axis('off')
    plt.show()

show_generated_images(generator, 25)

6. 결론

GAN을 활용한 사과와 오렌지 생성 모델을 구축하고 훈련시키는 과정에 대해 살펴보았습니다. 실질적인 데이터셋을 활용하여
PyTorch 프레임워크의 강력한 기능을 이용해 모델을 구현하는 방법을 배웠습니다. 다양한 분야에 응용 가능성이 있는 GAN의 힘을
경험한 만큼, 이를 통해 앞으로 더 발전된 모델을 만들어보시길 바랍니다.

더 많은 것들을 배우고 싶다면, GAN의 다양한 변종인 CycleGAN이나 StyleGAN을 학습해보는 것도 좋은 방법입니다.
이러한 고급 내용들을 통해 딥러닝 기술을 한층 더 넓혀가시길 바랍니다.