1. Generative Adversarial Nets(GAN) - code review

 

 

Generative Adversarial Nets(GAN) 논문 리뷰: cumulu-s.tistory.com/22

1. Generative Adversarial Nets(GAN) - paper review

 

이전 글에서는, Generative Adversarial Nets(GAN) 논문에 대해서 상세하게 정리해보았습니다.

 

 

이번 글에서는, 직접 코드를 하나하나 짜 보면서 논문에 나온 내용들이 어떻게 구현될 수 있는지를 확인해보겠습니다.

 

 

저는 Deep learning library 중 pytorch를 메인으로 사용하고 있어서, pytorch로 작성하였음을 알립니다.

 

 

그럼 시작해보겠습니다.

 

 

# Import packages
import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.optim as optim
from tqdm.auto import tqdm
import numpy as np
import matplotlib.pyplot as plt

먼저 필요한 패키지들을 불러옵니다. torch와 torchvision, numpy 등을 불러옵니다.

 

 

tqdm의 경우, for문이 얼만큼 돌았는지 나타내 주는 패키지로 여러 번의 반복문을 돌아야 하는 코드에서 유용하게 사용되는 패키지입니다.

 

 

USE_CUDA = torch.cuda.is_available()
DEVICE = torch.device("cuda" if USE_CUDA else "cpu")
print("사용하는 Device : ", DEVICE)

EPOCHS = 400
BATCH_SIZE = 200

torch.cuda.is_avilable()은 cuda를 사용할 수 있는지를 True / False로 알려줍니다. 즉 GPU를 사용할 수 있는지 확인하는 코드입니다.

 

DEVICE는 gpu가 가능하면 cuda로, 아니면 cpu로 뜨게 만들어줍니다.

 

그리고 학습을 진행할 에폭과 batch size를 지정해줍니다. 

 

batch size의 경우 Dataloader에서 나오는 데이터의 수를 결정하므로 여기서 먼저 지정해줍니다.

 

# code for MNIST dataset loading error at Google Colab.
from six.moves import urllib    
opener = urllib.request.build_opener()
opener.addheaders = [('User-agent', 'Mozilla/5.0')]
urllib.request.install_opener(opener)

해당 코드는, 혹시나 Google Colab에서 해당 코드를 돌리시는 분들을 위해 작성된 코드입니다.

 

Google Colab을 사용하지 않으신다면 넘어가셔도 좋습니다.

 

현재 torchvision에서 MNIST dataset을 불러올 때, 데이터셋 링크로부터 불러오는 코드가 거절당하는 경우가 있어 이를 세팅해주는 코드입니다.

 

# Transformer code
transformer = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.5,), (0.5,))
])


# Loading trainset, testset and trainloader, testloader
trainset = torchvision.datasets.MNIST(root = '/content/drive/MyDrive/MNIST', train = True,
                                        download = True, transform = transformer)

trainloader = torch.utils.data.DataLoader(trainset, batch_size = BATCH_SIZE, shuffle = True, num_workers = 2)


testset = torchvision.datasets.MNIST(root = '/content/drive/MyDrive/MNIST', train = False,
                                        download = True, transform = transformer)

testloader = torch.utils.data.DataLoader(testset, batch_size = BATCH_SIZE, shuffle = True, num_workers = 2)

transformer는 데이터셋을 받아서 변형해주는 도구입니다.

 

classification과 같은 task를 할 때는 여기에 Augmentation을 넣어주기도 합니다. 

 

transforms.ToTensor()는 numpy array를 torch tensor로 바꿔주는 역할을 합니다. 

 

보통 데이터셋들은 numpy array로 되어 있고, 이를 pytorch에서 사용하려면 torch tensor의 형태로 변환해줘야 합니다.

 

transforms.Normalize는 정규화를 해주는 코드입니다.

 

정규화를 진행하지 않으면 MNIST 데이터셋이 0부터 1의 값을 가지도록 조정이 되어있습니다.

 

따라서, 이에 평균 0.5, 표준편차 0.5로 정규화를 진행해줍니다.

 

정규화는 $\frac{x-\mu}{\sigma}$의 형태로 진행됩니다.

 

따라서 0의 값은 -1이 되고, 1의 값은 1이 됩니다. 이 과정을 통해 MNIST 데이터셋이 -1부터 1의 값을 가지도록 바꿔줍니다.

 

