使用生成对抗网络 (GAN) 可以执行生成机器学习。换句话说,你可以确保模型学会生成新数据,例如图像。
像这样:
在今天的文章中,你将创建一个简单的 GAN,也称为vanilla GAN。它类似于 Goodfellow 等人 (2014) 首次创建的生成对抗网络。
什么是 GAN?
在开始构建简单的 GAN 之前,最好先简单回顾一下 GAN 是什么。让我们来看看 GAN 的通用架构:
我们可以看到,GAN 由两个独立的模型组成。第一个模型称为生成器,它学习将噪声样本(通常取自标准正态分布)转换为假图像。然后,该图像被输入到鉴别器,鉴别器判断图像是假的还是真的。利用从这一判断中得出的损失,网络被联合优化,之后该过程再次开始。
你也可以将这个过程与造假者和警察的过程进行比较。生成器充当造假者,而警察的任务是抓捕他们。当警察抓到更多的假图像时,造假者必须学会制作更好的结果。这正是发生的事情:通过判别器在判断图像是假的还是真的方面变得越来越好,生成器最终在生成假图像方面也变得越来越好。因此,生成器在经过训练后可以单独用于生成图像。
现在,是时候开始构建 GAN 了。
使用 PyTorch 的简单 GAN(完整解释的代码示例)
导入依赖项
当你想要运行要创建的代码时,你需要确保一些依赖项已安装到你的环境中。这些依赖项如下:
现在,创建一个 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。
配置变量
既然你已经指定了导入的模块,现在是时候确定在整个训练过程中将使用的可配置变量了。以下是你将创建的内容以及为什么需要它们:
# 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)
整合所有内容
好了,我们现在有了两个不同的神经网络、一些导入和一些配置变量。是时候把它们整合在一起了!让我们从编写一些管理函数开始。
管理函数
回想一下,我们之前说过中间模型会被保存在一个文件夹中,并且还会生成图像。虽然我们实际上会在后面实现这些调用(即使用它们),但你现在就要编写它们。我们的管理函数包含五个定义:
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上也开始相对快速地迭代。
在最初的几个训练周期中,当我们打开为此训练运行创建的文件夹中的文件时,会看到从随机噪声迅速改进为略微可识别的数字。
在随后的训练周期中,随着越来越多的噪声消失,输出开始逐渐改善: