026. GAN(Generative Adversarial Network) 으로 간다

https://youtu.be/sucqskXRkss

위 링크는 ICLR에서 Ian Goddfellow가 GAN 및 2019년 기준 최신 연구를 발표한 영상입니다.

ICLR은 매년 전 세계에서 열리는 머신 러닝 컨퍼런스이고 Ian Goodfellow는 GAN의 아이디어를 세상에 알린 사람입니다.

워낙 이해하기 어려워서 유튜브를 서칭하며 공부하다가 보게 된 영상인데

GAN의 창시자가 직접 이야기해주니 감회가 남다른 영상입니다.

 

마지막 수업 시간에 GAN에 대한 수업을 들었습니다.

사실 GAN이 원리 때문에 내용도 어렵고 이미지를 학습하다보니 리소스도 많이 필요합니다.

그래서 수업 시간에 이를 이해하고 어떻게 구현하는 건지 알고 넘어가기 힘들었습니다.

그래도 Youtube와 책을 통해 이해는 하고 넘어가야겠지요..?

 

GAN(Generative Andversarial Network)를 이해하려면 일단 이름부터 살펴봐야 합니다.

첫 글자인 Generative는 이 신경망의 목적이 생성이라고 말해줍니다.

아마도 인풋 이미지의 패턴을 학습해 또렷한 이미지를 생성하는 게 목적인듯합니다.

왜냐하면 저번 글에서 이야기했던 Autoencoder는 Convolutional AE, Variational AE 등은 결국

원본과 비슷한 이미지를 만드는 목적으로 계속 발전해왔습니다..

그러나 어디까지나 비슷한 이미지에 불과했고 사람이 보기엔 충분히 구분이 가능했습니다.

 

이 한계를 타파할 방법으로 Adversarial(적대적) 학습을 제시한 것이 Ian Goodfellow입니다.

기존의 오토 인코더는 인코더를 거쳐 생성된 이미지와 인풋 이미지를 비교해 성능을 측정하고

이 성능을 높이는 방향으로 학습을 진행했습니다.

그런데 GAN은 이미지를 생성하는 Generator와 진짜 여부를 가리는 Discriminator가 교차로 

성능을 높이는 방향으로 학습을 진행합니다.

Generator와 Discriminator가 적대적인 관계로 제로섬 게임 틀 안에서 경쟁하며 신경망이 구현됩니다.

 

Generator는 랜덤한 분포의 데이터를 입력으로 받아 실제 이미지와 같은 데이터를 출력합니다.

이 랜덤한 분포의 데이터는 오토 인코더에서 인코더를 거친 잠재 표현(Latent)과 비슷해 보입니다.

언뜻 보면 Generator는 오토 인코더의 디코더와 비슷해 보입니다.

물론 이 둘은 같은 기능을 제공하지만 학습 방식이 매우 다릅니다.

 

Discriminator는 Generator가 생성한 이미지와 Trainingset에서 추출한 이미지를 입력으로 받아

가짜인지 진짜인지 구분하고 분류 성능을 높이는 학습을 진행합니다.

 

Generator와 Discriminator의 목표는 정반대입니다.

둘 중 하나의 성능 지표가 좋아지면 나머지 하나의 성능 지표가 떨어집니다.

그래서 GAN의 훈련은 두 단계로 이뤄집니다.

 

첫 번째로 Discriminator를 학습시킵니다.

Training Set에서 실제 이미지, Generator에서 가짜 이미지를 같은 수로 가져와서 합칩니다.

그리고 이 데이터를 학습시키는데, 이 단계에서 역전파는 Discriminator에만 적용해 최적화합니다.

 

두 번째로 Generator를 학습시킵니다.

먼저 Generator로 가짜 이미지를 생성하고 이 이미지를 Discriminator로 진짜 가짜를 판별합니다.

이 과정에선 가짜 이미지를 진짜로 labeling 해 Discriminator를 속여야 합니다.

만약 Discriminator가 가짜라 판정한다면 Generator의 성능에 학습할 여지가 있음을 알려주는 것입니다.

이 과정에서는 역시 Discriminator의 파라미터를 고정하고 Generator의 파라미터만 최적화합니다.

 

이 학습 과정에서는 Discriminator의 성능이 좋다면 역전파 과정에서

Discriminator에게 진짜 이미지를 더 많이 제공받게 됩니다.

서로 적대적인 목적을 가지고 경쟁하지만 학습 과정에서 서로의 학습을 돕습니다.

 

https://github.com/eriklindernoren/PyTorch-GAN/blob/master/implementations/gan/gan.py

