[Machine Learning] PyTorch로 DCGAN 구현하기 (Deep Convolutional Generative Adversarial Network)

지난 글에서 아무 기술이 들어가지 않은 Vanilla GAN을 살펴보았다.

GAN에 대한 설명은 다 지난 글에서 다뤄봤으니 바로 본론으로 넘어가자.

 

오늘은 간단한 기술을 하나 추가한 GAN을 살펴보자.

참고로 추가할 기술은 이미지 처리에 적합한 convolution이다.

---

DCGAN (Deep Convolutional GAN)

본 블로그에서 CNN에 대해 다뤘던 적이 있다.

CNN은 이미지의 구조적인 정보를 해치지 않기 때문에 이미지 처리에 적합한 네트워크라고 설명했었다.

 

이번에는 CNN을 GAN에 적용해서 이미지 판별 및 이미지 생성에 적합한 GAN을 제작해보자.

즉 Discriminator와 Generator을 다중 퍼셉트론이 아닌 CNN을 사용한다는 이야기이다.

 

이미 우리는 CNN와 GAN을 모두 알고 있으니 바로 코드로 구현해보자.

---

CNN을 적용한 Generator와 Discriminator

DCGAN의 제네레이터와 디스크리미네이터를 설정하는데 가장 어려운 점은 이미지 크기에 맞게 네트워크를 만드는 것이다.

이때 torchinfo을 사용하면 입출력의 텐서 사이즈를 미리 살펴볼 수 있다.

torchinfo의 사용법은 다음 글을 참조하자.

[PyTorch] torchinfo, summary 사용 방법 및 파라미터 / 여러 input 설정

[PyTorch] torchinfo, summary 사용 방법 및 파라미터 / 여러 input 설정

먼저 새롭게 바뀐 네트워크를 먼저 살펴보자.

아래 코드는 네트워크를 만들고 입출력을 테스트해보는 코드이다.

 

먼저 필요한 라이브러리를 임포트 하자.

import torch
import torch.nn as nn
from torchinfo import summary

먼저 제네레이터를 만들어보자.

제네레이터의 랜덤 벡터는 채널이 100, 가로 세로가 각각 1, 1으로 총 100개의 컴포넌트가 있는 랜덤 벡터이다.

class Generator(nn.Module):
    def __init__(self, nc=1, nz=100, ngf=64):
        super(Generator, self).__init__()
        
        self.main = nn.Sequential(
            # input is Z, going into a convolution
            nn.ConvTranspose2d(     nz, ngf * 8, 4, 1, 0, bias=False),
            nn.BatchNorm2d(ngf * 8),
            nn.ReLU(True),
            # state size. (ngf*8) x 4 x 4
            nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 4),
            nn.ReLU(True),
            # state size. (ngf*4) x 8 x 8
            nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 2),
            nn.ReLU(True),
            # state size. (ngf*2) x 16 x 16
            nn.ConvTranspose2d(ngf * 2, ngf, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf),
            nn.ReLU(True),
            
            nn.ConvTranspose2d(ngf, nc, kernel_size=1, stride=1, padding=2, bias=False),
            nn.Tanh()
        )

    def forward(self, input):
        output = self.main(input)
        return output

G = Generator()
print(summary(G, (64, 100, 1, 1)))

위 코드를 실행하면 제네레이터에 (64, 100, 1, 1)의 임의의 텐서를 넣었을 때의 출력을 보여준다.(64는 임의의 배치 사이즈이다. 다른 숫자로 둬도 무방하다.)

MNIST 이미지의 크기가 28*28이므로 우리도 그 크기에 맞게 텐서가 출력되도록 네트워크를 만들어 넣어야 한다.

출력은 아래와 같다.

