[Pytorch] 적대적 생성 모델(GAN) 구현하기 - MINST를 기반으로

다음 블로그의 글을 참고하였습니다.

https://ddongwon.tistory.com/124

[Pytorch] GAN 구현 및 학습

대표적인 손글씨 데이터셋인 MNIST를 기반으로 적대적 생성모델(GAN)을 구현하겠습니다.

 

0. GAN 개요

여기에서 GAN의 자세한 설명을 참조할 수 있습니다.

 

식별자(Discriminator) :

 - 입력받은 데이터가 실제 데이터인지, 위조 데이터인지 식별합니다.

 - 실제값을 받았을 때 참(1)에 가깝게, 위조값을 받았을 때 거짓(0)에 가깝게 출력하는 것을 목표로 합니다.

 

생성자(Generator) :

 - 입력받은 noise vector 를 기반으로 위조 데이터를 생성합니다.

 - 출력값을 discriminator에 입력했을 때 참(1)에 가까운 값이 출력되도록 하는 것을 목표로 합니다.

 

지금과 같은 경우 GAN을 이용해 이미지를 생성하기 때문에 generator가 생성하는 위조값은 이미지가 됩니다.

 

1. 모델 설정

Discriminator와 generator의 신경망 계층은 서로 진행 방향이 반대가 되도록 했습니다.

Discriminator

 - 입력값 : 길이 784(28×28)의 MNIST 이미지 데이터

 - L1 : 길이 512로 변환하는 선형계층 (활성화 함수 = Leaky Relu)

 - L2 : 길이 256으로 변환하는 선형계층 (활성화 함수 = Leaky Relu)

 - L3 : 길이 1로 변환하는 선형계층 (활성화 함수 = Sigmoid)

 - 출력값 : [0, 1] 범위의 실수

 

Generator

 - 입력값 : 100×1 크기의 noise vector

 - L1 : 길이 256로 변환하는 선형계층 (활성화 함수 = Relu)

 - L2 : 길이 512으로 변환하는 선형계층 (활성화 함수 = Relu)

 - L3 : 길이 784로 변환하는 선형계층 (활성화 함수 = Tanh)

 - 출력값 : 784×1 크기의 image tensor

2. Hyperparameter 설정

Loss function : BCE Loss (label이 0 혹은 1 두 가지 밖에 없으므로)

Optimizer : Adam

Max epoch : 200

batch size : 100

Learning rate : 0.0002

데이터셋 전처리 : dataset -> ToTensor() -> Normal(0.5, 0.5)

 

3. 코드

3.1. Import

import time
import os
import torch
import torch.optim as optim
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import MNIST
from torchvision.utils import save_image

time : 프로그램의 실행시간을 측정합니다.

os : 프로젝트 내의 경로(directory)를 다루기 위해 사용합니다.

torch : pytorch의 라이브러리 이름입니다.

torchvision : 이미지처리와 관련하여 별도로 묶여진 pytorch의 라이브러리 입니다.

3.2. Preprocess

# Device setting
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print("current device : {}".format(device))

# Output Directory setting
res_path = "results"
if not os.path.exists(res_path):
    os.makedirs(res_path)

# Hyperparamter
max_epoch = 200
batch_size = 100
lr = 0.0002

img_size = 784
noise_size = 100
hidden_size1 = 256
hidden_size2 = 512

loss = nn.BCELoss()

Device setting :

gpu를 사용 가능한 환경이라면 cuda 라이브러리를 사용합니다.

그렇지 않으면 cpu만 사용합니다.

 

Output Directory setting :

학습시 generator가 생성할 이미지를 저장할 경로를 설정합니다.

폴더의 이름은 res_path의 이름으로 지정합니다.

 

Hyperparameter :

1.절에서 설정한 모델 안의 계층 크기와

2.절에서 설정한 hyperparameter의 값을 정합니다.

3.3. Generator & Discriminator