GitHub - eriklindernoren/PyTorch-GAN: PyTorch implementations of Generative Adversarial Networks.

위 링크에 있는 코드를 참고해 구현해보겠습니다. 

위 링크의 깃허브는 GAN에 관련하여 정보가 많아서

공부할 겸 한 번 보는 것도 나쁘지 않습니다.

 

구현에 사용할 Pytorch와 dataset을 불러올 torchvision, 시각화할 matplotlib 그리고 os, numpy, time을 불러옵니다.

import torch
import torch.nn as nn
import torch.nn.functional as F

from torchvision import datasets
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from torchvision.transforms.functional import to_pil_image
import matplotlib.pylab as plt
%matplotlib inline

import os
import numpy as np
import time

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

데이터의 경로를 지정해줍니다.

# 데이터 경로 지정
path2data = './data'
os.makedirs(path2data, exist_ok=True) # 폴더 생성

MNIST 손글씨 데이터를 불러옵니다. 0.5, 0.5로 Normalize도 해줍니다.

# MNIST dataset 불러오기
train_ds = datasets.MNIST(path2data,
                          train=True,
                          transform=transforms.Compose([
                                                        transforms.ToTensor(),
                                                        transforms.Normalize([0.5],[0.5])]),
                          download=True)

데이터 로더를 생성해줍니다. batch_size는 32로 지정하고 shuffle은 True로 지정합니다.

# 데이터 로더 생성
train_dl = DataLoader(train_ds, batch_size=32, shuffle=True)

# check
for x, y in train_dl:
    print(x.shape, y.shape)
    break

torch.Size([32, 1, 28, 28]) torch.Size([32])

 

여러 개의 FC Layer를 통과하는 Generator를 생성합니다.

그리고 모델에 입력할 random으로 noise를 16개 생성해 시험으로 이미지를 만듭니다.

16개로 출력됐는지 확인합니다.

# generator: noise를 입력받아 이미지를 생성합니다.
class Generator(nn.Module):
    def __init__(self, params):
        super().__init__()
        self.nz = params['nz'] # 입력 노이즈 벡터 수, 100
        self.img_size = params['img_size'] # 이미지 크기, 1x28x28

        self.model = nn.Sequential(
            *self._fc_layer(self.nz, 128, normalize=False),
            *self._fc_layer(128,256),
            *self._fc_layer(256,512),
            *self._fc_layer(512,1024),
            nn.Linear(1024,int(np.prod(self.img_size))),
            nn.Tanh()
        )

    def forward(self, z):
        img = self.model(z)
        img = img.view(img.size(0), *self.img_size)
        return img

    # fc layer
    def _fc_layer(self, in_channels, out_channels, normalize=True):
        layers = []
        layers.append(nn.Linear(in_channels, out_channels)) # fc layer
        if normalize:
            layers.append(nn.BatchNorm1d(out_channels, 0.8)) # BN
        layers.append(nn.LeakyReLU(0.2)) # LeakyReLU
        return layers

# check
params = {'nz':100,
          'img_size':(1,28,28)}
x = torch.randn(16,100).to(device) # random noise
model_gen = Generator(params).to(device)
output = model_gen(x) # noise를 입력받아 이미지 생성
print(output.shape)

torch.Size([16, 1, 28, 28])

이미지를 인풋으로 받아 진짜 이미지와 가짜 이미지를 분류하는 신경망을 만듭니다. 

그리고 이미지를 입력해 분류하고 데이터의 사이즈를 프린트합니다.

# discriminator: 진짜 이미지와 가짜 이미지를 분류합니다.
class Discriminator(nn.Module):
    def __init__(self,params):
        super().__init__()
        self.img_size = params['img_size'] # 이미지 크기, 1x28x28

        self.model = nn.Sequential(
            nn.Linear(int(np.prod(self.img_size)), 512),
            nn.LeakyReLU(0.2),
            nn.Linear(512,256),
            nn.LeakyReLU(0.2),
            nn.Linear(256,1),
            nn.Sigmoid()
        )

    def forward(self, x):
        x = x.view(x.size(0),-1)
        x = self.model(x)
        return x

# check
x = torch.randn(16,1,28,28).to(device)
model_dis = Discriminator(params).to(device)
output = model_dis(x)
print(output.shape)

torch.Size([16, 1])

 

Generator와 Discriminator가 각각 학습할 때 하나의 파라미터는 고정해야 되기 때문에

가중치를 초기화하는 함수를 생성합니다.

그리고 만들어 놓은 신경망에 가중치 초기화 함수를 적용합니다.

