20.2. Глубокие свёрточные генеративно-состязательные сети

В Section 20.1 мы представили основные идеи, лежащие в основе работы GAN. Мы показали, что они могут брать примеры из некоторого простого распределения, которое легко сэмплировать, например, равномерного или нормального, и преобразовывать их в примеры, выглядящие так, будто они соответствуют распределению некоторого набора данных. И хотя наш пример с аппроксимацией двумерного гауссовского распределения хорошо проиллюстрировал суть, он не особенно впечатляет.

В этом разделе мы продемонстрируем, как можно использовать GAN для генерации фотореалистичных изображений. Мы будем опираться на глубокие свёрточные GAN (DCGAN), представленные в Radford et al. (2015). Мы заимствуем свёрточную архитектуру, которая зарекомендовала себя как очень эффективная для дискриминативных задач компьютерного зрения, и покажем, как с помощью GAN её можно использовать для генерации фотореалистичных изображений.

import warnings
import torch
import torchvision
from torch import nn
from d2l import torch as d2l

20.2.1. Раздел учебника: Набор данных по покемонам

Используемый нами набор данных представляет собой коллекцию спрайтов покемонов, полученных с сайта pokemondb. Сначала скачайте, распакуйте и загрузите этот набор данных.

#@save
d2l.DATA_HUB['pokemon'] = (d2l.DATA_URL + 'pokemon.zip',
                           'c065c0e2593b8b161a2d7873e42418bf6a21106c')

data_dir = d2l.download_extract('pokemon')
pokemon = torchvision.datasets.ImageFolder(data_dir)

Мы приводим размер каждого изображения к \(64\times 64\). Преобразование ToTensor проецирует значение пикселя в \([0, 1]\), тогда как наш генератор использует функцию tanh для получения выходных значений в диапазоне \([-1, 1]\). Следовательно, мы нормализуем данные, используя среднее значение \(0.5\) и стандартное отклонение \(0.5\), чтобы привести их в соответствие с требуемым диапазоном значений.

batch_size = 256
transformer = torchvision.transforms.Compose([
    torchvision.transforms.Resize((64, 64)),
    torchvision.transforms.ToTensor(),
    torchvision.transforms.Normalize(0.5, 0.5)
])
pokemon.transform = transformer
data_iter = torch.utils.data.DataLoader(
    pokemon, batch_size=batch_size,
    shuffle=True, num_workers=d2l.get_dataloader_workers())

Давайте визуализируем первые 20 изображений.

warnings.filterwarnings('ignore')
d2l.set_figsize((4, 4))
for X, y in data_iter:
    imgs = X[:20,:,:,:].permute(0, 2, 3, 1)/2+0.5
    d2l.show_images(imgs, num_rows=4, num_cols=5)
    break

20.2.2. Генератор

Генератор должен отображать переменную шума \(\mathbf z\in\mathbb R^d\), являющуюся вектором длины \(d\), в RGB-изображение с шириной и высотой \(64\times 64\). В Section 14.11 мы ввели полностью свёрточную сеть, использующую слой транспонированной свёртки (см. Section 14.10) для увеличения размера входа. Базовый блок генератора содержит слой транспонированной свёртки, за которым следуют батч-нормализация и активация ReLU.

class G_block(nn.Module):
    def __init__(self, out_channels, in_channels=3, kernel_size=4, strides=2,
                 padding=1, **kwargs):
        super(G_block, self).__init__(**kwargs)
        self.conv2d_trans = nn.ConvTranspose2d(in_channels, out_channels,
                                kernel_size, strides, padding, bias=False)
        self.batch_norm = nn.BatchNorm2d(out_channels)
        self.activation = nn.ReLU()

    def forward(self, X):
        return self.activation(self.batch_norm(self.conv2d_trans(X)))

По умолчанию транспонированный свёрточный слой использует ядро \(k_h = k_w = 4\), шаг \(s_h = s_w = 2\) и дополнение \(p_h = p_w = 1\). При входной форме \(n_h^{'} \times n_w^{'} = 16 \times 16\) блок генератора удваивает ширину и высоту входных данных.