class Generator(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = nn.Linear(noise_size, hidden_size1)
        self.linear2 = nn.Linear(hidden_size1, hidden_size2)
        self.linear3 = nn.Linear(hidden_size2, img_size)
        self.relu = nn.ReLU()
        self.tanh = nn.Tanh()
    
    def forward(self, x):
        x = self.relu(self.linear1(x))
        x = self.relu(self.linear2(x))
        x = self.tanh(self.linear3(x))
        
        return x

class Discriminator(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = nn.Linear(img_size, hidden_size2)
        self.linear2 = nn.Linear(hidden_size2, hidden_size1)
        self.linear3 = nn.Linear(hidden_size1, 1)
        self.relu = nn.LeakyReLU(0.2)
        self.sigmoid = nn.Sigmoid()
    
    def forward(self, x):
        x = self.relu(self.linear1(x))
        x = self.relu(self.linear2(x))
        x = self.sigmoid(self.linear3(x))
        
        return x

__init__( )에서는 각 계층과 활성화 함수를 정의하고 forward( )에서는 계층 순서대로 계산해줍니다.

Pytorch의 구조상 class의 instance를 호출하면 그 class의 forward( )가 바로 호출됩니다.

generator = Generator()

generator(args)  # = generator.forward(args)

3.4. Dataset

# Dataset - MNIST
transform = transforms.Compose([transforms.ToTensor(),
                                transforms.Normalize(0.5, 0.5)])

dataset = MNIST(root='../data',
                train=True,
                transform=transform,
                download=True)

torchvision.dataset.MNIST를 가져옵니다.

3.1.절 때 torchvision.dataset에서부터 MNIST를 import 했기 때문에 MNIST 라고만 쓸 수 있습니다.

 

MNIST(root, train, transform, download)		MNIST 데이터셋을 반환한다
root	string	데이터셋이 저장될 경로
train	bool	train용 데이터셋인지 (false : test용 데이터셋)
transform	transform	전처리 함수들
download	bool	데이터셋이 없다면 내려받을 것인지

 

MNIST가 설치되지 않았을 경우 실행시 다음과 같이 자동으로 내려받아집니다.

 

 

전처리는 두 단계를 거칩니다.

마찬가지로 3.1.절 때 torchvision에서부터 transforms를 import 했기 때문에 transforms라고만 쓸 수 있습니다.

ToTensor( )

데이터셋을 pytorch의 tensor 자료형으로 변환한다.

 

Normalize(mean, std)		정규분포로 변환한다.
mean	float	평균
std	float	표준편차

 

3.5. Initializing

# Dataloader
dataloader = DataLoader(dataset=dataset,
                        batch_size=batch_size,
                        shuffle=True)

# Model
generator = Generator().to(device)
discriminator = Discriminator().to(device)

# Optimizer
optim_G = optim.Adam(generator.parameters(), lr=lr)
optim_D = optim.Adam(discriminator.parameters(), lr=lr)

start = time.time()​

Dataloader :

MNIST를 데이터셋으로 하고 batch_size(100)를 가진 Dataloader를 생성합니다.

Model :

Generator( ) / Discriminator( )의 instance인 generator / discriminator를 선언합니다.

Optimizer :

각각 generator / discriminator의 매개변수를 관리하는 Adam optimizer를 생성합니다.

Time :

마지막으로 학습을 시작하기 전 시작 시간을 기록합니다.

 

3.6. Epoch iteration

1 epoch 마다 데이터셋 전체를 사용합니다. (Dataset iteration에 대해서는 3.7.절에서 설명합니다)

for epoch in range(max_epoch):
    print("epoch: {}".format(epoch+1))
    
    #-------------------#
    # Dataset iteration #
    #-------------------#

    # print performance
    print(" Epoch {}'s discriminator performance : {:.2f}  generator performance : {:.2f}"
          .format(epoch, performance_D, performance_G))

    # Save fake images in each epoch
    res_imgs = fake_imgs.reshape(batch_size, 1, 28, 28)
    save_image(res_imgs, os.path.join(res_path, 'GAN_result {}.png'.format(epoch + 1)))

elasped_time = time.time() - start
print("[Done] Time performance : {}".format(elasped_time))

각 dataset iteration이 끝난 후,

- 각 epoch마다의 performance를 출력하고 (performance에 대해서는 3.7.절에서 설명합니다)

- 생성된 fake images를 res_path의 경로에 저장합니다. 

- 이 때 fake image의 모양은 784×1 이기 때문에 원래 이미지 모양인 28×28로 바꿔야 합니다.

 

Epoch iteration이 모두 끝난 후, 종료시 시각을 측정하여 시작시 시각과의 차이를 계산합니다.

이를 time performance로 하여 출력합니다.

 

3.7. Dataset iteration

1개 batch를 처리할 때 진행되는 흐름입니다.

붉은 부분이 discriminator의 학습 단계이고, 푸른 부분이 generator의 학습 단계입니다.

보라색 부분은 discriminator와 generator가 공통으로 거치는 단계입니다.

 

GAN에서는 discriminator와 generator를 따로 학습시킵니다.

순서는 상관 없으나, 둘 중 하나의 모델이 갱신된 후 다른 하나가 갱신되어야 합니다.

 

보라색 부분은 공통 부분이지만, 한 번 선언해서 양 쪽에 중복으로 사용할 수는 없습니다.

실제로 보라색 부분은 한 번만 선언하고 실행하면 다음과 같은 에러가 출력됩니다.

이유는 3.7.3.절에서 설명합니다.

 

 

3.7.1. Batch initialization

    # 3.7. dataset iteration
    for i, (imgs, label) in enumerate(dataloader):
        label_real = torch.full((batch_size, 1), 1, dtype=torch.float32).to(device)
        label_fake = torch.full((batch_size, 1), 0, dtype=torch.float32).to(device)
        
        real_imgs = imgs.reshape(batch_size, -1).to(device)  # B * 784
    
        ## 3.7.2. generator training ##
        
        ## 3.7.3. discriminator training ##
        
        ## 3.7.4 performance ##

Dataloader는 호출값 하나 당 실제값을 나타내는 data와 그 데이터의 이름을 나타내는 label로 이루어져 있습니다.

위에서는 이를 (imgs, label)이라는 이름으로 받고 있습니다.

물론 반복 당 batch 하나가 처리되기 때문에 (imgs, label)은 데이터 하나가 아니라 batch_size개 입니다.

 

label_real : 1(참)으로 이루어져 있는 라벨 vector입니다.

label_fake : 0(거짓)으로 이루저여 있는 라벨 vector입니다.

real_imgs : batch_size × 784 크기의 MINST 데이터 batch입니다.

 

3.7.2. Generator training

        ## 3.7.2. Generator training ##
        optim_G.zero_grad()
        optim_D.zero_grad()
    
        z = torch.randn(batch_size, noise_size).to(device)
        fake_imgs = generator(z)
        
        loss_G = loss(discriminator(fake_imgs), label_real)
        
        loss_G.backward()
        optim_G.step()

먼저 generator / discriminator 매개변수의 미분값을 초기화합니다.

batch_size × noise_size 크기의 noise vector z를 만든 후

generator에 입력해 fake_imgs를 만들어냅니다. (보라색 부분)

 

Generator가 만든 이미지는 실제값과 비슷할 수록 좋습니다.

따라서 discriminator의 결과값이 1에 가까워야 하는 것이 목표입니다.

fake_imgs를 label_real과 대조해 discriminator에 입력하여 loss_G를 구합니다.

 

마지막으로 loss_G에 대해 backproagation을 수행한 뒤 매개변수를 갱신합니다.

 

3.7.3. Discriminator training

        ## 3.7.3. Discriminator training ##
        optim_G.zero_grad()
        optim_D.zero_grad()

        z = torch.randn(batch_size, noise_size).to(device)
        fake_imgs = generator(z)
        
        loss_fake = loss(discriminator(fake_imgs), label_fake)
        loss_real = loss(discriminator(real_imgs), label_real)
        loss_D = (loss_fake + loss_real) / 2
        
        loss_D.backward()
        optim_D.step()

똑같이 generator / discriminator 매개변수의 미분값을 초기화합니다.

batch_size × noise_size 크기의 noise vector z를 만든 후

generator에 입력해 fake_imgs를 만들어냅니다. (보라색 부분)

 

Discriminator는 실제값을 참으로, 위조값을 거짓으로 구별할수록 좋습니다.

따라서 discriminator의 결과값이 실제일 때는 1, 위조일 때는 0에 가까워야 하는 것이 목표입니다.

fake_imgs를 label_fake와 대조해 discriminator에 입력하여 loss_fake를 구합니다.

real_imgs를 label_real과 대조해 discriminator에 입력하여 loss_real을 구합니다.

이 둘의 평균을 계산해 loss_D를 구합니다.

 

마지막으로 loss_D에 대해 backpropagation을 수행한 뒤 매개변수를 갱신합니다.

 

* * *

여기서 보라색 부분을 다시 사용하지 않고 새로 선언했는데,

이는 앞서 loss_G의 backpropagation에서 fake_imgs가 computational graph에 연결되어

다시 사용하는 것이 불가한 것으로 보입니다.

실제로 discriminator에서 보라색 부분을 지우고 실행하면 다음과 같은 문구가 출력됩니다.

 

* * *

 

3.7.4. Performance

        ## 3.7.4. performance ##
        performance_D = discriminator(real_imgs).mean()
        performance_G = discriminator(fake_imgs).mean()
        
        if (i + 1) % 150 == 0:
            print("Epoch [ {}/{} ]  Step [ {}/{} ]  d_loss : {:.5f}  g_loss : {:.5f}"
                  .format(epoch, max_epoch, i+1, len(dataloader), loss_D.item(), loss_G.item()))

학습 이후, 실제 / 위조 이미지에 대한 discriminator의 출력값의 평균을 계산해 performance를 구합니다.

실제값에 대해 1에 가까울 수록 discriminator의 performance가 높으며,

위조값에 대해 1에 가까울 수록 generator의 performance가 높을 것입니다.

 

추가로 한 epoch에서 150, 300, 450, 600번째 batch 마다

현재 epoch / 현재 batch / loss 값을 출력하도록 합니다.

3.8. Main code

import time
import os
import torch
import torch.optim as optim
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import MNIST
from torchvision.utils import save_image

# Device setting
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print("current device : {}".format(device))

# Ouput Directory setting
res_path = "results"
if not os.path.exists(res_path):
    os.makedirs(res_path)

# hyperparamter
max_epoch = 200
batch_size = 100
lr = 0.0002

img_size = 784
noise_size = 100
hidden_size1 = 256
hidden_size2 = 512

loss = nn.BCELoss()
    
# nn.Module : Base class for all neural network modules.
class Generator(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = nn.Linear(noise_size, hidden_size1)
        self.linear2 = nn.Linear(hidden_size1, hidden_size2)
        self.linear3 = nn.Linear(hidden_size2, img_size)
        self.relu = nn.ReLU()
        self.tanh = nn.Tanh()
    
    def forward(self, x):
        x = self.relu(self.linear1(x))
        x = self.relu(self.linear2(x))
        x = self.tanh(self.linear3(x))
        
        return x

class Discriminator(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = nn.Linear(img_size, hidden_size2)
        self.linear2 = nn.Linear(hidden_size2, hidden_size1)
        self.linear3 = nn.Linear(hidden_size1, 1)
        self.relu = nn.LeakyReLU(0.2)
        self.sigmoid = nn.Sigmoid()
    
    def forward(self, x):
        x = self.relu(self.linear1(x))
        x = self.relu(self.linear2(x))
        x = self.sigmoid(self.linear3(x))
        
        return x

# Dataset - MNIST
transform = transforms.Compose([transforms.ToTensor(),
                                transforms.Normalize(0.5, 0.5)])

dataset = MNIST(root='../data',
                train=True,
                transform=transform,
                download=True)

# Dataloader
dataloader = DataLoader(dataset=dataset,
                        batch_size=batch_size,
                        shuffle=True)

# Model
generator = Generator().to(device)
discriminator = Discriminator().to(device)

# Optimizer
optim_G = optim.Adam(generator.parameters(), lr=lr)
optim_D = optim.Adam(discriminator.parameters(), lr=lr)

start = time.time()

for epoch in range(max_epoch):
    print("epoch: {}".format(epoch+1))
    
    # dataset iteration
    for i, (imgs, label) in enumerate(dataloader):
## 3.7.1. Batch initialization ##
        label_real = torch.full((batch_size, 1), 1, dtype=torch.float32).to(device)
        label_fake = torch.full((batch_size, 1), 0, dtype=torch.float32).to(device)
        
        real_imgs = imgs.reshape(batch_size, -1).to(device)  # B * 784
    
       ## 3.7.2. Generator training ##
        optim_G.zero_grad()
        optim_D.zero_grad()
    
        z = torch.randn(batch_size, noise_size).to(device)
        fake_imgs = generator(z)
        
        loss_G = loss(discriminator(fake_imgs), label_real)
        
        loss_G.backward()
        optim_G.step()
        
       ### 3.7.3. Discriminator training ##
        optim_G.zero_grad()
        optim_D.zero_grad()

        z = torch.randn(batch_size, noise_size).to(device)
        fake_imgs = generator(z)
        
        loss_fake = loss(discriminator(fake_imgs), label_fake)
        loss_real = loss(discriminator(real_imgs), label_real)
        loss_D = (loss_fake + loss_real) / 2
        
        loss_D.backward()
        optim_D.step()

        ## 3.7.4. performance ##
        performance_D = discriminator(real_imgs).mean()
        performance_G = discriminator(fake_imgs).mean()
        
        if (i + 1) % 150 == 0:
            print("Epoch [ {}/{} ]  Step [ {}/{} ]  d_loss : {:.5f}  g_loss : {:.5f}"
                  .format(epoch, max_epoch, i+1, len(dataloader), loss_D.item(), loss_G.item()))

    # print performance
    print(" Epoch {}'s discriminator performance : {:.2f}  generator performance : {:.2f}"
          .format(epoch, performance_D, performance_G))

    # Save fake images in each epoch
    res_imgs = fake_imgs.reshape(batch_size, 1, 28, 28)
    save_image(res_imgs, os.path.join(res_path, 'GAN_result {}.png'.format(epoch + 1)))

elasped_time = time.time() - start
print("[Done] Time performance : {}".format(elasped_time))

4. 결과

MNIST 기반 GAN의 출력 이미지. 좌측 상단에서 부터 1, 40, 80, 120, 160, 200 epoch 때의 결과값.

40 epoch 단위로 확인 했을 때 학습시간에 따라 생성되는 이미지가 MNIST의 이미지와 유사해지는 것을 확인했습니다.

단, noise vector가 무작위로 생성되었기 때문에 생성된 손글씨 숫자들 또한 무작위로 생성됩니다.