Deep Convolutional GAN

이번 페이지에서는 CNN 을 이용해 이미지 생성 모델의 시작을 알린 DCGAN (Deep Convolutional GAN) 을 이용해 생성형 모델을 학습하는 방법에 대해서 설명합니다.

Paper

논문에서는 안정적인 DCGAN 을 학습하기 위한 가이드라인을 아래와 같이 제시했습니다.

Architecture guidelines for stable Deep Convolutional GANs

• Replace anypooling layers with strided convolutions (discriminator) and fractional-strided convolutions (generator).
• Use batchnorm in both the generator and the discriminator.
• Remove fully connected hidden layers for deeper architectures.
• Use ReLU activation in generator for all layers except for the output, which uses Tanh.
• Use LeakyReLU activation in the discriminator for all layers

기본적인 학습방법은 Generative Adversarial Networks 논문을 따릅니다. 이를 위한 파이썬 수도 코드를 작성하면 아래와 같습니다.

for epoch in range(n_epochs):
    for batch in loader:
        #
        # 1. get discriminator loss from real data
        #
        ...
        #
        # 2. get discriminator loss from fake data
        #
        ...
        #
        # 3. get discriminator loss and update discriminator
        #
        ...
        #
        # 4. get generator loss and update generator
        #
        ...

이제부터 해당하는 내용들을 채워 보겠습니다.

Dataset

우선 튜토리얼에 들어가기에 앞서 사용할 데이터셋을 선언합니다. 데이터셋에 대한 자세한 설명은 CelebA 페이지에서 확인할 수 있습니다.

All models are trained with mini-batch stochastic gradient descent (SGD) with a mini-batch size of 128.

논문에서는 배치 사이즈를 128 을 사용해 학습했기 때문에 batch_size=128 로 주겠습니다.

import torchvision.transforms as T
from torchvision.datasets.celeba import CelebA
from torch.utils.data import DataLoader
from torchvision.utils import make_grid

transform = T.Compose(
    [
        T.Resize(64),
        T.CenterCrop(64),
        T.ToTensor(),
        T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
    ]
)
dataset = CelebA(
    "./datasets", download=True, transform=transform
)
loader = DataLoader(dataset, batch_size=128, shuffle=True)

Generator

Code

작성할 모델의 아키텍쳐는 논문에서 제시한 형태를 따라갑니다.

import torch
import torch.nn as nn