torchvision.datasets.MNIST를 통해서 MNIST 데이터셋을 받을 수 있고, root는 데이터셋을 받을 경로입니다. 

 

해당 경로는 제가 사용하고 있는 경로이므로, 직접 사용하실 때는 이 부분을 원하시는 경로로 변경해주셔야 합니다.

 

다음으로는 torch.utils.data.DataLoader를 정의해줍니다. DataLoader는 데이터 전체를 보관하고 있는 일종의 container 같은 것으로, 이를 통해 데이터셋 전체를 메모리에 올리지 않아도 되게 만들어줍니다.

 

즉, batch size만큼만 DataLoader에서 데이터를 내보내주는 방식을 사용하는 것이죠.

 

# sample check
sample, label = next(iter(trainloader))

# show grid image
def imshow_grid(img):
    img = torchvision.utils.make_grid(img.cpu().detach())
    img = (img + 1) / 2
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    ax = plt.gca()
    ax.axes.xaxis.set_visible(False)
    ax.axes.yaxis.set_visible(False)
    plt.show()


imshow_grid(sample[0:8])

 

해당 코드는, DataLoader로부터 데이터가 잘 나오고 있는지를 확인하는 코드입니다.

 

iter()를 통해서 trainloader를 iterator로 만들어주고, next()를 이용해 한 번 데이터를 뽑도록 만들어줍니다.

 

현재 mnist trainloader는 이미지 데이터와 label을 동시에 뱉어내기 때문에, sample , label의 형태로 두 개의 변수를 지정해줘야 합니다.

 

imshow_grid 함수는 이미지를 받아서 이를 grid 형태로 보여주는 함수입니다.

 

torchvision.utils.make_grid()를 이용하면 grid 형태의 이미지를 만들 수 있습니다.

 

여기서 (img + 1) / 2를 적용해주는 이유는, 우리가 0 ~ 1로 되어 있는 데이터를 -1 ~ 1로 만들어줬기 때문입니다.

 

따라서 여기에 다시 1을 더해주고, 2로 나눠주어 원래의 0 ~ 1 형태로 원상복구 해주기 위함입니다.

 

그리고 현재는 데이터셋이 torch tensor로 되어 있으므로, 이를 numpy()를 이용하여 다시 numpy array로 만들어줘야 합니다.

 

그리고 np.transpose를 사용해주는 이유는, mnist 데이터셋은 (channel, width, height) 형태의 shape를 가지고 있기 때문입니다.

 

따라서, 실제로 이를 plt.imshow()를 통해서 시각화하려면 차원의 순서를 바꿔줘야 합니다. 그래서 사용하는 것이고요.

 

이 코드를 통해 아래에 보이는 그림은 8개의 샘플만 뽑아서 시각화를 진행한 결과입니다.

 

# Discriminator class
class Dis_model(nn.Module):
    def __init__(self, image_size, hidden_space):
        super(Dis_model, self).__init__()
        self.features = nn.Sequential(
            nn.Linear(image_size, hidden_space),
            nn.ReLU(),
            nn.Linear(hidden_space, hidden_space),
            nn.ReLU(),
            nn.Linear(hidden_space, 1),
            nn.Sigmoid())
    
    def forward(self, input_x):
        x = self.features(input_x)
        return x

다음으로는 Discriminator class를 만들어주었습니다.

 

보통은 모델을 class형태로 만들어주기 때문에, 비록 GAN 모델 구현은 그리 복잡하진 않아서 이렇게 하지 않고 더 간단하게도 짤 수 있겠지만 그냥 습관적으로 class로 만들어보았습니다.

 

image size를 지정해주면, 이를 입력으로 받아서 hidden space 만큼의 fully connected layer를 만들어주고, 다시 hidden space - hidden space 만큼의 fully connected layer를 만들어준 다음, 마지막으로는 output이 1로 나오도록 해줍니다.

 

이때, activation function은 Sigmoid function을 사용하여 마지막 layer의 결괏값이 0부터 1 사이로 나오게 해 줍니다.

 

이를 통해서, Discriminator model은 image_size라고 하는 크기의 데이터를 받아서 해당 데이터가 만들어진 데이터인지(위조지폐인지) 혹은 실제 데이터셋에 존재하는 데이터인지(실제 지폐인지)를 판단할 수 있게 됩니다.

 