==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
Generator                                [64, 1, 28, 28]           --
├─Sequential: 1-1                        [64, 1, 28, 28]           --
│    └─ConvTranspose2d: 2-1              [64, 512, 4, 4]           819,200
│    └─BatchNorm2d: 2-2                  [64, 512, 4, 4]           1,024
│    └─ReLU: 2-3                         [64, 512, 4, 4]           --
│    └─ConvTranspose2d: 2-4              [64, 256, 8, 8]           2,097,152
│    └─BatchNorm2d: 2-5                  [64, 256, 8, 8]           512
│    └─ReLU: 2-6                         [64, 256, 8, 8]           --
│    └─ConvTranspose2d: 2-7              [64, 128, 16, 16]         524,288
│    └─BatchNorm2d: 2-8                  [64, 128, 16, 16]         256
│    └─ReLU: 2-9                         [64, 128, 16, 16]         --
│    └─ConvTranspose2d: 2-10             [64, 64, 32, 32]          131,072
│    └─BatchNorm2d: 2-11                 [64, 64, 32, 32]          128
│    └─ReLU: 2-12                        [64, 64, 32, 32]          --
│    └─ConvTranspose2d: 2-13             [64, 1, 28, 28]           64
│    └─Tanh: 2-14                        [64, 1, 28, 28]           --
==========================================================================================
Total params: 3,573,696
Trainable params: 3,573,696
Non-trainable params: 0
Total mult-adds (G): 26.61
==========================================================================================
Input size (MB): 0.03
Forward/backward pass size (MB): 126.23
Params size (MB): 14.29
Estimated Total Size (MB): 140.55
==========================================================================================

총 3,573,696개의 학습 가능한 파라미터를 가진 제네레이터가 완성되었다.

 

이제 비슷한 방법으로 디스크리미네이터도 만들어보자.

class Discriminator(nn.Module):
    def __init__(self, nc=1, ndf=64):
        super(Discriminator, self).__init__()
        
        self.main = nn.Sequential(
            # input is (nc) x 64 x 64
            nn.Conv2d(nc, ndf, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf) x 32 x 32
            nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 2),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf*2) x 16 x 16
            nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 4),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf*4) x 8 x 8
            nn.Conv2d(ndf * 4, 1, 4, 2, 1, bias=False),
            nn.Sigmoid()
        )

    def forward(self, input):
        output = self.main(input)
        return output.view(-1, 1).squeeze(1)
    
D = Discriminator()
print(summary(D, (64, 1, 28, 28)))

채널이 1, 가로 세로가 각각 28인 이미지를 입력받아 0에서 1사이의 출력(진짜는 1, 가짜는 0)을 나타내는 하나의 숫자만 나오도록 설계했다.

위 코드를 실행시켜 네트워크의 구조를 프린트해보면 아래와 같다.

==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
Discriminator                            [64]                      --
├─Sequential: 1-1                        [64, 1, 1, 1]             --
│    └─Conv2d: 2-1                       [64, 64, 14, 14]          1,024
│    └─LeakyReLU: 2-2                    [64, 64, 14, 14]          --
│    └─Conv2d: 2-3                       [64, 128, 7, 7]           131,072
│    └─BatchNorm2d: 2-4                  [64, 128, 7, 7]           256
│    └─LeakyReLU: 2-5                    [64, 128, 7, 7]           --
│    └─Conv2d: 2-6                       [64, 256, 3, 3]           524,288
│    └─BatchNorm2d: 2-7                  [64, 256, 3, 3]           512
│    └─LeakyReLU: 2-8                    [64, 256, 3, 3]           --
│    └─Conv2d: 2-9                       [64, 1, 1, 1]             4,096
│    └─Sigmoid: 2-10                     [64, 1, 1, 1]             --
==========================================================================================
Total params: 661,248
Trainable params: 661,248
Non-trainable params: 0
Total mult-adds (M): 726.19
==========================================================================================
Input size (MB): 0.20
Forward/backward pass size (MB): 15.20
Params size (MB): 2.64
Estimated Total Size (MB): 18.05
==========================================================================================

---

PyTorch로 만든 DCGAN

이제 새롭게 바뀐 네트워크를 만들었으니 코드를 직접 짜보자.

Vanilla GAN와 다른 점은 제네레이터의 입력인 랜덤 벡터의 크기가 약간 바뀌었다는 것 외에는 똑같다.

import os
import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
from torchvision.utils import save_image

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

# Directory to save the output of the generator
save_path = "./generated_images"
if not os.path.exists(save_path):
    os.makedirs(save_path)

# Set random seed for reproducibility
torch.manual_seed(0)

