[CV] DCGAN Practice

- 본 글에서는 CV분야에서 활용되는 DCGAN을 구현합니다.

- 실습에는 Colab을 사용합니다.

 

 

 

1. Import package

 

모델 구성에 필요한 패키지를 가져오고, 설정을 세팅합니다.

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from torch import optim
import matplotlib.pyplot as plt
import numpy as np
import scipy as sp
from tqdm import tqdm
import os
import random
import time
import datetime
from sklearn.datasets import fetch_openml
from sklearn.model_selection import train_test_split

print(torch.cuda.is_available()) # True

np.set_printoptions(precision=3)
np.set_printoptions(suppress=True)
random.seed(1234)
np.random.seed(1234)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

 

 

 

2. Load Dataset

 

CIFAR-10 데이터셋을 가져와 DataLoader 객체를 구현합니다.

from PIL import Image
from torch.utils import data
import torchvision
from torchvision import transforms

def create_dataloader(batch_size=64, num_workers=0):

	# transforms은 이미지 데이터를 변환하는데 도움을 주는 라이브러리입니다.
    # Compose는 일련의 변환과정을 하나의 Sequence로 만들어줍니다.
    # 데이터가 입력되면 Tensor로 변환해주고, Nomalize합니다.
    transform = transforms.Compose([transforms.ToTensor(),
                                    transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))])
	
    # torchvision의 CIFAR10 데이터셋을 Load하여 전처리를 합니다.
    trainset = torchvision.datasets.CIFAR10(root='./data/', train=True, transform=transform, download=True)
    testset = torchvision.datasets.CIFAR10(root='./data/', train=False, transform=transform, download=True)
	
    # 받아온 데이터셋을 DataLoader로 만듭니다.
    trainloader = data.DataLoader(dataset=trainset, batch_size=batch_size, shuffle=True, num_workers=num_workers)
    testloader = data.DataLoader(dataset=testset, batch_size=batch_size, shuffle=False, num_workers=num_workers)

    return trainloader, testloader

 

 

 

3. Define Model

 

DCGAN은 Convolution layer를 활용한 Generator와 Discriminator를 사용합니다. 따라서 Convolution layer를 활용한 Generator와 Discriminator를 만들겠습니다. 두 신경망의 구성은 다음 그림과 같습니다.

Generator는 (1,1,256)의 데이터를 입력으로 받아 Deconv Layer를 통과하여 (4,4,256) -> (8,8,128) -> (16,16,64) -> (32,32,3)로 반환됩니다. 

Discriminator는 (32,32,3)을 입력으로 받아 Conv Layer를 통과하여 (16,16,64) -> (8,8,128) -> (4,4,256) -> (1,1,1)로 변환됩니다.

 