# Generator class
class Gen_model(nn.Module):
    def __init__(self, latent_space, hidden_space, image_size):
        super(Gen_model, self).__init__()
        self.features = nn.Sequential(
            nn.Linear(latent_space, hidden_space),
            nn.ReLU(),
            nn.Linear(hidden_space, hidden_space),
            nn.ReLU(),
            nn.Linear(hidden_space, image_size),
            nn.Tanh())
        
    def forward(self, input_x):
        x = self.features(input_x)
        return x

이번에는 Generator model입니다.

 

Generator model은 latent_space라는 크기의 latent variable을 입력으로 받아서 hidden_space 크기로 output을 내고, 이를 다시 hidden_space 크기로 output을 내며, 마지막으로는 image_size 만큼의 결괏값을 내보내게 됩니다.

 

이를 통해서, Generator model은 latent_space라고 하는 크기의 noise input을 받아서 image_size라고 하는 크기의 데이터를 output으로 내게 됩니다.

 

im_size = 784
hidden_size = 256
latent_size = 100

Dis_net = Dis_model(image_size = im_size, hidden_space = hidden_size).to(DEVICE)
Gen_net = Gen_model(image_size = im_size, hidden_space = hidden_size, latent_space = latent_size).to(DEVICE)

d_optimizer = optim.Adam(Dis_net.parameters(), lr = 0.0002)
g_optimizer = optim.Adam(Gen_net.parameters(), lr = 0.0002)

이미지의 사이즈는 28 * 28 이므로 784로 지정하였고, hidden layer의 차원은 256, Generator에 들어가는 latent variable의 차원은 100차원으로 지정해주었습니다.

 

이미지의 사이즈는 데이터셋에 맞춰서 설정해주시면 되고 hidden layer의 차원이나 latent variable의 차원은 유동적으로 변경하셔도 됩니다.

 

하지만 이게 너무 작거나 하게 되면 Generator나 Discriminator가 충분한 capacity를 가질 수 없을 수도 있으니 과도하게 작게 하는 것은 여러모로 안 좋게 되지 않을까 하는 생각이 듭니다.

 

Discriminator 모델은 Dis_net이라는 이름으로 설정해주고, Generator 모델은 Gen_net이라는 이름으로 설정해주었습니다.

 

만약에 GPU를 사용하시는 분이라면, 반드시 .to(DEVICE)를 지정해주셔서 이를 GPU에서 연산하도록 설정해주셔야 합니다.

 

GPU로 하는 것과 CPU로 하는 것은 딥러닝에서의 연산 시간 차이가 심합니다. 

 

그리고 Discriminator 모델과 Generator 모델을 최적화해줄 optimization algorithm으로는 Adam을 선택했습니다.

 

모델이 두 개이므로, 두 개를 따로 만들어주셔야 합니다. 

 

학습률(learning rate)도 유동적으로 선택하시면 되는데, 저는 그냥 0.0002 정도로 잡았습니다.

 

# Start training
def train(generator, discriminator, train_loader, optimizer_d, optimizer_g):

    # Train version
    generator.train()
    discriminator.train()

    for data, target in train_loader:

        data, target = data.to(DEVICE), target.to(DEVICE)

        # ==========================================#
        # ==========Optimize discriminator==========#
        # ==========================================#

        # initialize discriminator optimizer
        optimizer_d.zero_grad()

        # Make noise samples for discriminator update
        noise_samples_d = torch.randn(BATCH_SIZE, latent_size).to(DEVICE)

        # real loss
        discri_value = discriminator(data.view(-1, 28*28))
        loss_real = -1 * torch.log(discri_value) # gradient ascent

        # fake loss
        gene_value = discriminator(generator(noise_samples_d))
        loss_fake = -1 * torch.log(1.0 - gene_value) # gradient ascent

        # Final loss
        loss_d = (loss_real + loss_fake).mean()

        loss_d.backward()
        optimizer_d.step()


        # ========================================= #
        # ==========Optimize generator============= #
        # ========================================= #

        # initialize generator optimizer
        optimizer_g.zero_grad()

        # Make noise samples for generator update
        noise_samples_g = torch.randn(BATCH_SIZE, latent_size).to(DEVICE)

        # calculate loss
        fake_value = discriminator(generator(noise_samples_g))
        loss_generator = -1 * torch.log(fake_value).mean() # provide much stronger gradients early in learning.
        
        loss_generator.backward()
        optimizer_g.step()

