Pytorch – DCGAN を Pytorch で実装して、有名人の顔画像を生成する

import random

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from IPython.display import display
from tqdm.notebook import tqdm

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
from torchvision import datasets, transforms

# 再現性確保のため、シードを固定する。
torch.manual_seed(0)
random.seed(0)

# モデルの入力サイズ
img_size = 64

# バッチサイズ
batch_size = 256

# ノイズの次元数
latent_dim = 100

# エポック数
n_epochs = 10

# 学習率
lr = 0.0002

# beta1 (Adam のパラメータ)
beta1 = 0.5

def get_device(gpu_id=-1):
    if gpu_id >= 0 and torch.cuda.is_available():
        return torch.device("cuda", gpu_id)
    else:
        return torch.device("cpu")


device = get_device(gpu_id=0)

# Transform を作成する。
transform = transforms.Compose(
    [
        transforms.Resize(img_size),
        transforms.CenterCrop(img_size),
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
    ]
)

# Dataset を作成する。
download_dir = "/data"  # ダウンロード先は適宜変更してください
dataset = datasets.CelebA(
    download_dir, transform=transform, download=True
)

# DataLoader を作成する。
dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)

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)

class Generator(nn.Module):
    def __init__(self, input_dim, n_features=64):
        super().__init__()
        self.main = nn.Sequential(
            # fmt: off
            # conv1
            nn.ConvTranspose2d(input_dim, n_features * 8,
                               kernel_size=4, stride=1, padding=0, bias=False),
            nn.BatchNorm2d(n_features * 8),
            nn.ReLU(inplace=True),
            # conv2
            nn.ConvTranspose2d(n_features * 8, n_features * 4,
                               kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(n_features * 4),
            nn.ReLU(inplace=True),
            # conv3
            nn.ConvTranspose2d(n_features * 4, n_features * 2,
                               kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(n_features * 2),
            nn.ReLU(inplace=True),
            # conv4
            nn.ConvTranspose2d(n_features * 2, n_features,
                               kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(n_features),
            nn.ReLU(inplace=True),
            # conv5
            nn.ConvTranspose2d(n_features, 3,
                               kernel_size=4, stride=2, padding=1, bias=False),
            nn.Tanh()
            # fmt: on
        )

    def forward(self, x):
        return self.main(x)

# Generator を作成する。
G = Generator(latent_dim).to(device)
# 初期化する。
G.apply(weights_init)

class Discriminator(nn.Module):
    def __init__(self, n_features=64):
        super().__init__()
        self.main = nn.Sequential(
            # fmt: off
            # conv1
            nn.Conv2d(3, n_features,
                      kernel_size=4, stride=2, padding=1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            # conv2
            nn.Conv2d(n_features, n_features * 2,
                      kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(n_features * 2),
            nn.LeakyReLU(0.2, inplace=True),
            # conv3
            nn.Conv2d(n_features * 2, n_features * 4,
                      kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(n_features * 4),
            nn.LeakyReLU(0.2, inplace=True),
            # conv4
            nn.Conv2d(n_features * 4, n_features * 8,
                      kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(n_features * 8),
            nn.LeakyReLU(0.2, inplace=True),
            # conv5
            nn.Conv2d(n_features * 8, 1,
                      kernel_size=4, stride=2, padding=0, bias=False),
            nn.Sigmoid(),
            nn.Flatten(),
            # fmt: on
        )

    def forward(self, x):
        return self.main(x)

# Discriminator を作成する。
D = Discriminator().to(device)
# 初期化する。
D.apply(weights_init)

# 損失関数を作成する。
criterion = nn.BCELoss()

# 学習過程の Generator が生成する画像を可視化するためのノイズ z
fixed_z = torch.randn(64, latent_dim, 1, 1, device=device)

# ラベル
real_label = 1
fake_label = 0

# オプティマイザを作成する。
G_optimizer = optim.Adam(G.parameters(), lr=lr, betas=(beta1, 0.999))
D_optimizer = optim.Adam(D.parameters(), lr=lr, betas=(beta1, 0.999))

def D_train(x_real):
    D.zero_grad()

    # 損失関数を計算する。
    # 本物のデータが入力の場合の Discriminator の損失関数を計算する。
    y_pred = D(x_real)
    y_real = torch.full_like(y_pred, real_label)
    loss_real = criterion(y_pred, y_real)

    # 偽物のデータが入力の場合の Discriminator の損失関数を計算する。
    z = torch.randn(batch_size, latent_dim, 1, 1, device=device)
    y_pred = D(G(z))
    y_fake = torch.full_like(y_pred, fake_label)
    loss_fake = criterion(y_pred, y_fake)

    loss = loss_real + loss_fake

    # 逆伝搬する。
    loss.backward()

    # パラメータを更新する。
    D_optimizer.step()

    return float(loss)

def G_train(x):
    G.zero_grad()

    # 損失関数を計算する。
    z = torch.randn(batch_size, latent_dim, 1, 1, device=device)
    y_pred = D(G(z))
    y = torch.full_like(y_pred, real_label)
    loss = criterion(y_pred, y)

    # 逆伝搬する。
    loss.backward()

    # パラメータを更新する。
    G_optimizer.step()

    return float(loss)

def generate_img(G, fixed_z):
    with torch.no_grad():
        # 画像を生成する。
        x = G(fixed_z)

    # 画像を格子状に並べる。
    img = torchvision.utils.make_grid(x.cpu(), nrow=8, normalize=True, pad_value=1)
    # テンソルを PIL Image に変換する。
    img = transforms.functional.to_pil_image(img)

    return img

def train_gan(n_epochs):
    G.train()
    D.train()

    history = []
    for i_epoch in range(1, n_epochs + 1):

        D_losses, G_losses = [], []
        for x, _ in tqdm(dataloader, desc=f"epoch {i_epoch}/{n_epochs}"):
            x = x.to(device)
            D_losses.append(D_train(x))
            G_losses.append(G_train(x))

        # 途中経過を確認するために画像を生成する。
        img = generate_img(G, fixed_z)

        # 途中経過を記録する。
        info = {
            "epoch": i_epoch,
            "D_loss": np.mean(D_losses),
            "G_loss": np.mean(G_losses),
            "img": img,
        }
        history.append(info)

        # 途中経過を出力する。
        display(img)

    history = pd.DataFrame(history)

    return history


history = train_gan(n_epochs)

def plot_history(history):
    fig, ax = plt.subplots()

    # 損失の推移を描画する。
    ax.set_title("Loss")
    ax.plot(history["epoch"], history["D_loss"], label="Discriminator")
    ax.plot(history["epoch"], history["G_loss"], label="Generator")
    ax.set_xlabel("Epoch")
    ax.legend()

    plt.show()


plot_history(history)

def create_animation(imgs):
    """gif アニメーションにして保存する。
    """
    imgs[0].save(
        "history.gif", save_all=True, append_images=imgs[1:], duration=500, loop=0
    )


# 各エポックの画像で gif アニメーションを作成する。
create_animation(history["img"])

# 一番最後の画像を表示する。
display(history["img"].iloc[-1])