Transpose Convolution(Deconv)에 대한 함수는 다음 내용을 참고해 주세요 :(https://pytorch.org/docs/stable/generated/torch.nn.ConvTranspose2d.html)

 

3-1. Convolution Block

# Discriminator에서 반복적으로 사용 할 Convolution Block입니다.
# Convolution -> Normaliztion -> Activation 과정을 하나의 Block으로 구현합니다.

def conv(c_in, c_out, k_size, stride=2, pad=1, bias=False, norm='bn', activation=None):
    layers = []
    
    # Conv.
    layers.append(nn.Conv2d(c_in, c_out, k_size, stride, pad, bias=bias))

    # Normalization
    if norm == 'bn':
        layers.append(nn.BatchNorm2d(c_out))
    elif norm == None:
        pass

    # Activation
    if activation == 'lrelu':
        layers.append(nn.LeakyReLU(0.2))
    elif activation == 'relu':
        layers.append(nn.ReLU())
    elif activation == 'tanh':
        layers.append(nn.Tanh())
    elif activation == 'sigmoid':
        layers.append(nn.Sigmoid())
    elif activation == None:
        pass

    return nn.Sequential(*layers)

 

3.2 Deconvolution Block

# Generator에서 반복적으로 사용 할 Deconvolution Block입니다.
# Deconvolution -> Normaliztion -> Activation 과정을 하나의 Block으로 구현합니다.

def deconv(c_in, c_out, k_size, stride=2, pad=1, output_padding=0, bias=False, norm='bn', activation=None):
    layers = []

    # Deconv.
    layers.append(nn.ConvTranspose2d(c_in, c_out, k_size, stride, pad, output_padding, bias=bias))

    # Normalization
    if norm == 'bn':
        layers.append(nn.BatchNorm2d(c_out))
    elif norm == None:
        pass

    # Activation
    if activation == 'lrelu':
        layers.append(nn.LeakyReLU(0.2))
    elif activation == 'relu':
        layers.append(nn.ReLU())
    elif activation == 'tanh':
        layers.append(nn.Tanh())
    elif activation == 'sigmoid':
        layers.append(nn.Sigmoid())
    elif activation == None:
        pass

    return nn.Sequential(*layers)

 

3.3 Discriminator

위에서 구현한 Convolution Block을 활용하여 Discriminator를 구현하도록 하겠습니다.

class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        model = []
		
        # Discriminator는 4개의 Conv으로 구성되며
        # 각 층의 차원은 위의 그림과 같이 구성됩니다.
        model += [conv(3, 64, 4, norm='bn', activation='lrelu'),
                  conv(64, 128, 4, norm='bn', activation='lrelu'),
                  conv(128, 256, 4, norm='bn', activation='lrelu'),
                  conv(256, 1, 4, 1, 0, norm=None, activation='sigmoid')]
        self.model = nn.Sequential(*model)

    def forward(self, x: torch.Tensor):
        # 입력데이터는 (Batch, 3, 32, 32)의 형태이고
        # 출력데이터는 (Batch, 1)의 형태가 됩니다.
        output = self.model(x).squeeze()

        return output

 

3.4 Generator

위에서 구현한 Deconvolution Block을 활용하여 Generator를 구현하도록 하겠습니다.

class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()
        model = []
		
        # Generator는 4개의 Deconv로 구성되며
        # 각 층의 차원은 위의 그림과 같이 구성됩니다.
        model += [deconv(256, 256, 4, 1, 0, norm='bn', activation='relu'),
                  deconv(256, 128, 4, norm='bn', activation='relu'),
                  deconv(128, 64, 4, norm='bn', activation='relu'),
                  deconv(64, 3, 4, norm=None, activation='tanh')]
        self.model = nn.Sequential(*model)

    def forward(self, z):
        # 입력데이터는 (Batch, 256, 1, 1)의 형태이고
        # 출력데이터는 (Batch, 3,32,32)의 형태가 됩니다.
        z = z.view(z.size(0), z.size(1), 1, 1)
        output = self.model(z)

        return output

 

 

 

4. Evaluator

 

기존의 지도학습 모델에서는 Loss나 validation accuracy를 통해 학습 진행 정도를 확인할 수 있었습니다. 하지만 GAN에서는 이러한 방법으로 학습 진행 정도를 확인할 수 없습니다.

 

GAN의 학습정도를 확인할 수 있는 방법은 크게 2가지인데, 첫 번째로 Fidelity(충실도)이고 두 번째로 Diversity(다양성)입니다. Fidelity는 이미지의 품질 정도를 의미하고 Diversity는 생성된 이미지가 얼마나 다양한지를 의미합니다. 두 성능을 측정하는 데 사용되는 기법은 여러 가지 있으며 이번 연습에는 Diversity를 측정하는 데 사용되는 Inception Score라는 metric을 사용하겠습니다.

 

Inception score는 먼저 Generator에서 생성된 N개의 이미지를 pre-trained된 Inception network(=googleNet)에 통과시켜 가짜 이미지의 Label별 확률의 평균이 얼마나 다양한지 측정합니다. 

 

Inception Score에 대한 자세한 내용은 다음을 참고하시면 됩니다 :

(https://arxiv.org/abs/1801.01973)

(https://cyc1am3n.github.io/2020/03/01/is_fid.html)

# pre-train된 Inception Network를 가져옵니다.
from torchvision.models.inception import inception_v3
from scipy.stats import entropy

class Inception_Score():
    
    # 입력으로 주어진 Dataset은 Generator로 만든 가짜 이미지 데이터입니다.
    def __init__(self, dataset):
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

        self.N = len(dataset)
        self.batch_size = 64

		# 입력 데이터를 받아 해상도를 높입니다.
        self.dataset = dataset
        self.dataloader = data.DataLoader(dataset=dataset, batch_size=self.batch_size, num_workers=1)
        self.transform = nn.Upsample(size=(299, 299), mode='bilinear').to(self.device)
		
        # Pre-train된 Inception Network를 받아 eval모드로 변경합니다.
        self.inception_model = inception_v3(pretrained=True, transform_input=False).to(self.device)
        self.inception_model.eval()

    def get_pred(self, x):
    	
        # 주어진 가짜 이미지에 대해 해상도를 높이고
        # 가짜 이미지의 라벨을 예측하는 모델을 통과시켜 라벨을 만듭니다.
        with torch.no_grad():
            x = self.transform(x)
            x = self.inception_model(x)
            return F.softmax(x, dim=1).data.cpu().numpy()

    def compute_score(self, splits=1):
        preds = np.zeros((self.N, 1000))
		
        # 가짜 이미지의 예측 라벨들을 저장합니다.
        for i, batch in tqdm(enumerate(self.dataloader)):
            batch = batch.to(self.device)
            batch_size_i = batch.size(0)
            preds[i * self.batch_size : i * self.batch_size + batch_size_i] = self.get_pred(batch)
		
        # 다음은 Inception Score 연산을 하기위한 코드입니다.
        inception_score = 0.0
        split_scores = []
        for k in tqdm(range(splits)):
            part = preds[k * (self.N // splits): (k + 1) * (self.N // splits), :]
            py = np.mean(part, axis=0)
            scores = []
            for i in range(part.shape[0]):
                pyx = part[i, :]
                scores.append(entropy(pyx, py))
            split_scores.append(np.exp(np.mean(scores)))
        inception_score = np.mean(split_scores)

        return inception_score

 

 

 

5. Train

 

dataloader, model, evaluator를 만들었습니다. 이제 GAN을 학습시키는 Trainer를 구현해야 합니다.

Trainer는 Discriminator Loss와 Generator Loss를 동시에 줄이는 방식으로 동작합니다. Discriminator의 목적함수는 \(-E_{x-p_{data}}[logD(x)] + E_z(log(1-D(G(z)))]\)로 표현할 수 있고, Generator의 목적함수는 \(-E_z[log(D(G(z))]\)로 표현할 수 있습니다.

 

이 두 목적함수를 활용해 Trainer를 구현하겠습니다.

# 정규화된 이미지를 원래의 상태로 돌려놓는 함수입니다.
def denorm(x):
    out = (x + 1) / 2
    return out.clamp(0, 1)
    
# 체크포인트를 저장하는 함수입니다.
def save_checkpoint(model, save_path, device):
    if not os.path.exists(os.path.dirname(save_path)):
        os.makedirs(os.path.dirname(save_path))
    torch.save(model.cpu().state_dict(), save_path)
    model.to(device)

# 체크포인트를 로드하는 함수입니다.
def load_checkpoint(model, checkpoint_path, device):
    if not os.path.exists(checkpoint_path):
        print("Invalid path!")
        return
    model.load_state_dict(torch.load(checkpoint_path))
    model.to(device)

# 폴더에 속한 이미지를 가져오기 위한 함수입니다.
class FolderDataset(data.Dataset):
    def __init__(self, folder):
        self.folder = folder
        self.image_list = os.listdir(folder)
        self.transform = transforms.Compose([transforms.ToTensor(),
                                             transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))])

    def __getitem__(self, index):
        image = Image.open(os.path.join(self.folder, self.image_list[index]))
        return self.transform(image)

    def __len__(self):
        return len(self.image_list)

# 데이터 학습과 테스트를 진행하는 클래스입니다.
class Trainer():
    def __init__(self, 
                 trainloader, 
                 testloader, 
                 generator, 
                 discriminator, 
                 criterion,
                 g_optimizer, 
                 d_optimizer, 
                 device):

        self.trainloader = trainloader
        self.testloader = testloader
        self.G = generator
        self.D = discriminator
        self.criterion = criterion
        self.g_optimizer = g_optimizer
        self.d_optimizer = d_optimizer
        self.device = device

        # Make directory to save the images & models for a specific checkpoint
        os.makedirs(os.path.join('./results/', 'images'), exist_ok=True)
        os.makedirs(os.path.join('./results/', 'checkpoints'), exist_ok=True)
        os.makedirs(os.path.join('./results/', 'evaluation'), exist_ok=True)
    
    # 학습의 진행은 Discriminator를 학습 후 Generator를 학습합니다.
    # 특정 iteration마다 테스트를 통해 Inception Score를 확인합니다.
    def train(self, epochs = 1):
        self.G.to(self.device)
        self.D.to(self.device)

        start_time = time.time()
        for epoch in range(epochs):
            for iter, (real_img, _) in enumerate(self.trainloader):
                self.G.train()
                self.D.train()
                
                batch_size = real_img.size(0)
                # 실제 이미지에 대한 라벨은 모두 1입니다.
                real_label = torch.ones(batch_size).to(self.device)
                
                # 가짜 이미지에 대한 라벨은 모두 0입니다.
                fake_label = torch.zeros(batch_size).to(self.device)
                
                # 실제 이미지입니다.
                real_img = real_img.to(self.device)
                
                # 가짜 이미지 생성을 위한 z를 생성합니다.
                z = torch.randn(real_img.size(0), 256).to(self.device)

                D_loss: torch.Tensor = None
                
                real_out = self.D(real_img)
                fake_img = self.G(z)
                
                # detach함수를 통해 Discriminator Loss가 Generator network의 parameter에
                # 영향을 주지 않도록 합니다.
                fake_out = self.D(fake_img.detach())
                
                # 위에서 정의한 Discriminator의 Loss입니다.
                D_loss = self.criterion(real_out, real_label) + self.criterion(fake_out, fake_label)

                # TEST의 통과가 맞는 구현을 보장하지는 못합니다. 일반적으로는 loss가 1.38~1.45 사이의 값이 나와야 합니다.
                if epoch == 0 and iter == 0:
                    assert D_loss.detach().allclose(torch.tensor(1.4036), atol=2e-1), \
                    f"Discriminator Loss of the model does not match expected result."
                    print("==Discriminator loss function test passed!==")
				
                # 역전파를 실행하여 학습을 진행합니다.
                self.D.zero_grad()
                D_loss.backward()
                self.d_optimizer.step()

                G_loss: torch.Tensor = None
                fake_img = self.G(z)
                fake_out = self.D(fake_img)
                
                # 위에서 정의한 Generator의 Loss입니다.
                G_loss = self.criterion(fake_out, real_label)

                # TEST의 통과가 맞는 구현을 보장하지는 못합니다. 일반적으로는 loss가 1.35~1.52 사이의 값이 나와야 합니다.
                if epoch == 0 and iter == 0:
                    assert G_loss.detach().allclose(torch.tensor(1.5178), atol=2e-1), \
                    f"Generator Loss of the model does not match expected result."
                    print("==Generator loss function test passed!==")
				
                # 역전파를 실행하여 학습을 진행합니다.
                self.G.zero_grad()
                G_loss.backward()
                self.g_optimizer.step()
			
            # 학습 시간과 epoch를 보여줍니다.
            end_time = time.time() - start_time
            end_time = str(datetime.timedelta(seconds=end_time))[:-7]
            print('Time [%s], Epoch [%d/%d], lossD: %.4f, lossG: %.4f'
                  % (end_time, epoch+1, epochs, D_loss.item(), G_loss.item()))

            # 매 epoch마다 만들어진 이미지를 저장합니다.
            fake_img = fake_img.reshape(fake_img.size(0), 3, 32, 32)
            torchvision.utils.save_image(denorm(fake_img), os.path.join('./results/', 'images', 'fake_image-{:03d}.png'.format(epoch+1)))
            
            # 특정 epoch마다 test를 실행합니다.
            if epoch % 10 == 0:
                self.test()

        save_checkpoint(self.G, os.path.join('./results', 'checkpoints', 'G_final.pth'), self.device)
        save_checkpoint(self.D, os.path.join('./results', 'checkpoints', 'D_final.pth'), self.device)
	
    # 테스트는 가짜 이미지를 생성해 가짜 이미지의 라벨을 예측하여 얻은
    # 라벨의 분포를 통해 Inception Score를 반환합니다.
    def test(self):
        print('Start computing Inception Score')
        self.G.eval()
        with torch.no_grad():
            for iter in tqdm.tqdm(range(5000)):
                z = torch.randn(1, 256).to(self.device)
                fake_img = self.G(z)
                torchvision.utils.save_image(denorm(fake_img), os.path.join('./results/', 'evaluation', 'fake_image-{:03d}.png'.format(iter)))

        # Compute the Inception score
        dataset = FolderDataset(folder = os.path.join('./results/', 'evaluation'))
        Inception = Inception_Score(dataset)
        score = Inception.compute_score(splits=1)
        print('Inception Score : ', score)

 

6. Visualization

다음은 학습 epoch에 따른 Generator가 만든 이미지입니다.

Epoch = 1

Epoch = 10

Epoch = 30

epoch = 50