class Generator(nn.Module):
    def __init__(self, num_channel=3, latent_dim=100, feature_dim=64):
        super().__init__()

        self.layer_1 = nn.Sequential(
            nn.ConvTranspose2d(latent_dim, feature_dim * 8, 4, 1, 0, bias=False),
            nn.BatchNorm2d(feature_dim * 8),
            nn.ReLU(True),
        )
        self.layer_2 = nn.Sequential(
            nn.ConvTranspose2d(feature_dim * 8, feature_dim * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(feature_dim * 4),
            nn.ReLU(True),
        )
        self.layer_3 = nn.Sequential(
            nn.ConvTranspose2d(feature_dim * 4, feature_dim * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(feature_dim * 2),
            nn.ReLU(True),
        )
        self.layer_4 = nn.Sequential(
            nn.ConvTranspose2d(feature_dim * 2, feature_dim, 4, 2, 1, bias=False),
            nn.BatchNorm2d(feature_dim),
            nn.ReLU(True),
        )
        self.last_layer = nn.Sequential(
            nn.ConvTranspose2d(feature_dim, num_channel, 4, 2, 1, bias=False),
            nn.Tanh(),
        )

    def forward(self, z):
        # decoding
        layer_1_out = self.layer_1(z)  # (N, 512, 4, 4)
        layer_2_out = self.layer_2(layer_1_out)  # (N, 256, 8, 8)
        layer_3_out = self.layer_3(layer_2_out)  # (N, 128, 16, 16)
        layer_4_out = self.layer_4(layer_3_out)  # (N, 64, 32, 32)
        # transform to rgb
        out = self.last_layer(layer_4_out)  # (N, 3, 64, 64)
        return out

안정적인 학습을 위해 위에서 제시된 항목들을 지켜서 아키텍쳐를 구성했습니다.

Architecture guidelines for stable Deep Convolutional GANs

• Replace any pooling layers with fractional-strided convolutions (generator).
• Use batchnorm in the generator.
• Remove fully connected hidden layers for deeper architectures.
• Use ReLU activation in generator for all layers except for the output, which uses Tanh.

Execution

이번에는 생성기의 내부 코드가 어떻게 동작하는 지 확인해 보도록 하겠습니다.

latent_dim = 100
generator = Generator(latent_dim=latent_dim)

z = torch.rand(128, latent_dim, 1, 1) # (N, 100, 1, 1)

with torch.no_grad(): # decoding
    layer_1_out = generator.layer_1(z) # (N, 512, 4, 4)
    layer_2_out = generator.layer_2(layer_1_out) # (N, 256, 8, 8)
    layer_3_out = generator.layer_3(layer_2_out) # (N, 128, 16, 16)
    layer_4_out = generator.layer_4(layer_3_out) # (N, 64, 32, 32) # transform to rgb
    out = generator.last_layer(layer_4_out) # (N, 3, 64, 64)

• 입력 z
  • 생성기의 입력값은 생성할 이미지들의 정보가 있는 latent space 에서 추출한 값입니다.
  • 이 텐서의 차원은 (N, 100, 1, 1) 입니다.
• 실행을 한 후 각 중간 텐서들의 shape 을 확인하면 다음과 같이 됩니다.
  1. (N, 512, 4, 4)
  2. (N, 256, 8, 8)
  3. (N, 128, 16, 16)
  4. (N, 64, 32, 32)
• 출럭 out
  • 최종 출력으로 나오는 데이터는 이미지입니다.
  • 그렇기 때문에 텐서의 차원은 (N, 3, 64, 64) 입니다.
  • 각각 배치수, 채널수, 넓이, 높이 입니다.

Visualize

학습되지 않은 생성기가 생성하는 이미지를 확인하면 아래처럼 의미없는 노이즈들이 출력됩니다.

import numpy as np
import matplotlib.pyplot as plt
from torchvision.utils import make_grid


out_x_grid = make_grid(out, nrow=12).numpy()

plt.figure(figsize=(8, 8))
plt.title("First batch")

# make data range to 0~1
out_x_grid = (out_x_grid * 0.5) + 0.5
plt.imshow(np.transpose(out_x_grid, (1, 2, 0)))

Discriminator

이제 생성된 데이터 가짜 데이터와 진짜 데이터를 분류할 모델을 작성해 보겠습니다.

Code

class Discriminator(nn.Module):
    def __init__(self, num_channel=3, feature_dim=64):
        super().__init__()
        self.layer_1 = nn.Sequential(
            nn.Conv2d(num_channel, feature_dim, 4, 2, 1, bias=False),
            nn.BatchNorm2d(feature_dim),
            nn.LeakyReLU(0.2, inplace=True),
        )
        self.layer_2 = nn.Sequential(
            nn.Conv2d(feature_dim, feature_dim * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(feature_dim * 2),
            nn.LeakyReLU(0.2, inplace=True),
        )
        self.layer_3 = nn.Sequential(
            nn.Conv2d(feature_dim * 2, feature_dim * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(feature_dim * 4),
            nn.LeakyReLU(0.2, inplace=True),
        )
        self.layer_4 = nn.Sequential(
            nn.Conv2d(feature_dim * 4, feature_dim * 8, 4, 2, 1, bias=False),
            nn.BatchNorm2d(feature_dim * 8),
            nn.LeakyReLU(0.2, inplace=True),
        )

        self.last_layer = nn.Sequential(
            nn.Conv2d(feature_dim * 8, 1, 4, 1, 0, bias=False), nn.Sigmoid()
        )

    def forward(self, x):
        # encoding
        layer_1_out = self.layer_1(x)  # (N, 64, 32, 32)
        layer_2_out = self.layer_2(layer_1_out)  # (N, 128, 16, 16)
        layer_3_out = self.layer_3(layer_2_out)  # (N, 256, 8, 8)
        layer_4_out = self.layer_4(layer_3_out)  # (N, 512, 4, 4)
        # classify
        out = self.last_layer(layer_4_out).squeeze()  # (N)
        return out

안정적인 학습을 위해 위에서 제시된 항목들을 지켜서 아키텍쳐를 구성했습니다.

Architecture guidelines for stable Deep Convolutional GANs

• Replace any pooling layers with strided convolutions (discriminator).
• Use batchnorm in the discriminator.
• Remove fully connected hidden layers for deeper architectures.
• Use LeakyReLU activation in the discriminator for all layers

Execution

이번에는 분류기의 내부 코드가 어떻게 동작하는 지 확인해 보도록 하겠습니다.

discriminator = Discriminator()

for batch in loader:
    x = batch[0]
    break

with torch.no_grad(): # encoding
    layer_1_out = discriminator.layer_1(x) # (N, 64, 32, 32)
    layer_2_out = discriminator.layer_2(layer_1_out) # (N, 128, 16, 16)
    layer_3_out = discriminator.layer_3(layer_2_out) # (N, 256, 8, 8)
    layer_4_out = discriminator.layer_4(layer_3_out) # (N, 512, 4, 4) # classify
    out = discriminator.last_layer(layer_4_out).squeeze() # (N)

• 입력 x
  • 분류기의 입력값은 생성된 혹은 진짜 이미지 입니다.
  • 그렇기 때문에 입력 텐서의 차원은 (N, 3, 64, 64) 입니다.
• 실행을 한 후 각 중간 텐서들의 shape 을 확인하면 다음과 같이 됩니다.
  1. (N, 64, 32, 32)
  2. (N, 128, 16, 16)
  3. (N, 256, 8, 8)
  4. (N, 512, 4, 4)
• 출럭 out
  • 최종 출력으로 나오는 데이터는 각 이미지들의 진짜 혹은 가짜를 구분하기 위한 확률값입니다.
  • 이 때, discriminator.last_layer 레이어를 통해서 나온 결과는 (N, 1, 1, 1) 의 shape 을 갖습니다.
  • squeeze() : 좀 더 활용하기 쉽게 squeeze 옵션을 통해서 1로 이루어진 모든 차원을 제거해 (N, 1) shape 으로 변경합니다.

분류된 결과물을 간단하게 확인해보면 0~1 사이의 숫자로 이루어진 확률값입니다.

out[:10]

tensor([0.4386, 0.5597, 0.5724, 0.4969, 0.4873, 0.4456, 0.5770, 0.5647, 0.4118,
        0.4347])

Train

이제 모델을 학습하고 학습된 모델을 이용해 이미지를 생성하는 결과를 확인해 보겠습니다.

학습 알고리즘은 논문에서 제시한 방법 중 일부를 따릅니다.

제시된 내용은 아래와 같습니다.

1. All models are trained with mini-batch stochastic gradient descent (SGD) with a mini-batch size of 128.
2. All weights were initialized from a zero-centered Normal distribution with standard deviation 0.02
3. In the LeakyReLU, the slope of the leak was set to 0.2 in all models.
4. While previous GAN work has used momentum to accelerate training, we used the Adam optimizer (Kingma & Ba, 2014) with tuned hyperparameters. We found the suggested learning rate of 0.001, to be too high, using 0.0002 instead.
5. We found leaving the momentum term β1 at the suggested value of 0.9 resulted in training oscillation and instability while reducing it to 0.5 helped stabilize training.

이 중 에서 사용할 방법은 아래와 같습니다.

1. 학습 배치 사이즈는 128
2. 모델 파라미터들을 분산이 0.02 인 정규 분포로 초기화
3. LeakyReLU 의 slope 는 0.2
4. Adam 학습기를 사용하고 leraning_rate 는 0.0002 를 사용
5. β1 은 0.5로 설정

Weight Initialization

파라미터 초기화를 위한 함수를 작성합니다.

def weights_init(m):
    classname = m.__class__.__name__
    if classname.find("Conv") != -1:
        nn.init.normal_(m.weight.data, 0.0, 0.02)
    elif classname.find("BatchNorm") != -1:
        nn.init.normal_(m.weight.data, 1.0, 0.02)
        nn.init.constant_(m.bias.data, 0)

각 모델에 적용합니다.

_ = generator.apply(weights_init)
_ = discriminator.apply(weights_init)

Convolution Network 는 원활한 학습을 위해서는 gpu 가 필요합니다. GPU 가 없는 경우 학습에 다소 시간이 소요될 수 있습니다.

아래 코드를 이용해 device 를 선언합니다. 만약 gpu 가 사용 가능한 경우 device(type='cuda') 메세지가 나옵니다.

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

device(type='cuda')

위에서 선언한 모델을 gpu 메모리로 옮기겠습니다.

_ = discriminator.to(device)
_ = generator.to(device)

Logger

학습 과정을 tensorboard 에 저장하기 위한 writer 입니다.

from torch.utils.tensorboard import SummaryWriter
# tensorboard logger
writer = SummaryWriter()

Loss

GAN 모델의 학습은 분류기의 분류 결과로 부터 시작됩니다. 이 때 진짜로 구분은 1 가짜로 구분은 0으로 하는 이진 분류로서 BinaryCrossEntropy 를 이용해 계산합니다.

# loss function
bce_loss_fn = nn.BCELoss()

Optimizer

논문에서는 SGD 를 Optimizer 로 제시했지만 Adam 이 좀 더 안정적인 학습을 할 수 있기에 Adam 으로 선언합니다.

import torch.optim as optim


# optimizer
discriminator_opt = optim.Adam(
    discriminator.parameters(), lr=0.0002, betas=(0.5, 0.999)
)
generator_opt = optim.Adam(generator.parameters(), lr=0.0002, betas=(0.5, 0.999))

아래 코드를 작성 후 학습을 진행합니다.

Code

from tqdm import tqdm


# meta data
n_epochs = 5

# for visualization
iteration = 0
fixed_z = torch.randn((64, latent_dim, 1, 1)).to(device)

# for iteration history
G_loss_history = []
D_loss_history = []

# train code
for epoch in range(n_epochs):
    epoch_G_loss = 0.0
    epoch_D_loss = 0.0

    for batch in tqdm(loader, desc=f"Epoch {epoch + 1}/{n_epochs}"):
        generator.train()
        discriminator.train()

        x = batch[0].to(device)
        batch_size = x.size(0)

        #
        # 1. get discriminator loss from real data
        #
        discriminator_opt.zero_grad()
        real_D_score = discriminator(x)
        real_target = torch.ones_like(real_D_score).to(device)
        real_D_loss = bce_loss_fn(real_D_score, real_target)
        #
        # 2. get discriminator loss from fake data
        #
        z = torch.randn((batch_size, latent_dim, 1, 1)).to(device)
        fake_data = generator(z)

        fake_D_score = discriminator(fake_data)
        fake_target = torch.zeros_like(fake_D_score).to(device)
        fake_D_loss = bce_loss_fn(fake_D_score, fake_target)
        #
        # 3. get discriminator loss and update discriminator
        #
        D_loss = real_D_loss + fake_D_loss
        D_loss.backward()
        discriminator_opt.step()

        #
        # 4. get generator loss and update generator
        #
        generator_opt.zero_grad()
        z = torch.randn((batch_size, latent_dim, 1, 1)).to(device)
        generated_data = generator(z)

        generate_D_score = discriminator(generated_data)
        generate_target = torch.ones_like(generate_D_score).to(device)

        G_loss = bce_loss_fn(generate_D_score, generate_target)

        G_loss.backward()
        generator_opt.step()

        # iteration logging
        G_loss_history += [G_loss.item()]
        D_loss_history += [D_loss.item()]

        epoch_D_loss += D_loss.item() * len(x) / len(loader.dataset)
        epoch_G_loss += G_loss.item() * len(x) / len(loader.dataset)

        writer.add_scalar("Loss/generator", G_loss, iteration)
        writer.add_scalar("Loss/discriminator", D_loss, iteration)

        # iteration checkpoint
        if iteration % 1000 == 0:
            # visualization
            with torch.no_grad():
                generated_x = generator(fixed_z)
                generated_x = generated_x.detach().cpu()
                generated_x_grid = make_grid(generated_x, nrow=8).numpy()
                generated_x_grid = (generated_x_grid * 0.5) + 0.5
                generated_x_grid = np.transpose(generated_x_grid, (1, 2, 0))
                plt.imsave(f"dcgan-step-{iteration}.png", generated_x_grid)
            # model save
            torch.save(generator, f"dcgan-generator-step-{iteration}.pt")
        iteration += 1
    log_string = f"Loss at epoch {epoch + 1} - D_loss : {epoch_D_loss:.4f}, G_loss : {epoch_G_loss:.4f}"
    print(log_string)

# save last model
with torch.no_grad():
    generated_x = generator(fixed_z)
    generated_x = generated_x.detach().cpu()
    generated_x_grid = make_grid(generated_x, nrow=8).numpy()
    generated_x_grid = (generated_x_grid * 0.5) + 0.5
    generated_x_grid = np.transpose(generated_x_grid, (1, 2, 0))
    plt.imsave(f"dcgan-step-{iteration}.png", generated_x_grid)
torch.save(generator, f"dcgan-generator-step-{iteration}.pt")

Epoch 1/5: 100%|██████████| 1272/1272 [03:20<00:00,  6.34it/s]
Loss at epoch 1 - D_loss : 0.6342, G_loss : 7.3766
Epoch 2/5: 100%|██████████| 1272/1272 [03:18<00:00,  6.41it/s]
Loss at epoch 2 - D_loss : 0.6417, G_loss : 3.8137
Epoch 3/5: 100%|██████████| 1272/1272 [03:18<00:00,  6.42it/s]
Loss at epoch 3 - D_loss : 0.6967, G_loss : 2.9740
Epoch 4/5: 100%|██████████| 1272/1272 [03:17<00:00,  6.44it/s]
Loss at epoch 4 - D_loss : 0.7547, G_loss : 2.5252
Epoch 5/5: 100%|██████████| 1272/1272 [03:18<00:00,  6.41it/s]
Loss at epoch 5 - D_loss : 0.7395, G_loss : 2.4700

제일 처음에는 데이터 로더에서 실제 이미지를 분류기를 통해서 1로 분류하는 loss 를 계산합니다.

# train code
for epoch in range(n_epochs):
    ...
    for batch in tqdm(loader, desc=f"Epoch {epoch + 1}/{n_epochs}"):
        ...
        x = batch[0].to(device)
        #
        # 1. get discriminator loss from real data
        #
        discriminator_opt.zero_grad()
        real_D_score = discriminator(x)
        real_target = torch.ones_like(real_D_score).to(device)
        real_D_loss = bce_loss_fn(real_D_score, real_target)
        ...

이어서 임의의 latent z 를 선언하고 생성기를 통해서 가짜 이미지를 만듭니다. 이후 생성된 이미지를 분류기가 0으로 분류하는 loss 를 계산합니다.

# train code
for epoch in range(n_epochs):
    ...
    for batch in tqdm(loader, desc=f"Epoch {epoch + 1}/{n_epochs}"):
        ...
        #
        # 2. get discriminator loss from fake data
        #
        z = torch.randn((batch_size, latent_dim, 1, 1)).to(device)
        fake_data = generator(z)

        fake_D_score = discriminator(fake_data)
        fake_target = torch.zeros_like(fake_D_score).to(device)
        fake_D_loss = bce_loss_fn(fake_D_score, fake_target)
        ...

두 loss 를 합쳐서 분류기가 생성된 이미지와 실제 이미지를 잘 구분할 수 있도록 학습합니다.

# train code
for epoch in range(n_epochs):
    ...
    for batch in tqdm(loader, desc=f"Epoch {epoch + 1}/{n_epochs}"):
        ...
        #
        # 3. get discriminator loss and update discriminator
        #
        D_loss = real_D_loss + fake_D_loss
        D_loss.backward()
        discriminator_opt.step()
        ...

앞선 과정에서 분류기에 대한 부분을 학습했다면 이제 생성기가 더 잘 생성할 수 있도록 학습해야 합니다. 이를 위해서 생성된 이미지를 분류기로 분류하고 이에 대한 loss 를 앞선 분류기에서 생성된 이미지에 대해서 0으로 준 것과 반대로 1로 설정합니다. 이를 통해서 분류기를 속이기 위한 gradient 값을 계산해서 생성기를 업데이트합니다.

# train code
for epoch in range(n_epochs):
    ...
    for batch in tqdm(loader, desc=f"Epoch {epoch + 1}/{n_epochs}"):
        ...
        #
        # 4. get generator loss and update generator
        #
        generator_opt.zero_grad()
        z = torch.randn((batch_size, latent_dim, 1, 1)).to(device)
        generated_data = generator(z)

        generate_D_score = discriminator(generated_data)
        generate_target = torch.ones_like(generate_D_score).to(device)

        G_loss = bce_loss_fn(generate_D_score, generate_target)

        G_loss.backward()
        generator_opt.step()
        ...

위에서 학습의 결과물을 logging 하기 위한 코드입니다.

# train code
for epoch in range(n_epochs):
    ...
    for batch in tqdm(loader, desc=f"Epoch {epoch + 1}/{n_epochs}"):
        ...
        # iteration logging
        G_loss_history += [G_loss.item()]
        D_loss_history += [D_loss.item()]

        epoch_D_loss += D_loss.item() * len(x) / len(loader.dataset)
        epoch_G_loss += G_loss.item() * len(x) / len(loader.dataset)

        writer.add_scalar("Loss/generator", G_loss, iteration)
        writer.add_scalar("Loss/discriminator", D_loss, iteration)
        ...

아래 코드는 고정된 z 에 대해서 모델이 학습하면서 생성하는 그림이 어떻게 바뀌는지 저장하기 위한 코드입니다.

# for visualization
iteration = 0
fixed_z = torch.randn((64, latent_dim, 1, 1)).to(device)

# train code
for epoch in range(n_epochs):
    ...
    for batch in tqdm(loader, desc=f"Epoch {epoch + 1}/{n_epochs}"):
        ...
        # iteration checkpoint
        if iteration % 1000 == 0:
            # visualization
            with torch.no_grad():
                generated_x = generator(fixed_z)
                generated_x = generated_x.detach().cpu()
                generated_x_grid = make_grid(generated_x, nrow=8).numpy()
                generated_x_grid = (generated_x_grid * 0.5) + 0.5
                generated_x_grid = np.transpose(generated_x_grid, (1, 2, 0))
                plt.imsave(f"dcgan-step-{iteration}.png", generated_x_grid)
        ...

Loss history

Tensorboard 에 로깅된 학습 정보를 확인하면 다음과 같습니다.

Visualize

학습을 하면서 생성기가 어떻게 이미지를 생성하는 지 확인해 보겠습니다.

import os

img_list = list(sorted(filter(lambda x: x.startswith("dcgan-step"), os.listdir("."))))
nrows = (len(img_list) // 3) + 1

fig, axes = plt.subplots(ncols=3, nrows=nrows, figsize=(5 * nrows, 15))
for idx, fname in enumerate(img_list):
    array = plt.imread(fname)
    axes[idx // 3, idx % 3].imshow(array)
    axes[idx // 3, idx % 3].axis("off")
    axes[idx // 3, idx % 3].set_title(fname.replace(".png", ""))

plt.tight_layout()

Generation

이번에는 학습한 모델의 디코더를 이용해 랜덤한 latent 값을 주었을 때 어떤 이미지를 생성하는 지 확인해 보겠습니다. 다만 이번 생성기의 경우 BatchNorm 을 사용하기 때문에 학습과 평가 모드에 따라서 동작이 달라집니다. 학습이 끝난 모델을 사용하기 위해서는 eval() 모드를 설정해주어야 합니다.

generator.eval()
with torch.no_grad():    
    random_z = torch.randn((64, latent_dim, 1, 1)).to(device)
    generated_x = generator(random_z)

generated_x_grid = make_grid(generated_x, nrow=8).cpu().numpy()
# make data range to 0~1
generated_x_grid = (generated_x_grid * 0.5) + 0.5
plt.figure(figsize=(8, 8))
plt.title("Generated batch")
plt.imshow(np.transpose(generated_x_grid, (1, 2, 0)))

References

• https://pytorch.org/tutorials/beginner/dcgan_faces_tutorial.html#implementation
• https://arxiv.org/pdf/1511.06434