# 가중치 초기화
def initialize_weights(model):
    classname = model.__class__.__name__
    # fc layer
    if classname.find('Linear') != -1:
        nn.init.normal_(model.weight.data, 0.0, 0.02)
        nn.init.constant_(model.bias.data, 0)
    # batchnorm
    elif classname.find('BatchNorm') != -1:
        nn.init.normal_(model.weight.data, 1.0, 0.02)
        nn.init.constant_(model.bias.data, 0)

# 가중치 초기화 적용
model_gen.apply(initialize_weights);
model_dis.apply(initialize_weights);

손실 함수를 적용하고 최적화 파라미터를 지정해 각각의 optimizer를 만듭니다.

# 손실 함수
loss_func = nn.BCELoss()

from torch import optim

# 최적화 파라미터
lr = 2e-4
beta1 = 0.5

opt_dis = optim.Adam(model_dis.parameters(),lr=lr,betas=(beta1,0.999))
opt_gen = optim.Adam(model_gen.parameters(),lr=lr,betas=(beta1,0.999))

진짜 이미지와 가짜 이미지의 label을 지정하고 노이즈와 epoch를 지정합니다.

그리고 loss_history dic을 만듭니다. 

real_label = 1.
fake_label = 0.
nz = params['nz']
num_epochs = 100

loss_history={'gen':[],
              'dis':[]}

Generator와 Discriminator를 각각 절차에 맡게 훈련시키는 코드를 작성합니다.

시각화를 위해 loss_history에 loss값들을 입력합니다.

batch_count = 0
start_time = time.time()
model_dis.train()
model_gen.train()

for epoch in range(num_epochs):
    for xb, yb in train_dl:
        ba_si = xb.size(0)

        xb = xb.to(device)
        yb_real = torch.Tensor(ba_si,1).fill_(1.0).to(device)
        yb_fake = torch.Tensor(ba_si,1).fill_(0.0).to(device)

        # Generator
        model_gen.zero_grad()
        noise = torch.randn(ba_si,nz, device=device) # 노이즈 생성
        out_gen = model_gen(noise) # 가짜 이미지 생성
        out_dis = model_dis(out_gen) # 가짜 이미지 판별

        loss_gen = loss_func(out_dis, yb_real)
        loss_gen.backward()
        opt_gen.step()

        # Discriminator
        model_dis.zero_grad()

        out_real = model_dis(xb) # 진짜 이미지 판별
        out_fake = model_dis(out_gen.detach()) # 가짜 이미지 판별
        loss_real = loss_func(out_real, yb_real)
        loss_fake = loss_func(out_fake, yb_fake)
        loss_dis = (loss_real + loss_fake) / 2

        loss_dis.backward()
        opt_dis.step()

        loss_history['gen'].append(loss_gen.item())
        loss_history['dis'].append(loss_dis.item())

        batch_count += 1
        if batch_count % 1000 == 0:
            print('Epoch: %.0f, G_Loss: %.6f, D_Loss: %.6f, time: %.2f min' %(epoch, loss_gen.item(), loss_dis.item(), (time.time()-start_time)/60))

학습을 하며 둘 다 점점 loss가 줄어드는 것을 볼 수 있습니다.

# plot loss history
plt.figure(figsize=(10,5))
plt.title('Loss Progress')
plt.plot(loss_history['gen'], label='Gen. Loss')
plt.plot(loss_history['dis'], label='Dis. Loss')
plt.xlabel('batch count')
plt.ylabel('Loss')
plt.legend()
plt.show()

 

# 가중치 저장
path2models = './models/'
os.makedirs(path2models, exist_ok=True)
path2weights_gen = os.path.join(path2models, 'weights_gen.pt')
path2weights_dis = os.path.join(path2models, 'weights_dis.pt')

torch.save(model_gen.state_dict(), path2weights_gen)
torch.save(model_dis.state_dict(), path2weights_dis)

가중치들을 불러와 노이즈로 이미지를 생성합니다.

# 가중치 불러오기
weights = torch.load(path2weights_gen)
model_gen.load_state_dict(weights)

# evaluation mode
model_gen.eval()

# fake image 생성
with torch.no_grad():
    fixed_noise = torch.randn(16, 100, device=device)
    img_fake = model_gen(fixed_noise).detach().cpu()
print(img_fake.shape)

구현한 코드로 만든 Generator가 꽤 괜찮은 퀄리티의 이미지를 생성합니다.

# 가짜 이미지 시각화
plt.figure(figsize=(10,10))
for ii in range(16):
    plt.subplot(4,4,ii+1)
    plt.imshow(to_pil_image(0.5*img_fake[ii]+0.5),cmap='gray')
    plt.axis('off')