使用 Pytorch 构建 Vanilla GAN

一、说明

使用生成对抗网络 (GAN) 可以执行生成机器学习。换句话说，您可以确保模型学会生成新数据，例如图像。

像这样：

在今天的文章中，您将创建一个简单的 GAN，也称为vanilla GAN。它类似于 Goodfellow 等人 (2014) 首次创建的生成对抗网络。阅读本文后，您将：1）了解什么是 GAN 以及它如何工作。2）能够使用 Python 和 PyTorch 构建一个简单的 GAN。3）已经产生了你的第一个 GAN 结果。
让我们看看吧。

二、什么是 GAN？

在开始构建简单的 GAN 之前，最好先简单回顾一下 GAN 是什么。如果你想更详细地了解 GAN 的行为，请务必阅读我对 GAN 的简单介绍。不过，我们也会在这里简要介绍一下。让我们来看看 GAN 的通用架构：

我们可以看到，GAN 由两个独立的模型组成。第一个模型称为生成器，它学习将噪声样本（通常取自标准正态分布）转换为假图像。然后，该图像被输入到鉴别器，鉴别器判断图像是假的还是真的。利用从这一判断中得出的损失，网络被联合优化，之后该过程再次开始。

你也可以将这个过程与造假者和警察的过程进行比较。生成器充当造假者，而警察的任务是抓捕他们。当警察抓到更多的假图像时，造假者必须学会制作更好的结果。这正是发生的事情：通过判别器在判断图像是假的还是真的方面变得越来越好，生成器最终在生成假图像方面也变得越来越好。因此，生成器在经过训练后可以单独用于生成图像。

现在，是时候开始构建 GAN 了。请注意，如果您希望在生产中使用 GAN，则更倾向于使用更现代的方法，例如 DCGAN（未来文章）（原因很简单，因为最初，原始 GAN 没有使用任何卷积层）。但是，如果您想从 GAN 开始，下面将生成的示例是一个非常好的起点 — 之后您可以继续使用 DCGAN 及更多。让我们来看看！

三、使用 PyTorch 的简单 GAN（完整解释的代码示例）

现在让我们看一下构建一个简单的生成对抗网络，它看起来像 Goodfellow 等人 (2014) 提出的原始 GAN。

导入依赖项
当您想要运行要创建的代码时，您需要确保一些依赖项已安装到您的环境中。这些依赖项如下：

基于 3.x 的 Python 版本，您将使用它来运行这些脚本。
PyTorch 及其对应版本的 Torchvision 用于使用 MNIST 数据训练神经网络。
NumPy 用于数字处理。
Matplotlib 用于可视化图像。
现在，创建一个 Python 文件或基于 Python 的 Notebook，并导入以下内容：

import os
import torch
from torch import nn
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader
from torchvision import transforms
import numpy as np
import matplotlib.pyplot as plt
import uuid

对于某些操作系统功能，您将需要os。uuid将用于生成唯一的运行标识符，这对于保存中间模型和生成的图像很有用；即用于内务管理。torch将用于训练神经网络，因此您需要导入其nn库。MNIST数据集将被使用，因此需要导入，并且将使用 加载DataLoader。加载数据时，您将把它转换成张量格式并对图像进行规范化，需要transforms。最后，对于数字处理和可视化，您需要numpy和matplotlib.pyplot。

3.1 配置变量

现在您已经指定了导入，现在是时候确定将在整个训练过程中使用的可配置变量了。以下是您将创建的内容以及您需要它的原因：