(20.2.1)\[\begin{split}\begin{aligned} n_h^{'} \times n_w^{'} &= [(n_h k_h - (n_h-1)(k_h-s_h)- 2p_h] \times [(n_w k_w - (n_w-1)(k_w-s_w)- 2p_w]\\ &= [(k_h + s_h (n_h-1)- 2p_h] \times [(k_w + s_w (n_w-1)- 2p_w]\\ &= [(4 + 2 \times (16-1)- 2 \times 1] \times [(4 + 2 \times (16-1)- 2 \times 1]\\ &= 32 \times 32 .\\ \end{aligned}\end{split}\]
x = torch.zeros((2, 3, 16, 16))
g_blk = G_block(20)
g_blk(x).shape

Если изменить транспонированный свёрточный слой на ядро размера \(4\times 4\), шаг \(1\times 1\) и нулевое дополнение, то при входном размере \(1 \times 1\) ширина и высота выхода увеличатся на 3 соответственно.

x = torch.zeros((2, 3, 1, 1))
g_blk = G_block(20, strides=1, padding=0)
g_blk(x).shape

Генератор состоит из четырёх базовых блоков, которые увеличивают ширину и высоту входа с 1 до 32. При этом латентная переменная сначала проецируется в \(64\times 8\) каналов, а затем количество каналов уменьшается вдвое на каждом шаге. В завершение используется слой транспонированной свёртки для формирования выходных данных. Он дополнительно удваивает ширину и высоту, чтобы соответствовать требуемой \(64\times 64\) форме, и уменьшает размер канала до \(3\). Функция активации tanh применяется для проецирования выходных значений в диапазон \((-1, 1)\).

n_G = 64
net_G = nn.Sequential(
    G_block(in_channels=100, out_channels=n_G*8,
            strides=1, padding=0),                  # Output: (64 * 8, 4, 4)
    G_block(in_channels=n_G*8, out_channels=n_G*4), # Output: (64 * 4, 8, 8)
    G_block(in_channels=n_G*4, out_channels=n_G*2), # Output: (64 * 2, 16, 16)
    G_block(in_channels=n_G*2, out_channels=n_G),   # Output: (64, 32, 32)
    nn.ConvTranspose2d(in_channels=n_G, out_channels=3,
                       kernel_size=4, stride=2, padding=1, bias=False),
    nn.Tanh())  # Output: (3, 64, 64)

Сгенерируем латентную переменную размерности 100, чтобы проверить выходную форму генератора.

x = torch.zeros((1, 100, 1, 1))
net_G(x).shape

20.2.3. Дискриминатор

Дискриминатор — это обычная свёрточная нейронная сеть, за исключением того, что в качестве функции активации используется leaky ReLU. Учитывая \(\alpha \in[0, 1]\), его определение имеет вид

(20.2.2)\[\begin{split}\textrm{leaky ReLU}(x) = \begin{cases}x & \textrm{if}\ x > 0\\ \alpha x &\textrm{otherwise}\end{cases}.\end{split}\]

Как видно, при \(\alpha=0\) это обычная функция ReLU, а при \(\alpha=1\) — тождественная функция. При \(\alpha \in (0, 1)\) leaky ReLU представляет собой нелинейную функцию, которая выдает ненулевое значение для отрицательного входа. Она призвана решить проблему «засыпания» ReLU, при которой нейрон может постоянно выдавать отрицательное значение и, следовательно, не может продвигаться вперед, поскольку градиент функции ReLU равен нулю.

alphas = [0, .2, .4, .6, .8, 1]
x = torch.arange(-2, 1, 0.1)
Y = [nn.LeakyReLU(alpha)(x).detach().numpy() for alpha in alphas]
d2l.plot(x.detach().numpy(), Y, 'x', 'y', alphas)

Базовым блоком дискриминатора является слой свёртки, за которым следуют слой батч-нормализации и функция активации leaky ReLU. Гиперпараметры слоя свёртки аналогичны гиперпараметрам слоя транспонированной свёртки в блоке генератора.

class D_block(nn.Module):
    def __init__(self, out_channels, in_channels=3, kernel_size=4, strides=2,
                padding=1, alpha=0.2, **kwargs):
        super(D_block, self).__init__(**kwargs)
        self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size,
                                strides, padding, bias=False)
        self.batch_norm = nn.BatchNorm2d(out_channels)
        self.activation = nn.LeakyReLU(alpha, inplace=True)

    def forward(self, X):
        return self.activation(self.batch_norm(self.conv2d(X)))

Базовый блок с настройками по умолчанию уменьшает ширину и высоту входных данных вдвое, как мы продемонстрировали в Section 7.3. Например, при заданной форме входа \(n_h = n_w = 16\), форме ядра \(k_h = k_w = 4\), шаге \(s_h = s_w = 2\) и дополнении \(p_h = p_w = 1\), форма выхода будет следующей:

(20.2.3)\[\begin{split}\begin{aligned} n_h^{'} \times n_w^{'} &= \lfloor(n_h-k_h+2p_h+s_h)/s_h\rfloor \times \lfloor(n_w-k_w+2p_w+s_w)/s_w\rfloor\\ &= \lfloor(16-4+2\times 1+2)/2\rfloor \times \lfloor(16-4+2\times 1+2)/2\rfloor\\ &= 8 \times 8 .\\ \end{aligned}\end{split}\]
x = torch.zeros((2, 3, 16, 16))
d_blk = D_block(20)
d_blk(x).shape

Дискриминатор — это зеркальное отражение генератора.

n_D = 64
net_D = nn.Sequential(
    D_block(n_D),  # Output: (64, 32, 32)
    D_block(in_channels=n_D, out_channels=n_D*2),  # Output: (64 * 2, 16, 16)
    D_block(in_channels=n_D*2, out_channels=n_D*4),  # Output: (64 * 4, 8, 8)
    D_block(in_channels=n_D*4, out_channels=n_D*8),  # Output: (64 * 8, 4, 4)
    nn.Conv2d(in_channels=n_D*8, out_channels=1,
              kernel_size=4, bias=False))  # Output: (1, 1, 1)

Последним слоем является слой свёртки с выходным каналом \(1\), который формирует одно предсказанное значение.

x = torch.zeros((1, 3, 64, 64))
net_D(x).shape

20.2.4. Обучение

В отличие от базовой GAN в Section 20.1, мы используем одинаковую скорость обучения для генератора и дискриминатора, поскольку они схожи. Кроме того, мы изменяем \(\beta_1\) в Adam (Section 12.10) с \(0.9\) на \(0.5\). Это снижает сглаживание импульса, то есть экспоненциально взвешенного скользящего среднего прошлых градиентов, чтобы лучше реагировать на быстро меняющиеся градиенты, поскольку генератор и дискриминатор соревнуются друг с другом. Кроме того, случайно сгенерированный шум Z представляет собой 4-мерный тензор, и мы используем GPU для ускорения вычислений.

def train(net_D, net_G, data_iter, num_epochs, lr, latent_dim,
          device=d2l.try_gpu()):
    loss = nn.BCEWithLogitsLoss(reduction='sum')
    for w in net_D.parameters():
        nn.init.normal_(w, 0, 0.02)
    for w in net_G.parameters():
        nn.init.normal_(w, 0, 0.02)
    net_D, net_G = net_D.to(device), net_G.to(device)
    trainer_hp = {'lr': lr, 'betas': [0.5,0.999]}
    trainer_D = torch.optim.Adam(net_D.parameters(), **trainer_hp)
    trainer_G = torch.optim.Adam(net_G.parameters(), **trainer_hp)
    animator = d2l.Animator(xlabel='epoch', ylabel='loss',
                            xlim=[1, num_epochs], nrows=2, figsize=(5, 5),
                            legend=['discriminator', 'generator'])
    animator.fig.subplots_adjust(hspace=0.3)
    for epoch in range(1, num_epochs + 1):
        # Train one epoch
        timer = d2l.Timer()
        metric = d2l.Accumulator(3)  # loss_D, loss_G, num_examples
        for X, _ in data_iter:
            batch_size = X.shape[0]
            Z = torch.normal(0, 1, size=(batch_size, latent_dim, 1, 1))
            X, Z = X.to(device), Z.to(device)
            metric.add(d2l.update_D(X, Z, net_D, net_G, loss, trainer_D),
                       d2l.update_G(Z, net_D, net_G, loss, trainer_G),
                       batch_size)
        # Show generated examples
        Z = torch.normal(0, 1, size=(21, latent_dim, 1, 1), device=device)
        # Normalize the synthetic data to N(0, 1)
        fake_x = net_G(Z).permute(0, 2, 3, 1) / 2 + 0.5
        imgs = torch.cat(
            [torch.cat([
                fake_x[i * 7 + j].cpu().detach() for j in range(7)], dim=1)
             for i in range(len(fake_x)//7)], dim=0)
        animator.axes[1].cla()
        animator.axes[1].imshow(imgs)
        # Show the losses
        loss_D, loss_G = metric[0] / metric[2], metric[1] / metric[2]
        animator.add(epoch, (loss_D, loss_G))
    print(f'loss_D {loss_D:.3f}, loss_G {loss_G:.3f}, '
          f'{metric[2] / timer.stop():.1f} examples/sec on {str(device)}')

Мы обучаем модель всего несколько эпох исключительно в демонстрационных целях. Для достижения лучшей производительности переменную num_epochs следует установить в большее значение.

latent_dim, lr, num_epochs = 100, 0.005, 20
train(net_D, net_G, data_iter, num_epochs, lr, latent_dim)

20.2.5. Итоги

  • Архитектура DCGAN содержит четыре свёрточных слоя для дискриминатора и четыре свёрточных слоя с «дробным шагом» для генератора.

  • Дискриминатор состоит из 4 слоёв с использованием свёрток с шагом, батч-нормализации (за исключением входного слоя) и функций активации leaky ReLU.

  • Leaky ReLU — это нелинейная функция, которая выдает ненулевое значение при отрицательном входе. Она призвана решить проблему «умирающего ReLU» и облегчает прохождение градиентов через архитектуру.

20.2.6. Упражнения

  1. Что произойдет, если мы используем стандартную функцию активации ReLU вместо leaky ReLU?

  2. Примените DCGAN к набору данных Fashion-MNIST и выясните, какие категории распознаются хорошо, а какие — плохо.

Обсуждения