​ 在本教程中，我们将在PyTorch中构建一个简单的DCGAN，并在手写数据集上对它进行训练。我们将讨论PyTorch DataLoader，以及如何使用它将图像数据提供给PyTorch神经网络进行训练。PyTorch是本教程的重点，所以我假设您熟悉GAN的工作方式。

要求

1. python版本为3.7或更高。
2. PyTorch 1.5不知道如何安装? 可以参考github项目https://github.com/zergtant/pytorch-handbook
3. Matplotlib 3.1或更高版本
4. 22分钟左右的时间。

目前的任务

创建一个函数G: Z→X, Z ~ N₁₆(0, 1)和X ~ MNIST。

​ 也就是说，训练一个GAN，让它接收16维随机噪声，并生成看起来像来自MNIST数据集的真实样本的图像。

来自MNIST数据集的数字示例(来源:Josef Steppan)

但在我们开始之前…

​ 我们来整理一下。首先请安装所需的Python版本和上述库。然后，创建您的项目目录。我这里是DCGAN。在该目录中，创建一个名为data的目录。然后，定位到这个GitHub目录(https://github.com/myleott/mnist_png)并下载mnist_png.tar.gz。这个压缩文件包含MNIST数据集，为70000个单独的png文件。当然，我们可以使用PyTorch内置的MNIST数据集，但这样您就不能了解如何加载具体的图像数据进行训练。解压缩文件并将mnist_png目录放入数据目录中。你的项目目录应该是这样的:

我们的项目目录，包括图像文件和Python脚本。0/、1/等中的数千个图像文件没有显示。

​ 最后，在dcgan_mni .py脚本中添加以下内容:

import osimport torch
from torch import nn
from torch import optim
import torchvision as tv
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader

​ 好了，现在我们可以开始了。

生成器

​ 将以下内容添加到您的dcgan_mni .py脚本:

class Generator(nn.Module):
    def __init__(self, latent_dim=100, batchnorm=True):
        """A generator for mapping a latent space to a sample space.
        The sample space for this generator is single-channel, 28x28 images
        with pixel intensity ranging from -1 to +1.
        Args:
            latent_dim (int): latent dimension ("noise vector")
            batchnorm (bool): Whether or not to use batch normalization
        """
        super(Generator, self).__init__()
        self.latent_dim = latent_dim
        self.batchnorm = batchnorm
        self._init_modules()

    def _init_modules(self):
        """Initialize the modules."""
        # Project the input
        self.linear1 = nn.Linear(self.latent_dim, 256*7*7, bias=False)
        self.bn1d1 = nn.BatchNorm1d(256*7*7) if self.batchnorm else None
        self.leaky_relu = nn.LeakyReLU()

        # Convolutions
        self.conv1 = nn.Conv2d(
                in_channels=256,
                out_channels=128,
                kernel_size=5,
                stride=1,
                padding=2,
                bias=False)
        self.bn2d1 = nn.BatchNorm2d(128) if self.batchnorm else None

        self.conv2 = nn.ConvTranspose2d(
                in_channels=128,
                out_channels=64,
                kernel_size=4,
                stride=2,
                padding=1,
                bias=False)
        self.bn2d2 = nn.BatchNorm2d(64) if self.batchnorm else None

        self.conv3 = nn.ConvTranspose2d(
                in_channels=64,
                out_channels=1,
                kernel_size=4,
                stride=2,
                padding=1,
                bias=False)
        self.tanh = nn.Tanh()

    def forward(self, input_tensor):
        """Forward pass; map latent vectors to samples."""
        intermediate = self.linear1(input_tensor)
        intermediate = self.bn1d1(intermediate)
        intermediate = self.leaky_relu(intermediate)

        intermediate = intermediate.view((-1, 256, 7, 7))

        intermediate = self.conv1(intermediate)
        if self.batchnorm:
            intermediate = self.bn2d1(intermediate)
        intermediate = self.leaky_relu(intermediate)

        intermediate = self.conv2(intermediate)
        if self.batchnorm:
            intermediate = self.bn2d2(intermediate)
        intermediate = self.leaky_relu(intermediate)

        intermediate = self.conv3(intermediate)
        output_tensor = self.tanh(intermediate)
        return output_tensor

​ 生成器继承nn.Module，它是PyTorch神经网络的基类。生成器有三种方法:

Generator.init

​ 构造函数，它存储实例变量并调用_init_layers。这里没什么可说的。

Generator._init_modules

​ 这个方法实例化PyTorch模块(或在其他框架中调用的“层”)。这些包括:

• 一个线性(“完全连接”)模块,将向量空间映射到一个7×7×256 = 1254维空间。我们将看到,这个12554长度张量被重新塑造为a(256,7,7)的“图像”张量(通道×高×宽)。在pytorch中,通道在空间维度之前。
• 一个一维的指定的的批处理模块。
• ReLU模块。
• 一个二维的卷积层。
• 两个二维反卷积层;这用于放大图像。请注意一个卷积层的外通道是如何成为下一个卷积层的内通道的。
• 两个二维批归一化层。
• 一个Tanh模块作为输出激活。我们将重新标定图像到范围[-1,1]，所以我们的生成器输出激活应该反映这一点。

​ 这些可以在剩余的__init__方法中实例化，但是我喜欢将模块实例化与构造函数分开。对于如此简单的模型来说，这是微不足道的，但是当模型变得越来越复杂时，这有助于保持代码的简单。

Generator.forward

​ 这就是我们的生成器从随机噪声中生成样本的方法。输入张量被传递给第一个模块，输出被传递给下一个模块，输出被传递给下一个模块，以此类推。这是相当直接的，但我想提请您注意两个有趣的特点:

1. 注意这一行intermediate = intermediate.view((-1, 256, 7, 7))。与Keras不同，PyTorch不使用显式的“重塑”模块;相反，我们使用PyTorch操作“视图”手动重塑张量。其他简单的PyTorch操作也可以在前传过程中应用，比如将一个张量乘以2,PyTorch不会眨眼睛。
2. 注意forward方法中的if语句。PyTorch使用一种逐行定义策略，这意味着在向前传球期间动态构建计算图。这使得PyTorch极其灵活;没有什么可以阻止您向向前传递添加循环，或者随机选择要使用的几个模块中的一个。

鉴别器

​ 将以下内容添加到您的dcgan_mni .py脚本:

class Discriminator(nn.Module):
    def __init__(self):
        """A discriminator for discerning real from generated images.
        Images must be single-channel and 28x28 pixels.
        Output activation is Sigmoid.
        """
        super(Discriminator, self).__init__()
        self._init_modules()

    def _init_modules(self):
        """Initialize the modules."""
        self.conv1 = nn.Conv2d(
                in_channels=1,
                out_channels=64,
                kernel_size=5,
                stride=2,
                padding=2,
                bias=True)
        self.leaky_relu = nn.LeakyReLU()
        self.dropout_2d = nn.Dropout2d(0.3)

        self.conv2 = nn.Conv2d(
                in_channels=64,
                out_channels=128,
                kernel_size=5,
                stride=2,
                padding=2,
                bias=True)

        self.linear1 = nn.Linear(128*7*7, 1, bias=True)
        self.sigmoid = nn.Sigmoid()

    def forward(self, input_tensor):
        """Forward pass; map samples to confidence they are real [0, 1]."""
        intermediate = self.conv1(input_tensor)
        intermediate = self.leaky_relu(intermediate)
        intermediate = self.dropout_2d(intermediate)

        intermediate = self.conv2(intermediate)
        intermediate = self.leaky_relu(intermediate)
        intermediate = self.dropout_2d(intermediate)

        intermediate = intermediate.view((-1, 128*7*7))
        intermediate = self.linear1(intermediate)
        output_tensor = self.sigmoid(intermediate)

        return output_tensor

​ 我不会详细介绍这点，因为它非常类似于生成器。但仔细阅读它，确保你理解它的作用。

DCGAN

​ 将以下内容添加到您的dcgan_mni .py脚本:

class DCGAN():
    def __init__(self, latent_dim, noise_fn, dataloader,
                 batch_size=32, device='cpu', lr_d=1e-3, lr_g=2e-4):
        """A very basic DCGAN class for generating MNIST digits
        Args:
            generator: a Ganerator network
            discriminator: A Discriminator network
            noise_fn: function f(num: int) -> pytorch tensor, (latent vectors)
            dataloader: a pytorch dataloader for loading images
            batch_size: training batch size. Must match that of dataloader
            device: cpu or CUDA
            lr_d: learning rate for the discriminator
            lr_g: learning rate for the generator
        """
        self.generator = Generator(latent_dim).to(device)
        self.discriminator = Discriminator().to(device)
        self.noise_fn = noise_fn
        self.dataloader = dataloader
        self.batch_size = batch_size
        self.device = device
        self.criterion = nn.BCELoss()
        self.optim_d = optim.Adam(self.discriminator.parameters(),
                                  lr=lr_d, betas=(0.5, 0.999))
        self.optim_g = optim.Adam(self.generator.parameters(),
                                  lr=lr_g, betas=(0.5, 0.999))
        self.target_ones = torch.ones((batch_size, 1), device=device)
        self.target_zeros = torch.zeros((batch_size, 1), device=device)

    def generate_samples(self, latent_vec=None, num=None):
        """Sample images from the generator.
        Images are returned as a 4D tensor of values between -1 and 1.
        Dimensions are (number, channels, height, width). Returns the tensor
        on cpu.
        Args:
            latent_vec: A pytorch latent vector or None
            num: The number of samples to generate if latent_vec is None
        If latent_vec and num are None then use self.batch_size
        random latent vectors.
        """
        num = self.batch_size if num is None else num
        latent_vec = self.noise_fn(num) if latent_vec is None else latent_vec
        with torch.no_grad():
            samples = self.generator(latent_vec)
        samples = samples.cpu()  # move images to cpu
        return samples

    def train_step_generator(self):
        """Train the generator one step and return the loss."""
        self.generator.zero_grad()

        latent_vec = self.noise_fn(self.batch_size)
        generated = self.generator(latent_vec)
        classifications = self.discriminator(generated)
        loss = self.criterion(classifications, self.target_ones)
        loss.backward()
        self.optim_g.step()
        return loss.item()

    def train_step_discriminator(self, real_samples):
        """Train the discriminator one step and return the losses."""
        self.discriminator.zero_grad()

        # real samples
        pred_real = self.discriminator(real_samples)
        loss_real = self.criterion(pred_real, self.target_ones)

        # generated samples
        latent_vec = self.noise_fn(self.batch_size)
        with torch.no_grad():
            fake_samples = self.generator(latent_vec)
        pred_fake = self.discriminator(fake_samples)
        loss_fake = self.criterion(pred_fake, self.target_zeros)

        # combine
        loss = (loss_real + loss_fake) / 2
        loss.backward()
        self.optim_d.step()
        return loss_real.item(), loss_fake.item()

    def train_epoch(self, print_frequency=10, max_steps=0):
        """Train both networks for one epoch and return the losses.
        Args:
            print_frequency (int): print stats every `print_frequency` steps.
            max_steps (int): End epoch after `max_steps` steps, or set to 0
                             to do the full epoch.
        """
        loss_g_running, loss_d_real_running, loss_d_fake_running = 0, 0, 0
        for batch, (real_samples, _) in enumerate(self.dataloader):
            real_samples = real_samples.to(self.device)
            ldr_, ldf_ = self.train_step_discriminator(real_samples)
            loss_d_real_running += ldr_
            loss_d_fake_running += ldf_
            loss_g_running += self.train_step_generator()
            if print_frequency and (batch+1) % print_frequency == 0:
                print(f"{batch+1}/{len(self.dataloader)}:"
                      f" G={loss_g_running / (batch+1):.3f},"
                      f" Dr={loss_d_real_running / (batch+1):.3f},"
                      f" Df={loss_d_fake_running / (batch+1):.3f}",
                      end='\r',
                      flush=True)
            if max_steps and batch == max_steps:
                break
        if print_frequency:
            print()
        loss_g_running /= batch
        loss_d_real_running /= batch
        loss_d_fake_running /= batch
        return (loss_g_running, (loss_d_real_running, loss_d_fake_running))
view raw
dcgan_mnist_gan.py hosted with  by GitHub

DCGAN.init

​ 让我们逐行查看构造函数:

self.generator = Generator(latent_dim).to(device)
self.discriminator = Discriminator().to(device)

​ 构造函数的前两行(非docstring)实例化生成器和Discriminator，将它们移动到指定的地方，并将它们存储为实例变量。通常是“cpu”，或者“cuda”，如果你想使用gpu。

self.noise_fn = noise_fn

​ 接下来，我们将noise_fn存储为一个实例变量;noise_fn函数以整数作为输入，并以PyTorch张量的形式返回num潜在向量作为输出，带有(num, latent_dim)的形状。这个PyTorch张量必须在指定的设备上。

self.dataloader = dataloader

​ 我们存储dataloader，一个torch.utils.data.DataLoader对象作为实例变量;稍后将对此进行更多介绍。

self.batch_size = batch_size
self.device = device

​ 存储为实例变量。

self.criterion = nn.BCELoss()
self.optim_d = optim.Adam(self.discriminator.parameters(), lr=lr_d, betas=(0.5, 0.999))
self.optim_g = optim.Adam(self.generator.parameters(), lr=lr_g, betas=(0.5, 0.999))

​ 将损失函数设置为交叉熵，并实例化生成器和鉴别器的Adam优化器。pytorch的优化器需要知道他们在优化什么。对于鉴别器，这意味着鉴别器网络中的所有可训练参数。因为我们的Discriminator类继承自nn.Module中，它有parameters()方法，该方法返回所有实例变量中的所有可训练参数，这些实例变量也是PyTorch模块。生成器也是一样。

self.target_ones = torch.ones((batch_size, 1), device=device)
self.target_zeros = torch.zeros((batch_size, 1), device=device)

​ 为训练的目标，设置为指定的设备。记住，鉴别器试图将真实样本分类为1，将生成样本分类为0，而生成器试图让鉴别器将生成样本错误分类为1。我们在这里定义并存储它们，这样我们就不必在每个训练步骤中重新创建它们。

DCGAN.generate_samples

​ 用于生成示例的辅助方法。注意，这里使用了no_grad上下文管理器，它告诉PyTorch不要跟踪梯度，因为这个方法不用于训练网络。还要注意的是，无论指定的设备是什么，返回的张量都被设置为cpu，这对于进一步的使用是必要的，比如显示样本或将它们保存到磁盘上。

DCGAN.train_step_generator

​ 此方法执行生成器的一个epoch，并以浮点数的形式返回损失。让我们一步步来看看:

self.generator.zero_grad()

​ 清除生成器的梯度是必要的。因为PyTorch会自动跟踪梯度和计算网络。而我们现在不需要这些。

latent_vec = self.noise_fn(self.batch_size)
generated = self.generator(latent_vec)
classifications = self.discriminator(generated)
loss = self.criterion(classifications, self.target_ones)

​ 获取一批潜在向量，利用它们生成样本，判别每个样本的真实程度，然后利用交叉熵计算损失。注意，通过将这些网络链接在一起，我们创建了一个单一的计算图，从潜在向量开始，包括生成器和鉴别器网络，并以损失结束。

loss.backward()
self.optim_g.step()

​ PyTorch的主要优点之一是它可以自动跟踪计算图形及其梯度。通过loss调用反向传播，PyTorch应用反向传播并计算损失相对于计算图中的每个参数的梯度。通过调用生成器中优化器的step方法，生成器的参数(只有生成器的参数)将略微向其梯度的负方向移动。

return loss.item()

​ 最后，我们得到损失。使用item方法很重要，这样我们将返回一个浮点数而不是一个PyTorch张量。如果我们返回了张量，Python垃圾收集器将无法清理底层的计算图，我们将很快耗尽内存。

DCGAN.train_step_discriminator

​ 这个方法与train_step_generator非常相似，但是有两个显著的区别。第一:

with torch.no_grad():
    fake_samples = self.generator(latent_vec)

​ 这里使用上下文管理器no_grad来告诉PyTorch不要跟踪梯度。这不是必须的，但减少了不必要的计算。第二:

loss = (loss_real + loss_fake) / 2

​ 这条线真的很酷。loss_real为真实样本的鉴别器损失(附加其计算图)，loss_fake为虚假样本的鉴别器损失(及计算图)。PyTorch能够使用+运算符将这些图形组合成一个计算图形。然后我们将反向传播和参数更新应用到组合计算图。如果您不认为这是简单的，试着在另一个框架中重写它。

DCGAN.train_epoch

​ 这个函数进行一次训练生成器和鉴别器的epoch，也就是在整个数据集上进行一次遍历。我们绕一圈后再回到这个问题上。

main

​ 添加以下代码到您的脚本:

def main():
    import matplotlib.pyplot as plt
    from time import time
    batch_size = 32
    epochs = 100
    latent_dim = 16
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    transform = tv.transforms.Compose([
            tv.transforms.Grayscale(num_output_channels=1),
            tv.transforms.ToTensor(),
            tv.transforms.Normalize((0.5,), (0.5,))
            ])
    dataset = ImageFolder(
            root=os.path.join("data", "mnist_png", "training"),
            transform=transform
            )
    dataloader = DataLoader(dataset,
            batch_size=batch_size,
            shuffle=True,
            num_workers=2
            )
    noise_fn = lambda x: torch.randn((x, latent_dim), device=device)
    gan = DCGAN(latent_dim, noise_fn, dataloader, device=device, batch_size=batch_size)
    start = time()
    for i in range(10):
        print(f"Epoch {i+1}; Elapsed time = {int(time() - start)}s")
        gan.train_epoch()
    images = gan.generate_samples() * -1
    ims = tv.utils.make_grid(images, normalize=True)
    plt.imshow(ims.numpy().transpose((1,2,0)))
    plt.show()


if __name__ == "__main__":
    main()
view raw
dcgan_mnist_main.py hosted with  by GitHub

​ 该函数构建、训练和展示GAN。

import matplotlib.pyplot as plt
from time import time
batch_size = 32
epochs = 100
latent_dim = 16

​ 导入pyplot(用于可视化)和time(用于为训练计时)。将训练批处理大小设置为32,epoch数设置为100，隐藏层维度设置为16。

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

​ 这一行检查是否有cuda设备可用。如果有，则分配该设备;否则，分配cpu。

transform = tv.transforms.Compose([
            tv.transforms.Grayscale(num_output_channels=1),
            tv.transforms.ToTensor(),
            tv.transforms.Normalize((0.5,), (0.5,))
            ])

​ dataloader使用这种复合变换对图像进行预处理。我们之前下载的MNIST数据集是.png文件;当PyTorch从磁盘加载它们时，必须对它们进行处理，以便我们的神经网络能够正确地使用它们。变换的顺序是:

• Grayscale(num_output_channels=1):将图像转换为灰度图。加载时，MNIST数字为RGB格式，有三个通道。Greyscale将这三种减少为一种。
• ToTensor():将图像转换为点tensor张量，其尺寸通道×高度×宽度。这也将重新调整像素值，从0到255之间的整数到0.0到1.0之间的浮点值。
• Normalize((0.5，)，(0.5，)):将像素值从范围[0.0,1.0]缩放到[-1.0,1.0]。第一个参数是所属，第二个参数是使用量，应用于每个像素的函数为:

​ 因为这个转换是对每个通道应用的，所以它是一个元组。RGB图像的等效变换是

Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))

