【pytorchで深層生成モデル#13】学習テクニックまとめ(DCGAN)

# [1]
!nvidia-smi

# [2]
# データ作成に使用するライブラリ
from torchvision import datasets
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

# モデル作成に使用するライブラリ
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F

# よく使用するライブラリ
import matplotlib.pyplot as plt
import numpy as np

# 学習テクニックライブラリ
import torchsummary # モデルを確認するライブラリ
import pickle # モデルの保存と読み込みするライブラリ
import torchvision.utils as vutils # 画像データを保存するライブラリ
from torch.utils.data import Dataset # データセット作成に使用するライブラリ
from pathlib import Path # データセット作成に使用するライブラリ
from PIL import Image # データセット作成に使用するライブラリ

torch.manual_seed(1)

# [3]
batch_size = 100
n_channel = 100
n_epoch = 10

# [4]
!mkdir './data'
dataset = datasets.MNIST(root='.', download=True)

for idx, (img, _) in enumerate(dataset):
    img.save('./data/{:05d}.jpg'.format(idx))

# [5]
!find /content/data/* -type f | wc -l

# [6]
class Image_dataset(Dataset):

    def __init__(self, img_dir, transform=None):
        self.img_paths = self._get_img_paths(img_dir)
        self.transform = transform

    def __getitem__(self, index):
        path = self.img_paths[index]
        img = Image.open(path)
        if self.transform is not None:
            img = self.transform(img)

        return img

    def _get_img_paths(self, img_dir):
        img_dir = Path(img_dir)
        img_paths = [p for p in img_dir.iterdir() if p.suffix in [".jpg", ".jpeg", ".png", ".bmp"]]
        return img_paths

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

# [7]
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
dataset = Image_dataset("/content/data", transform)
dataloader = DataLoader(dataset, batch_size=100)

# [8]
# データの読み込み
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
device

# [9]
# Google ドライブにマウント
from google.colab import drive
drive.mount('/content/gdrive')

# [10]
%cd '/content/gdrive/MyDrive/'

# [11]
!mkdir './result_image'
!mkdir './result_model'

# [12]
class Generator(nn.Module):

    def __init__(self):

        super(Generator, self).__init__()

        self.layers = nn.ModuleDict({
            'layer0': nn.Sequential(
                nn.ConvTranspose2d(n_channel, 512, 3, 1, 0),
                nn.BatchNorm2d(512),
                nn.ReLU()
            ), 
            'layer1': nn.Sequential(
                nn.ConvTranspose2d(512, 256, 3, 2, 0),
                nn.BatchNorm2d(256),
                nn.ReLU()
            ),
            'layer2': nn.Sequential(
                nn.ConvTranspose2d(256, 128, 4, 2, 1),
                nn.BatchNorm2d(128),
                nn.ReLU()
            ),
            'layer3': nn.Sequential(
                nn.ConvTranspose2d(128, 1, 4, 2, 1),
                nn.Tanh()
            )
        })

    def forward(self, z):

        for layer in self.layers.values(): 
            z = layer(z)
        return z

# [13]
class Discriminator(nn.Module):

    def __init__(self):

        super(Discriminator, self).__init__()

        self.layers = nn.ModuleDict({
            'layer0': nn.Sequential(
                nn.Conv2d(1, 128, 4, 2, 1),
                nn.LeakyReLU(negative_slope=0.2)
            ),
            'layer1': nn.Sequential(
                nn.Conv2d(128, 256, 4, 2, 1),
                nn.BatchNorm2d(256),
                nn.LeakyReLU(negative_slope=0.2)
            ),
            'layer2': nn.Sequential(
                nn.Conv2d(256, 512, 3, 2, 0),
                nn.BatchNorm2d(512),
                nn.LeakyReLU(negative_slope=0.2)
            ),
            'layer3': nn.Sequential(
                nn.Conv2d(512, 1, 3, 1, 0),
                nn.Sigmoid()
            )    
        })

    def forward(self, x):
        for layer in self.layers.values():
            x = layer(x)
        return x.squeeze()

# [14]
generator = Generator().to(device)
discriminator = Discriminator().to(device)

# [15]
criterion = nn.BCELoss()

optimizerG = optim.Adam(generator.parameters(), lr = 0.0002, betas=(0.5, 0.999), weight_decay=1e-5)
optimizerD = optim.Adam(discriminator.parameters(), lr = 0.0002, betas=(0.5, 0.999), weight_decay=1e-5)

# [16]
print(generator)

# [17]
torchsummary.summary(generator, (100, 1, 1))

# [18]
print(discriminator)

# [19]
torchsummary.summary(discriminator, (1, 28, 28))

# [20]
pre = 0

# []
generator = pickle.load(open('./result_model/generator_{}.pth'.format(pre), 'rb'))
discriminator = pickle.load(open('./result_model/discriminator_{}.pth'.format(pre), 'rb'))
# generator.load_state_dict(torch.load('./result_model/generator_{}.pth'.format(pre)))
# discriminator.load_state_dict(torch.load('./result_model/discriminator_{}.pth'.format(pre)))

# [21]
G_losses = []
D_losses = []
D_x_list = []
D_G_z1_list = []
D_G_z2_list = []

# 学習のループ
for epoch in range(pre, n_epoch):

    for x in dataloader:

      # 前準備
      real_image = x.to(device) # 本物の画像データ
      noise = torch.randn(batch_size, n_channel, 1, 1, device=device) # ノイズ作成
      real_target = torch.full((batch_size,), 1., device=device) # 本物ラベル
      fake_target = torch.full((batch_size,), 0., device=device) # 偽物ラベル
      
      # discriminatorの学習(本物画像の学習)
      discriminator.zero_grad() 
      y = discriminator(real_image)
      errD_real = criterion(y, real_target)
      D_x = y.mean().item()

      # discriminatorの学習(偽物画像の学習)
      fake_image = generator(noise)
      y = discriminator(fake_image.detach())
      errD_fake = criterion(y, fake_target)
      D_G_z1 = y.mean().item()

      # discriminatorの更新
      errD = errD_real + errD_fake
      errD.backward()
      optimizerD.step()

      # generatorの学習
      generator.zero_grad()
      y = discriminator(fake_image)
      errG = criterion(y, real_target)
      errG.backward()
      D_G_z2 = y.mean().item()
      optimizerG.step()

      # 損失関数のリスト作成
      D_losses.append(errD.item())
      G_losses.append(errG.item())
      D_x_list.append(D_x)
      D_G_z1_list.append(D_G_z1)
      D_G_z1_list.append(D_G_z2)

    # 画像データとモデルの保存
    if (epoch + 1) % 2 == 0:  
      vutils.save_image(fake_image.detach(), './result_image/fake_samples_{:03d}.png'.format(epoch + 1), normalize=True, nrow=10)
      pickle.dump(generator, open('./result_model/generator_{}.pth'.format(epoch + 1), 'wb'))
      pickle.dump(discriminator, open('./result_model/discriminator_{}.pth'.format(epoch + 1), 'wb'))
      # torch.save(generator.state_dict(), './result_model/generator_{}.pth'.format(epoch + 1))
      # torch.save(discriminator.state_dict(), './result_model/discriminator_{}.pth'.format(epoch + 1))

    print('Epoch:{}/{}, Loss_D: {:.3f}, Loss_G: {:.3f}, D(x): {:.3f}, D(G(z)): {:.3f}/{:.3f}'
    .format(epoch + 1, n_epoch, errD.item(), errG.item(), D_x, D_G_z1, D_G_z2))

# [22]
# fake_image作成
generator.eval
noise = torch.randn(50, n_channel, 1, 1, device=device)
sample_images = generator(noise)

# fake_image可視化
fig = plt.figure(figsize=(20,20))
plt.subplots_adjust(wspace=0.1, hspace=-0.8)
for i in range(50):
    ax = fig.add_subplot(5, 10, i+1, xticks=[], yticks=[])
    ax.imshow(sample_images[i,].view(28,28).cpu().detach(), "gray")