【指南】使用PyTorch构建Vanilla GAN

2024年11月25日 由 alex 发表 84 0

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


像这样:


10

11

12


在今天的文章中,你将创建一个简单的 GAN,也称为vanilla GAN。它类似于 Goodfellow 等人 (2014) 首次创建的生成对抗网络。


什么是 GAN?

在开始构建简单的 GAN 之前,最好先简单回顾一下 GAN 是什么。让我们来看看 GAN 的通用架构:


13


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


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


现在,是时候开始构建 GAN 了。


使用 PyTorch 的简单 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加载它。在加载数据时,你需要将其转换为Tensor格式并对图像进行归一化,这需要用到transforms。最后,对于数字处理和可视化,你将需要numpy和matplotlib.pyplot。


配置变量

既然你已经指定了导入的模块,现在是时候确定在整个训练过程中将使用的可配置变量了。以下是你将创建的内容以及为什么需要它们:

  • 迭代次数(epochs):每个训练过程都包含对整个训练集的固定次数的迭代,即迭代次数。我们将其设置为50,但你可以选择任何数字。请注意,50次可以产生可接受的结果;更多次数可能会进一步提高结果。
  • 噪声维度:回想一下,生成器将被输入一个变量,该变量作为从多维潜在分布中抽取的样本。简而言之,我们从一个景观中采样,这个景观最终将形成一个形状,以便生成器产生良好的示例。这个景观的维度以及从中采样的向量的维度将由NOISE_DIMENSION定义。
  • 批处理大小(batch size):在一个迭代次数内,我们分批通过网络向前馈送数据,即不是一次性全部馈送。原因很简单——因为否则它们将无法全部装入内存。我们将批处理大小设置为128个样本,但根据你的系统硬件,这个数值可以更高。
  • 是否在GPU上训练:根据GPU的可用性,你可以选择使用它进行训练——否则将使用你的CPU。
  • 唯一运行标识符:与日常管理相关。你会看到,在训练过程中,中间模型和图像将被存储在磁盘上,以便你可以跟踪训练进度。为此,将创建一个带有唯一标识符的文件夹;因此需要UNIQUE_RUN_ID。
  • 每n个批次后打印统计信息:在向网络馈送小批次数据后,将每n个批次打印一次统计信息。目前,我们将其设置为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


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


构建生成器

既然我们已经编写了一些准备代码,现在是时候构建实际的生成器了!与我们今天将要创建的、基本上遵循原始GAN的深度卷积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)


你可以看到,这是一个常规的PyTorch nn.Module类,因此只需将数据馈送到模型中即可完成前向传播,该模型在self.layers中指定为一个基于nn.Sequential的神经网络。在我们的案例中,你将编写四个上采样块。中间块由一个nn.Linear(或全连接)层、一个用于批归一化的BatchNorm1d层以及Leaky ReLU组成。由于批归一化层会使其失效,因此将偏置设置为False。


最后一个上采样层将已经具有512个神经元的中间数量转换为GENERATOR_OUTPUT_IMAGE_SHAPE,即28 * 28 * 1 = 784。通过Tanh函数,输出被归一化到[-1, 1]范围内。


构建判别器

判别器甚至比生成器还要简单。它是一个独立的神经网络,正如你通过其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)


整合所有内容

好了,我们现在有了两个不同的神经网络、一些导入和一些配置变量。是时候把它们整合在一起了!让我们从编写一些管理函数开始。


管理函数

回想一下,我们之前说过中间模型会被保存在一个文件夹中,并且还会生成图像。虽然我们实际上会在后面实现这些调用(即使用它们),但你现在就要编写它们。我们的管理函数包含五个定义:

  1. 获取设备。回想一下,你为TRAIN_ON_GPU指定了True或False。这个定义将检查你是否想使用GPU以及它是否可用,否则将指示PyTorch使用你的CPU。
  2. 创建运行目录,该函数利用UNIQUE_RUN_ID为唯一运行生成一个目录。
  3. 生成图像,该函数将使用某个生成器(通常是你最近训练过的生成器)生成16个示例,并将它们存储到磁盘上。
  4. 保存模型,该函数将生成器和判别器的当前状态保存到磁盘上。
  5. 打印训练进度,该函数将在屏幕上打印当前的损失值。


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))


准备数据集

好的,在完成了管理函数之后,接下来是编写准备数据集的功能。这将是一个多阶段的过程。首先,我们从torchvision中加载MNIST数据集。在加载时,样本将被转换为Tensor格式,并归一化到[-1, 1]范围内,以便它们与生成器生成的图像直接兼容。


然而,在加载了所有数据之后,我们仍然需要对它们进行分批处理——回想一下,你不会一次性将所有图像馈送到网络中,而是会以分批的方式进行。你还需要对图像进行打乱。为了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,用于生成具有noise_dimension维数的number_of_images个噪声向量,并将它们放到你之前配置的设备上。


在每个训练步骤的开始时,必须有效地将梯度清零,这将通过调用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包括多个训练周期,而每个训练周期又包括多个训练步骤。既然你已经为单个训练步骤编写了一些代码,那么现在是时候为执行一个训练周期编写代码了。如下所示,你将遍历由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.py。


你应该会看到模型即使在CPU上也开始相对快速地迭代。


在最初的几个训练周期中,当我们打开为此训练运行创建的文件夹中的文件时,会看到从随机噪声迅速改进为略微可识别的数字。


14

15

16

17

18

19

20

21


在随后的训练周期中,随着越来越多的噪声消失,输出开始逐渐改善:


22

23

24

25

26

27

28

29



文章来源:https://medium.com/@francescofranco_39234/building-a-vanilla-gan-with-pytorch-ffdf26275b70
欢迎关注ATYUN官方公众号
商务合作及内容投稿请联系邮箱:bd@atyun.com
评论 登录
热门职位
Maluuba
20000~40000/月
Cisco
25000~30000/月 深圳市
PilotAILabs
30000~60000/年 深圳市
写评论取消
回复取消