dataset = ImageFolder(
            root=os.path.join("data", "mnist_png", "training"),
            transform=transform
            )

​ 在这里，我们通过指定数据集的目录和要应用的转换来创建数据集。用于创建DataLoader:

dataloader = DataLoader(dataset,
            batch_size=batch_size,
            shuffle=True,
            num_workers=2
            )

​ DataLoader是一个对象，它从数据集加载数据。这里我们指定批量大小,告诉dataloader打乱每个epoch之间的数据集,并使用两个工作进程(如果您使用的是Windows,这将导致问题,可以将num_workers设置为0),遍历这个dataloader和每次迭代它将返回一个元组包含:

1. 对应一批(32个样本)灰度(1通道)MNIST图像(28×28像素)的形状(32,1,28,28)的PyTorch张量。
2. 从0到9的形状(32，)的PyTorch张量，对应于该图像的标号(digit)。这些类标签是从目录结构中获取的，因为所有的0都在目录0中，所有的1都在目录1中，等等。

noise_fn = lambda x: torch.rand((x, latent_dim), device=device)

​ 用于产生随机、正态分布噪声的函数。

gan = DCGAN(latent_dim, noise_fn, dataloader, device=device)
start = time()
for i in range(10):
    print(f"Epoch {i+1}; Elapsed time = {int(time() - start)}s")
    gan.train_epoch()

