５ステップでできるPyTorch - DCGAN

# 概要

PyTorchを使って、以下の５ステップでDCGANを作成します。

1. データの準備  
2. Generatorの作成
3. Discriminatorの作成
4. 訓練関数の作成
5. DCGANの訓練スタート


当記事は、
DCGANの理論は他の方に任せて、簡単・シンプルなコードで、サクッと動かすことを目的としています。

コード・サンプルデータセットは**[GitHub](https://github.com/hokuto-HIRANO/5Steps-PyTorch-DCGAN)**に載せています。

# 1. データの準備

## 1.1 データのダウンロード・前処理

**前処理済みのサンプルデータセット(sample_data)を[GitHub](https://github.com/hokuto-HIRANO/5Steps-PyTorch-DCGAN)に用意しているので、以下は読み飛ばしても大丈夫です。**

データは、アニメのキャラクターの顔を集めたデータセット[AnimeFace Character Dataset](http://www.nurs.or.jp/~nagadomi/animeface-character-dataset/README.html)を使用します。

しかし、AnimeFace Character Datasetは顔よりも広めに切り取ってあるので、[lbpcascade_animeface](https://github.com/nagadomi/lbpcascade_animeface)を使用して、さらに実際に顔の部分だけを切り出し、画像サイズを64 x 64 に整形しておきます。

![face_0_119_15.png](https://qiita-image-store.s3.amazonaws.com/0/247972/d40eefec-83f3-1bde-de9a-84f404d6a097.png)　　　==>　　　![0.jpg](https://qiita-image-store.s3.amazonaws.com/0/247972/fb03aa38-0b0f-44ce-c855-69b4f1cd7ea2.jpeg)

以上で、計4257枚のデータセットが出来上がりました。

## 1.2 データの読み込み

[GitHub](https://github.com/hokuto-HIRANO/5Steps-PyTorch-DCGAN)より**sample_data**をダウンロードしているテイで進めます。

```.py
import torch
from torch import nn, optim
from torchvision.utils import save_image
from torch.utils.data import Dataset, DataLoader, TensorDataset
from torchvision import transforms, datasets
import tqdm
from statistics import mean


# === 1. データの読み込み ===
# datasetrの準備
dataset = datasets.ImageFolder("sample_data/",
    transform=transforms.Compose([
        transforms.ToTensor()
]))

batch_size = 64

# dataloaderの準備
data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
```


# 2. Generatorの作成

潜在特徴100次元ベクトルzから、3(チャネル) x 64 x 64の画像を生成するモデルを作成します。

Generator ( z )　　=>　　画像 ( 3x64x64 )

```.py
class Generator(nn.Module):
    def __init__(self):
        super().__init__()
        self.main = nn.Sequential(

            nn.ConvTranspose2d(100, 256, 4, 1, 0, bias=False),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),

            nn.ConvTranspose2d(256, 128, 4, 2, 1, bias=False),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),

            nn.ConvTranspose2d(128, 64, 4, 2, 1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),

            nn.ConvTranspose2d(64, 32, 4, 2, 1, bias=False),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),

            nn.ConvTranspose2d(32, 3, 4, 2, 1, bias=False),
            nn.Tanh()
        )

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

# 3. Discriminatorの作成
3 x 64 x 64 の画像を、１次元のスカラーに変換するモデルを作成します。

Discriminator ( 画像 )　　=>　　１次元スカラー

```.py
class Discriminator(nn.Module):
    def __init__(self):
        super().__init__()
        self.main = nn.Sequential(
            
            nn.Conv2d(3, 32, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),

            nn.Conv2d(32, 64, 4, 2, 1, bias=False),
            nn.BatchNorm2d(64),
            nn.LeakyReLU(0.2, inplace=True),

            nn.Conv2d(64, 128, 4, 2, 1, bias=False),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.2, inplace=True),

            nn.Conv2d(128, 256, 4, 2, 1, bias=False),
            nn.BatchNorm2d(256),
            nn.LeakyReLU(0.2, inplace=True),
            
            nn.Conv2d(256, 1, 4, 1, 0, bias=False),
        )

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


# 4. 訓練関数の作成

GANには、GeneratorとDiscriminatorの２つのモデルがあり、それぞれ交互に訓練していきます。

少し長くなりますが、以下のような流れで訓練していきます。（コードを見た方が分かりやすいかも）

(１) 潜在特徴100次元ベクトルzから、Generatorを使用して偽画像を生成する
(２) 偽画像をDiscriminatorで識別し、偽画像を実画像と騙せるようにGeneratorを学習する
(３) 実画像をDiscriminatorで識別する
(４) 偽画像を偽画像、実画像を実画像と識別できるようにDiscriminatorを学習する

```.py
model_G = Generator().to("cuda:0")
model_D = Discriminator().to("cuda:0")

params_G = optim.Adam(model_G.parameters(),
    lr=0.0002, betas=(0.5, 0.999))
params_D = optim.Adam(model_D.parameters(),
    lr=0.0002, betas=(0.5, 0.999))

# 潜在特徴100次元ベクトルz
nz = 100

# ロスを計算するときのラベル変数
ones = torch.ones(batch_size).to("cuda:0") # 正例 1
zeros = torch.zeros(batch_size).to("cuda:0") # 負例 0
loss_f = nn.BCEWithLogitsLoss()

# 途中結果の確認用の潜在特徴z
check_z = torch.randn(batch_size, nz, 1, 1).to("cuda:0")


# 訓練関数
def train_dcgan(model_G, model_D, params_G, params_D, data_loader):
    log_loss_G = []
    log_loss_D = []
    for real_img, _ in tqdm.tqdm(data_loader):
        batch_len = len(real_img)


        # == Generatorの訓練 ==
        # 偽画像を生成
        z = torch.randn(batch_len, nz, 1, 1).to("cuda:0")
        fake_img = model_G(z)

        # 偽画像の値を一時的に保存 => 注(１)
        fake_img_tensor = fake_img.detach()

        # 偽画像を実画像（ラベル１）と騙せるようにロスを計算
        out = model_D(fake_img)
        loss_G = loss_f(out, ones[: batch_len])
        log_loss_G.append(loss_G.item())

        # 微分計算・重み更新 => 注（２）
        model_D.zero_grad()
        model_G.zero_grad()
        loss_G.backward()
        params_G.step()


        # == Discriminatorの訓練 ==
        # sample_dataの実画像
        real_img = real_img.to("cuda:0")
        
        # 実画像を実画像（ラベル１）と識別できるようにロスを計算
        real_out = model_D(real_img)
        loss_D_real = loss_f(real_out, ones[: batch_len])

        # 計算省略 => 注（１）
        fake_img = fake_img_tensor

        # 偽画像を偽画像（ラベル０）と識別できるようにロスを計算
        fake_out = model_D(fake_img_tensor)
        loss_D_fake = loss_f(fake_out, zeros[: batch_len])

        # 実画像と偽画像のロスを合計
        loss_D = loss_D_real + loss_D_fake
        log_loss_D.append(loss_D.item())

        # 微分計算・重み更新 => 注（２）
        model_D.zero_grad()
        model_G.zero_grad()
        loss_D.backward()
        params_D.step()

    return mean(log_loss_G), mean(log_loss_D)
```

注（１）：pytorchでは同じTensorに対して２回バックプロパゲーションを計算できないので、一時的に保存しておいたTensorを使用して、無駄な計算（偽画像の生成）を省略する

注（２）：計算グラフがGeneratorとDiscriminatorの両方に関わっているので、両モデルの勾配を初期化してから、微分計算・重み更新を行う

# 5. DCGANの訓練スタート

その後、epoch数だけ訓練関数を実行します。

```.py
for epoch in range(300):
    train_dcgan(model_G, model_D, params_G, params_D, data_loader)
    
    # 訓練途中のモデル・生成画像の保存
    if epoch % 10 == 0:
        torch.save(
            model_G.state_dict(),
            "Weight_Generator/G_{:03d}.pth".format(epoch),
            pickle_protocol=4)
        torch.save(
            model_D.state_dict(),
            "Weight_Discriminator/D_{:03d}.pth".format(epoch),
            pickle_protocol=4)

        generated_img = model_G(check_z)
        save_image(generated_img,
                   "Generated_Image/{:03d}.jpg".format(epoch))
```

# 訓練後Generatorが作った画像例

![290.jpg](https://qiita-image-store.s3.amazonaws.com/0/247972/1515fedd-a8a3-06a2-88d3-bc8a6e6cef5c.jpeg)

データやモデル、ハイパラなど、特にこだわっていない中でこの程度の出来でした。

# 参考

- [現場で使える！PyTorch開発入門 深層学習モデルの作成とアプリケーションへの実装](https://www.amazon.co.jp/%E7%8F%BE%E5%A0%B4%E3%81%A7%E4%BD%BF%E3%81%88%E3%82%8B%EF%BC%81PyTorch%E9%96%8B%E7%99%BA%E5%85%A5%E9%96%80-%E6%B7%B1%E5%B1%A4%E5%AD%A6%E7%BF%92%E3%83%A2%E3%83%87%E3%83%AB%E3%81%AE%E4%BD%9C%E6%88%90%E3%81%A8%E3%82%A2%E3%83%97%E3%83%AA%E3%82%B1%E3%83%BC%E3%82%B7%E3%83%A7%E3%83%B3%E3%81%B8%E3%81%AE%E5%AE%9F%E8%A3%85-AI-TECHNOLOGY-%E4%B8%96%E6%A9%8B/dp/479815718X)

- [pytorchで書いたDCGANでアニメキャラの顔を生成する](https://qiita.com/phyblas/items/bcab394e1387f15f66ee)
- [データセット：AnimeFace Character Dataset](http://www.nurs.or.jp/~nagadomi/animeface-character-dataset/README.html)
- [前処理スクリプト：lbpcascade_animeface](https://github.com/nagadomi/lbpcascade_animeface)