해당 함수는 모델의 학습을 수행합니다.

 

먼저, 모델을 학습하려면 반드시 generator.train()와 같이 이 모델은 학습을 할 것이다 라는 시그널을 보내줘야 합니다.

 

반대로 학습을 하지 않고 추론만 진행하려고 한다면, .eval()를 적용해주시면 됩니다.

 

학습은 train_loader에서 데이터를 받아서 진행하므로, data와 target을 train_loader에서 전달받습니다.

 

이때, data와 target도 모델과 동일하게 .to(DEVICE)를 통해서 GPU 연산을 적용하도록 만들어줘야 합니다.

 

 

혹시나 궁금하시다면, 해당 부분에 .to(DEVICE)를 빼시면 오류가 날 겁니다. (예전에 이런 적이 있었는데, 오류명까지는 기억이 안 나고 오류 내용에 cuda 설정 어쩌고저쩌고 하면서 설명이 잘 나옵니다.)

 

먼저 optimizer_d.zero_grad()로 gradient를 초기화해주고

 

torch.randn을 이용해서 정해진 사이즈의 noise variable을 만들어줍니다. 우리는 정규분포 기준으로 만들어줍니다.

 

 

우리가 BATCH_SIZE 기준으로 데이터를 받고 연산을 진행하기 때문에, noise variable도 크기는 BATCH_SIZE에 맞춰줍니다.

 

다음으로는, 논문에 이 파트를 구현해주는 코드입니다.

 

먼저, 위 식의 첫 번째 항을 만들어줍니다.

 

discri_value는 데이터를 먼저 (200, 784) 형태로 만들어줍니다. 이는 .view()를 이용하시면 됩니다.

 

 

(200, 784)로 만들어준 실제 MNIST 데이터를 discriminator의 input으로 넣어줍니다. 

 

이렇게 하면 discriminator가 해당 데이터를 기반으로 이게 위조지폐인지 실제 지폐인지에 대한 판단을 하겠죠?

 

그 값이 바로 discri_value입니다. 

 

여기에 log 값을 취해준 게 loss_real이 되고, 실제 논문에서는 gradient ascending을 하라고 했기 때문에 여기에 -1을 곱해줘서 구현을 진행했습니다.

 

다음으로는, 두 번째 항을 만들어줍니다.

 

gene_value는 아까 만들어낸 noise sample을 generator에 넣어서 새롭게 데이터를 만들어냅니다. 

 

이를 discriminator에 넣어서 해당 데이터가 위조지폐인지 실제 지폐인지 판단한 결과 값이 gene_value입니다.

 

여기에 log 값을 취해주고, gradient ascending을 위해 -1을 곱해서 이를 loss_fake라고 지정하였습니다.

 

마지막으로, loss_real과 loss_fake를 합해주고, 이를 평균 내는 것이 논문의 내용이였으므로 .mean()을 해줍니다.

 

 

loss_d.backward()는 backpropagation을 진행하는 코드이고, optimizer_d.step()은 optimizer를 한 단계 진행해줍니다.

 

다음으로는 Generator 학습을 진행해줘야 합니다.

 

논문상으로는 해당 파트입니다.

 

 

원래는 위의 식을 사용해야 하나, 논문에서는 초기의 학습을 위해서 다음과 같이 변경하도록 하고 있습니다.

 

 

즉 $log(1-D(G(z^{(i)})))$를 최소화하는 것이 아니라, $logD(G(z))$를 최대화하는 쪽으로 학습하라는 것입니다.

 

 

동일하게, generator의 optimizer gradient를 0으로 갱신해주고

 

다시 한번 noise sample을 만들어줍니다. 이 부분은 위 식에서 $z^{(i)}$로 표현됩니다.

 

 

이를 이용해서 generator의 input으로 투입하고, 이를 다시 discriminator에 넣어서 해당 데이터가 위조지폐인지 실제 지폐인지 discriminator가 판단합니다.

 