​ 建立和训练GAN。

DCGAN.train_epoch, again:

​ 既然我们已经讨论了什么是DataLoader，让我们再来看看这个。方法虽然冗长，但我想重点在两行:

for batch, (real_samples, _) in enumerate(self.dataloader):
   real_samples = real_samples.to(self.device)

​ 在这里，我们遍历dataloader。我们将dataloader包装在迭代器中，这样我们就可以跟踪编号，但是正如您所看到的，dataloader确实按照承诺返回了一个元组。我们将这批图像张量分配给real_samples，并忽略标签，因为我们不需要它们。然后，在循环中，我们将real_samples移动到指定的网络。重要的是模型的输入和模型本身在同一个设备上;如果你忘记了，不要担心，PyTorch一定会让你知道的!另外，不要担心dataloader“快用完了”。一旦我们遍历了整个数据集，循环将结束，但如果我们尝试再次遍历它，它将从开始开始(首先移动图像，因为我们在创建dataloader时指定了这一点)。

让我们试着运行一下?

​ 如果复制和粘贴正确，运行脚本应该会显示几分钟的训练统计数据，然后是一些生成的数字。希望它是这样的:

​ 如果它们看起来很糟糕，试着再运行一次(GANs是出了名的不稳定)。如果它仍然不行，在下面加一个提示，我们将看看我们是否不能调用它。

​ 为了乐趣，我修改了这个脚本，看看生成器在每10个epoch之后能够做什么。以下是结果。

​ 我认为这对于1000个epoch来说已经很不错了。以下是那些训练步骤的损失，分为10个“阶段”。

结论

• 本教程中描述的DCGAN显然非常简单，但它应该足以让您开始在PyTorch中实现更复杂的GANs。
• 在我做一个关于GAN的教程之前，你能修改这个脚本来制作一个条件GAN吗?
• 完整的脚本可以在这里找到。https://github.com/ConorLazarou/pytorch-generative-models/blob/master/GAN/DCGAN/dcgan_mnist.py

作者：Conor Lazarou

deephub翻译组：孟翔杰