# Generator Network
class Generator(nn.Module):
    def __init__(self, nc=1, nz=100, ngf=64):
        super(Generator, self).__init__()
        
        self.main = nn.Sequential(
            # input is Z, going into a convolution
            nn.ConvTranspose2d(nz, ngf * 8, 4, 1, 0, bias=False),
            nn.BatchNorm2d(ngf * 8),
            nn.ReLU(True),
            # state size. (ngf*8) x 4 x 4
            nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 4),
            nn.ReLU(True),
            # state size. (ngf*4) x 8 x 8
            nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 2),
            nn.ReLU(True),
            # state size. (ngf*2) x 16 x 16
            nn.ConvTranspose2d(ngf * 2, ngf, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf),
            nn.ReLU(True),
            
            nn.ConvTranspose2d(ngf, nc, kernel_size=1, stride=1, padding=2, bias=False),
            nn.Tanh()
        )

    def forward(self, input):
        output = self.main(input)
        return output

# Discriminator Network
class Discriminator(nn.Module):
    def __init__(self, nc=1, ndf=64):
        super(Discriminator, self).__init__()
        
        self.main = nn.Sequential(
            # input is (nc) x 64 x 64
            nn.Conv2d(nc, ndf, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf) x 32 x 32
            nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 2),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf*2) x 16 x 16
            nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 4),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf*4) x 8 x 8
            nn.Conv2d(ndf * 4, 1, 4, 2, 1, bias=False),
            nn.Sigmoid()
        )

    def forward(self, input):
        output = self.main(input)
        return output.view(-1, 1).squeeze(1)

# Create discriminator and generator objects
D = Discriminator().to(device)
G = Generator().to(device)
# Define loss function and optimizer for discriminator
criterion = nn.BCELoss()
optimizer_D = torch.optim.Adam(D.parameters(), lr=0.0002, betas=(0.5, 0.999))

# Define loss function and optimizer for generator
optimizer_G = torch.optim.Adam(G.parameters(), lr=0.0002, betas=(0.5, 0.999))

# Load MNIST dataset
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
trainset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True)

# Fixed noise settings to monitor the training process
fixed_noise = torch.randn(100, 100, 1, 1).to(device)

# Training loop
num_epochs = 200
for epoch in range(num_epochs):
    for i, data in enumerate(trainloader):
        # Train discriminator with real data
        D.zero_grad()
        real_images, _ = data
        batch_size = real_images.size(0)
        label = torch.full((batch_size,), 1.).to(device)
        output = D(real_images.to(device))
        errD_real = criterion(output, label)
        errD_real.backward()

        # Train discriminator with fake data
        noise = torch.randn(batch_size, 100, 1, 1).to(device)
        fake_images = G(noise).to(device)
        label.fill_(0.)
        output = D(fake_images.detach())
        errD_fake = criterion(output, label)
        errD_fake.backward()
        errD = errD_real + errD_fake
        optimizer_D.step()

        # Train generator
        G.zero_grad()
        label.fill_(1.)
        output = D(fake_images)
        errG = criterion(output, label)
        errG.backward()
        optimizer_G.step()

        # Print statistics
        if i % 100 == 0:
            print('[%d/%d][%d/%d]\tLoss_D: %.4f\tLoss_G: %.4f'
                  % (epoch + 1, num_epochs, i, len(trainloader), errD.item(), errG.item()))

    # Save generated images
    with torch.no_grad():
        save_image(G(fixed_noise), f"generated_images/{epoch+1}.png", nrow=10, normalize=True)

---

결과 확인

마지막 200 에포크에서 제네레이터의 결과를 보자.

200-epoch에서 Generator의 출력

확실히 Vanilla GAN보다 깨끗하고 선명한 것을 확인할 수 있다.

---

결과를 보며 문득 이런 생각이 들 수 있다.

GAN을 MNIST 데이터를 사용해 0부터 9까지 여러 클래스를 학습시켰는데, 제네레이터에서 원하는 클래스의 데이터를 선택해서 생성할 수 없을까라는 생각이 들 수 있다.

다음 포스팅에서는 원하는 데이터만 골라 생성하는 GAN을 살펴보도록 하자.