그리고 이 값에 log을 취한 뒤에 평균을 취해주고, 최대화를 해야 하므로 여기에 -1을 곱해서 진행합니다.

 

 

마지막으로 loss_generator를 backpropagation 해주고, optimizer_g도 한 단계 진행하도록 해줍니다.

 

 

 

다음으로는 실제 학습을 진행해주는 코드를 다뤄보겠습니다.

 

for epoch in tqdm(range(EPOCHS)):
    train(Gen_net, Dis_net, trainloader, d_optimizer, g_optimizer)

    if (epoch+1)%20 == 0:
        print('epoch %i / 400' % (epoch+1))
        
        noise_sam = torch.randn(16, latent_size).to(DEVICE)
        imshow_grid(Gen_net(noise_sam).view(-1, 1, 28, 28))
        print("\n")

 

위에서 언급한 train이라는 코드가 바로 학습을 진행하는 코드이고, 여기에 Gen_net과 Dis_net, trainloader, d_optimizer, g_optimizer를 인수로 넣어서 학습을 진행해줍니다. 

 

그리고 학습이 잘 되는지를 확인해보기 위해서, 20 에폭마다 16개의 noise variable을 만들어 이를 Generator network에 통과시켜서 나오게 되는 데이터를 시각화해줍니다.

 

epoch 20 ~ 60

에폭 80 ~ 120

에폭 140 ~ 180

에폭 200 ~ 240

에폭 260 ~ 300

에폭 320 ~ 360

에폭 380 ~ 400

 

400 에폭을 학습해감에 따라서, 만들어내는 숫자들이 깔끔해지는 것을 보실 수 있습니다.

 

특히, 가장 초기 20 에폭쯤에서는 완전히 random noise 느낌의 이미지를 만들어내지만 점점 갈수록 모양을 잡아나가는 것을 볼 수 있습니다.

 

물론 아직도 원본 데이터만큼 엄청나게 좋은 데이터를 만든다고 보기는 어렵지만, 굉장히 간단한 모델과 간단한 코드를 가지고도 어느 정도의 생성 성능을 보여줄 수 있음을 확인할 수 있었습니다.

 

해당 모델은 2014년 모델로, 그 이후에 다양한 생성 모델 방법론들이 나오면서 지금은 매우 고화질의 이미지들도 만들어내긴 하지만 이렇게 단순한 코드로도 어느정도 생성을 할 수 있다는 것이 놀랍다고 생각했습니다.

 

마지막으로, 학습이 완료된 모델을 가지고 test를 진행하는 코드를 다뤄보겠습니다.

 

# For test after training
vis_loader = torch.utils.data.DataLoader(testset, 16, True)
img_vis, label_vis = next(iter(vis_loader))
imshow_grid(img_vis)

아까 만들어준 testset을 가지고 DataLoader를 만들어주고, 16개만 뽑아내 줍니다. 

 

이를 통해 맨 처음에 정의한 imshow_grid 함수를 통해 16개만 샘플로 그리드 그림을 그려줍니다.

 

# Make samples by using trained generator model.
sample_noise = torch.randn(16, latent_size).to(DEVICE)
imshow_grid(Gen_net(sample_noise).view(-1, 1, 28, 28))

 

 

16개의 sample noise를 만들어주고, 이를 400 에폭 학습이 끝난 Generator network에 투입시켜서 이미지를 생성하고 이를 (-1, 1, 28, 28) 형태로 데이터를 다시 만들어 이를 plotting 해줍니다.

 

 

물론 그림의 성능이 그렇게 썩 좋지는 않지만 그래도 나름대로 완전한 랜덤 값에서 이미지를 만들어내고 있음을 확인할 수 있었습니다.

 

 

 

여기까지 GAN 코드를 직접 짜고 돌려보면서, 정말로 이미지 생성이 잘 되는지 확인해보았습니다.

 

위에서 설명한 코드는 제 Github 주소에 있으니, 직접 돌려보고 싶은 분들은 이를 확인해주시면 되겠습니다.

 

github.com/PeterKim1/paper_code_review

 

다음에는 또 다른 논문으로 찾아오겠습니다.

 

 

Reference

 

intelligence.korea.ac.kr/members/wschoi/seminar/tutorial/mnist/pytorch/gan/GAN-%ED%8A%9C%ED%86%A0%EB%A6%AC%EC%96%BC/