迭代次数：每个训练过程包含整个训练集的固定迭代次数，即迭代次数。我们将其设置为 50，但您可以选择任意数字。请注意，50 将产生可接受的结果；更多迭代次数可能会进一步改善结果。
噪声维度：回想一下，生成器将被输入一个变量，该变量作为多维潜在分布的样本。这些话很难说清楚，我们从最终会成形的景观中采样，以便生成器能够生成好的示例。这个景观的维数以及从中采样的向量将定义为NOISE_DIMENSION。
批处理大小：在一个时期内，我们分批将数据馈送到网络 — 即不是一次性全部馈送。原因很简单 — 因为否则数据无法装入内存。我们将批处理大小设置为 128 个样本，但这个大小可以更高，具体取决于您系统的硬件。
是否在 GPU 上进行训练：根据 GPU 的可用性，您可以选择使用它进行训练 - 否则将使用您的 CPU。
唯一运行标识符：与内务管理相关。您将看到，在训练过程中，中间模型和图像将存储在磁盘上，以便您可以跟踪训练进度。为此，将创建一个具有唯一标识符的文件夹；因此UNIQUE_RUN_ID。
打印第 n 个批次后的统计数据：在通过网络向前馈送小批次后，将在每个批次后打印统计数据n-th。目前，我们将其设置为 50。
优化器学习率和优化器 beta。生成器和鉴别器的优化器将使用学习率和 beta 值进行初始化。我们根据先前的研究将它们设置为被认为能够产生可接受结果的值。
生成器输出的输出形状将用于初始化生成器的最后一层和鉴别器的第一层。它必须是单个图像的所有形状维度的乘积。在我们的例子中，MNIST 数据集有28x28x1图像。

# Configurable variables
NUM_EPOCHS = 50
NOISE_DIMENSION = 50
BATCH_SIZE = 128
TRAIN_ON_GPU = True
UNIQUE_RUN_ID = str(uuid.uuid4())
PRINT_STATS_AFTER_BATCH = 50
OPTIMIZER_LR = 0.0002
OPTIMIZER_BETAS = (0.5, 0.999)
GENERATOR_OUTPUT_IMAGE_SHAPE = 28 * 28 * 1

3.2 、PyTorch 加速

您可以使用一些方法来使您的 PyTorch 代码运行得更快：这就是您接下来要编写这些加速的原因。

# Speed ups
torch.autograd.set_detect_anomaly(False)
torch.autograd.profiler.profile(False)
torch.autograd.profiler.emit_nvtx(False)
torch.backends.cudnn.benchmark = True

3.3 构建生成器

现在我们已经编写了一些准备代码，是时候构建实际的生成器了！与深度卷积 GAN（本质上遵循您今天将创建的vanilla GAN）相反，此生成器不使用卷积层。以下是生成器的代码：

class Generator(nn.Module):
  """
    Vanilla GAN Generator
  """
  def __init__(self,):
    super().__init__()
    self.layers = nn.Sequential(
      # First upsampling
      nn.Linear(NOISE_DIMENSION, 128, bias=False),
      nn.BatchNorm1d(128, 0.8),
      nn.LeakyReLU(0.25),
      # Second upsampling
      nn.Linear(128, 256, bias=False),
      nn.BatchNorm1d(256, 0.8),
      nn.LeakyReLU(0.25),
      # Third upsampling
      nn.Linear(256, 512, bias=False),
      nn.BatchNorm1d(512, 0.8),
      nn.LeakyReLU(0.25),
      # Final upsampling
      nn.Linear(512, GENERATOR_OUTPUT_IMAGE_SHAPE, bias=False),
      nn.Tanh()
    )

  def forward(self, x):
    """Forward pass"""
    return self.layers(x)

您可以看到它是一个常规的 PyTorchnn.Module类，因此forward只需将数据输入到模型（在 中指定self.layers为nn.Sequential基于神经网络）即可执行传递。在我们的例子中，您将编写四个上采样块。中间块由一个nn.Linear（或密集连接的）层、一个BatchNorm1d用于批量标准化的层和 Leaky ReLU 组成。偏差设置为，False因为批量标准化层会使其无效。

最后的上采样层将中间数量的 512 个神经元转换为GENERATOR_OUTPUT_IMAGE_SHAPE，即28 * 28 * 1 = 784。使用 Tanh，输出被标准化为 范围[-1, 1]。

3.4 构建鉴别器

鉴别器比生成器更简单。它是一个单独的神经网络，从它的nn.Module类定义可以看出。它只是组成一个完全连接的神经网络，接受维度输入GENERATOR_OUTPUT_IMAGE_SHAPE（即生成器输出），并将其转换为[0, 1]Sigmoid 正则化预测，以判断图像是真是假。

class Discriminator(nn.Module):
“”"
Vanilla GAN Discriminator
“”"
def init(self):
super().init()
self.layers = nn.Sequential(
nn.Linear(GENERATOR_OUTPUT_IMAGE_SHAPE, 1024),
nn.LeakyReLU(0.25),
nn.Linear(1024, 512),
nn.LeakyReLU(0.25),
nn.Linear(512, 256),
nn.LeakyReLU(0.25),
nn.Linear(256, 1),
nn.Sigmoid()
)

def forward(self, x):
“”“Forward pass”“”
return self.layers(x)

将所有东西合并为一个
好的，我们现在有两个不同的神经网络、一些导入和一些配置变量。是时候将所有内容合并为一个了！让我们从编写一些管理函数开始。

内部管理功能
回想一下，我们之前说过，中间模型将保存在一个文件夹中，并且还会生成图像。虽然我们稍后会实际实现这些调用，即使用它们，但您现在已经要编写它们了。我们的管理函数包含五个定义：

获取设备。回想一下，你为指定了True或。此定义将检查你是否要使用 GPU 以及它是否可用，否则指示 PyTorch 使用你的 CPU。FalseTRAIN_ON_GPU
为一次运行创建目录利用UNIQUE_RUN_ID为唯一运行生成一个目录。
生成图像将使用某些生成器（通常是您最近训练的生成器）生成 16 个示例并将它们存储到磁盘。
保存模型会将生成器和鉴别器的当前状态保存到磁盘。
打印训练进度会在屏幕上打印当前的损失值。

def get_device():
  """ Retrieve device based on settings and availability. """
  return torch.device("cuda:0" if torch.cuda.is_available() and TRAIN_ON_GPU else "cpu")
    
    
def make_directory_for_run():
  """ Make a directory for this training run. """
  print(f'Preparing training run {UNIQUE_RUN_ID}')
  if not os.path.exists('./runs'):
    os.mkdir('./runs')
  os.mkdir(f'./runs/{UNIQUE_RUN_ID}')


def generate_image(generator, epoch = 0, batch = 0, device=get_device()):
  """ Generate subplots with generated examples. """
  images = []
  noise = generate_noise(BATCH_SIZE, device=device)
  generator.eval()
  images = generator(noise)
  plt.figure(figsize=(10, 10))
  for i in range(16):
    # Get image
    image = images[i]
    # Convert image back onto CPU and reshape
    image = image.cpu().detach().numpy()
    image = np.reshape(image, (28, 28))
    # Plot
    plt.subplot(4, 4, i+1)
    plt.imshow(image, cmap='gray')
    plt.axis('off')
  if not os.path.exists(f'./runs/{UNIQUE_RUN_ID}/images'):
    os.mkdir(f'./runs/{UNIQUE_RUN_ID}/images')
  plt.savefig(f'./runs/{UNIQUE_RUN_ID}/images/epoch{epoch}_batch{batch}.jpg')


def save_models(generator, discriminator, epoch):
  """ Save models at specific point in time. """
  torch.save(generator.state_dict(), f'./runs/{UNIQUE_RUN_ID}/generator_{epoch}.pth')
  torch.save(discriminator.state_dict(), f'./runs/{UNIQUE_RUN_ID}/discriminator_{epoch}.pth')


def print_training_progress(batch, generator_loss, discriminator_loss):
  """ Print training progress. """
  print('Losses after mini-batch %5d: generator %e, discriminator %e' %
        (batch, generator_loss, discriminator_loss))

四、准备数据集

好的，整理完毕后，是时候开始编写用于准备数据集的功能了。这将是一个多阶段的过程。首先，我们MNIST从中加载数据集torchvision。加载后，样本将转换为 Tensor 格式并在范围内进行归一化，[-1, 1]以便它们与 Generator 生成的图像直接兼容。

但是，在加载完所有数据后，我们仍然需要对其进行批处理——请记住，您不会一次性将所有图像输入网络，而是会以批处理的方式进行。您还将对图像进行打乱。为了提高 PyTorch 的效率，工作者的数量将为 4，并pin_memory设置为 True。完成后，DataLoader将返回，以便可以使用。

def prepare_dataset():
  """ Prepare dataset through DataLoader """
  # Prepare MNIST dataset
  dataset = MNIST(os.getcwd(), download=True, train=True, transform=transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
  ]))
  # Batch and shuffle data with DataLoader
  trainloader = torch.utils.data.DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=4, pin_memory=True)
  # Return dataset through DataLoader
  return trainloader

五、初始化函数

您需要的其他一些定义与联合训练过程中将使用的模型、损失函数和优化器有关。

在 中initialize_models，您将初始化生成器和鉴别器，将它们移动到已配置的设备，然后返回。初始化二元交叉熵损失将在 中执行initialize_loss，最后，生成器和鉴别器的优化器将在 中初始化initialize_optimizers。再次，您将在稍后使用这些。

def initialize_models(device = get_device()):
  """ Initialize Generator and Discriminator models """
  generator = Generator()
  discriminator = Discriminator()
  # Move models to specific device
  generator.to(device)
  discriminator.to(device)
  # Return models
  return generator, discriminator


def initialize_loss():
  """ Initialize loss function. """
  return nn.BCELoss()


def initialize_optimizers(generator, discriminator):
  """ Initialize optimizers for Generator and Discriminator. """
  generator_optimizer = torch.optim.AdamW(generator.parameters(), lr=OPTIMIZER_LR,betas=OPTIMIZER_BETAS)
  discriminator_optimizer = torch.optim.AdamW(discriminator.parameters(), lr=OPTIMIZER_LR,betas=OPTIMIZER_BETAS)
  return generator_optimizer, discriminator_optimizer

六、前向和后向传递

使用初始化的模型，您将执行前向和后向传递。为此以及整个训练步骤，您将需要接下来创建的三个定义。第一个generate_noise用于在您之前配置的设备上生成维度number_of_images的噪声向量。noise_dimension

必须在每次训练步骤开始时有效地将梯度归零，并通过调用 来完成efficient_zero_grad()。最后，使用forward_and_backward，将使用一些模型、损失函数、数据和相应的目标计算前向和后向传递。然后返回损失的数值。

def generate_noise(number_of_images = 1, noise_dimension = NOISE_DIMENSION, device=None):
  """ Generate noise for number_of_images images, with a specific noise_dimension """
  return torch.randn(number_of_images, noise_dimension, device=device)


def efficient_zero_grad(model):
  """ 
    Apply zero_grad more efficiently
    Source: https://betterprogramming.pub/how-to-make-your-pytorch-code-run-faster-93079f3c1f7b
  """
  for param in model.parameters():
    param.grad = None


def forward_and_backward(model, data, loss_function, targets):
  """
    Perform forward and backward pass in a generic way. Returns loss value.
  """
  outputs = model(data)
  error = loss_function(outputs, targets)
  error.backward()
  return error.item()

七、执行训练步骤

现在我们已经定义了前向和后向传递的函数，现在是时候创建一个用于执行训练步骤的函数了。

回想一下，GAN 的训练步骤涉及多次前向和后向传递：一次使用鉴别器对真实图像进行传递，一次使用鉴别器对虚假图像进行传递，之后进行优化。然后，再次使用虚假图像来优化生成器。

下面，您将把这个过程编码为四个中间步骤。首先，您将准备一些东西，例如为真实数据和虚假数据设置标签值。在第二步中，训练鉴别器，然后在第三步中训练生成器。最后，您将在第四步中合并一些损失值并返回它们。

def perform_train_step(generator, discriminator, real_data, \
  loss_function, generator_optimizer, discriminator_optimizer, device = get_device()):
  """ Perform a single training step. """
  
  # 1. PREPARATION
  # Set real and fake labels.
  real_label, fake_label = 1.0, 0.0
  # Get images on CPU or GPU as configured and available
  # Also set 'actual batch size', whih can be smaller than BATCH_SIZE
  # in some cases.
  real_images = real_data[0].to(device)
  actual_batch_size = real_images.size(0)
  label = torch.full((actual_batch_size,1), real_label, device=device)
  
  # 2. TRAINING THE DISCRIMINATOR
  # Zero the gradients for discriminator
  efficient_zero_grad(discriminator)
  # Forward + backward on real images, reshaped
  real_images = real_images.view(real_images.size(0), -1)
  error_real_images = forward_and_backward(discriminator, real_images, \
    loss_function, label)
  # Forward + backward on generated images
  noise = generate_noise(actual_batch_size, device=device)
  generated_images = generator(noise)
  label.fill_(fake_label)
  error_generated_images =forward_and_backward(discriminator, \
    generated_images.detach(), loss_function, label)
  # Optim for discriminator
  discriminator_optimizer.step()
  
  # 3. TRAINING THE GENERATOR
  # Forward + backward + optim for generator, including zero grad
  efficient_zero_grad(generator)
  label.fill_(real_label)
  error_generator = forward_and_backward(discriminator, generated_images, loss_function, label)
  generator_optimizer.step()
  
  # 4. COMPUTING RESULTS
  # Compute loss values in floats for discriminator, which is joint loss.
  error_discriminator = error_real_images + error_generated_images
  # Return generator and discriminator loss so that it can be printed.
  return error_generator, error_discriminator

执行纪元
回想一下，训练 GAN 包含多个 epoch，而这些 epoch 本身又包含多个训练步骤。现在您已经为单个训练步骤编写了一些代码，是时候编写执行 epoch 的代码了。如下所示，您将迭代由 创建的批次DataLoader。使用每个批次执行一个训练步骤，并在必要时打印统计数据。

每个时期之后，模型都会被保存，并且 CUDA 内存会被清除。

def perform_epoch(dataloader, generator, discriminator, loss_function, \
    generator_optimizer, discriminator_optimizer, epoch):
  """ Perform a single epoch. """
  for batch_no, real_data in enumerate(dataloader, 0):
    # Perform training step
    generator_loss_val, discriminator_loss_val = perform_train_step(generator, \
      discriminator, real_data, loss_function, \
      generator_optimizer, discriminator_optimizer)
    # Print statistics and generate image after every n-th batch
    if batch_no % PRINT_STATS_AFTER_BATCH == 0:
      print_training_progress(batch_no, generator_loss_val, discriminator_loss_val)
      generate_image(generator, epoch, batch_no)
  # Save models on epoch completion.
  save_models(generator, discriminator, epoch)
  # Clear memory after every epoch
  torch.cuda.empty_cache()

开始训练过程
最后——最后一个定义！

在这个定义中，你将把所有内容合并在一起，以便能够真正进行训练。

首先，您需要确保为这次独特的运行创建一个新目录。然后，您需要将随机数生成器的种子设置为一个固定数字，这样初始化向量中的可变性就不会成为任何异常的原因。然后，您将检索准备好的（即经过打乱和分批的）数据集；初始化模型、损失和优化器；最后通过迭代指定的时期数来训练模型。

为了确保您的脚本开始运行，您将调用它train_dcgan()作为代码的最后一部分。

def train_dcgan():
  """ Train the DCGAN. """
  # Make directory for unique run
  make_directory_for_run()
  # Set fixed random number seed
  torch.manual_seed(42)
  # Get prepared dataset
  dataloader = prepare_dataset()
  # Initialize models
  generator, discriminator = initialize_models()
  # Initialize loss and optimizers
  loss_function = initialize_loss()
  generator_optimizer, discriminator_optimizer = initialize_optimizers(generator, discriminator)
  # Train the model
  for epoch in range(NUM_EPOCHS):
    print(f'Starting epoch {epoch}...')
    perform_epoch(dataloader, generator, discriminator, loss_function, \
      generator_optimizer, discriminator_optimizer, epoch)
  # Finished :-)
  print(f'Finished unique run {UNIQUE_RUN_ID}')


if __name__ == '__main__':
  train_dcgan()

Python GAN完整模型代码
为方便起见，可以在我的Github 机器学习存储库中找到完整的模型代码。

八、结果

现在，是时候运行你的模型了，例如使用python gan.py。

您应该看到模型开始相对较快地迭代，即使在 CPU 上也是如此。

在第一个时期，当我们打开为本次训练运行创建的文件夹中的文件时，我们可以看到随机噪声迅速改善为稍微可识别的数字：

周期 0，批次 0

纪元 0，批次 50

纪元 1，Bacth 0

第 1 阶段批次 50

第 2 阶段，第 0 批次

Epcoh 2，批次 50

第 3 阶段，第 0 批次

第 3 阶段，第 50 批次
在随后的几个时期中，随着越来越多的噪音消失，输出开始改善：瞧，你的第一个 GAN 已经完成了！

原文地址：https://blog.csdn.net/gongdiwudu/article/details/144147344

免责声明：本站文章内容转载自网络资源，如本站内容侵犯了原著者的合法权益，可联系本站删除。更多内容请关注自学内容网（zxcms.com）！