PyTorch框架实战系列（4）——生成式对抗网络GAN

深度学习在图像识别问题的应用上发展较早，可以说已经很成熟，从这类问题入行的确是一个非常好的选择。

但是个人总感觉还没摸到真正人工智能的前沿，比如计算机视觉领域中的一些应用：图像转换、增强现实、图像合成、风格迁移、图像修复等，具体比如把照片转成某种油画风格，还有某短视频软件把人脸慢慢转变到老的样子，在感叹的同时不免会产生探究其原理的兴趣。

在研究PyTorch框架的过程中，找到了这一类应用的基础模型：生成式对抗网络GAN。

本例官网教程：https://pytorch.org/tutorials/beginner/dcgan_faces_tutorial.html

中文教程：http://pytorch123.com/SixthSection/Dcgan/

1、生成式对抗网络——GAN

生成式对抗网络模型于2014年首次在论文Generative Adversarial Nets中描述，GANs属于无监督学习的深度学习模型，其主要作用是让模型从训练样本中捕获特征分布，从而生成同样特征分布的数据，比如图片和声音。最初是用来生成样本做数据增强或者模拟噪声攻击的，目的都是提高图像分类模型的能力。

一个GAN包含了两个不同的网络模型，一个叫生成器（Generator），一个叫判别器（Discriminator），两者是对抗关系，所以称为对抗网络。简单来说，生成器不断生成图像，而判别器则判断输入的图像是否来自训练样本。这个过程中，判别器通过学习训练样本的特征分布，不断的在提高自己的分辨能力；生成器又在通过捕获判别器学习到的特征来反向生成具备这些特征的图像，就是说生成更像训练样本的图像，更有效的去迷惑判别器。两者相互学习，又相互对抗，最终的理想平衡点就是，判别器对来自训练样本和生成器产出的图像，判断其为真的概率都是50%。就是说，它分辨不出来了。

这让我想起了图灵关于人工智能的设想——图灵测试：当30%的测试者不能分辨回答者是人还是计算机时，就可以认为计算机具备人工智能了。

GANs的应用场景：

GAN模型从14年提出后，就被众多人工智能爱好者疯狂研究，出现了上百种不同的应用的模型。主要应用场景包括：图像生成、图像转换、场景合成、人脸合成、文本到图像的合成、风格迁移、年龄变换、图像超分辨率、图像域转换、图像修复，另外还有音乐、声音的生成合成等。

一篇GANs应用汇总博文：https://blog.csdn.net/qq_25737169/article/details/80874717

GAN zoo：https://github.com/hindupuravinash/the-gan-zoo

2、深度卷积生成式对抗网络——DCGAN

DCGAN是本例要实现的模型。

GAN最初是由一系列感知机组成的，DCGAN是对GAN的直接扩展，利用卷积概念重新构造了模型，在判别器中使用卷积层操作，在生成器中使用逆卷积操作，首先在论文Unsupervised Representation Learning With Deep Convolutional Generative Adversarial Networks中提出，下面按照该论文实现。

1、符号定义

首先对论文中的符号进行定义，后面也以这些符号表述，不再解释其意义。

D：判别器，二元分类器

x：表示输入来自训练样本的图片数据

D(x)：图片数据输入判别器，输出其是真实数据（来自训练样本）的概率

G：生成器

z：潜在空间向量，从标准正态分布采样，即高斯白噪声

G(z)：噪声输入生成器，由逆卷积操作生成一张图片

D(G(z))：生成器产生的图片输入判别器，输出其是真实数据的概率

2、数学原理

假设只有一个训练样本x，和一个假图片G(z)：判别器D的目标是让D(x)尽量大，接近于1，同时让D(G(z))尽量小，接近于0；而生成器G的目标是让D(G(z))尽量大。

为实现模型自动学习，必须定义一个合适的损失函数，即模型越接近目标，损失函数越接近0。本例DCGAN包括两个模型，如何将两者的目标归结到同一个损失函数呢？

因为D的输出是[0, 1]的概率，那么假图片的目标改为1-D(G(z))就可以实现目标同向，而且不会损失梯度。同时因为log1=0，log0是负无穷，正好可以满足损失函数与两个模型目标的一致。

所以DCGAN选用的是二进制交叉熵损失函数BCEloss

上图中，x为模型输出，y是目标标签。每个样本为一个分量，当y设为1时，分项是变为 - (logx)，当y设为0时，分项成了 - [log(1 - x)]。所有分项计算结果取平均值，就是最后的损失值。

对于本例，判别器D的单个分量就可以表示为 - [logD(x) + log(1 - D(G(z)))]，G的分量表示为 - [logD(G(z))]。这样，两个模型的目标就都归结到损失函数了。

3、获取样本数据

本例使用Celeb-A Faces数据集，可以到百度网盘下载：https://pan.baidu.com/s/1eSNpdRG#list/path=%2F

只下载Img文件夹下的img_align_celeba.zip文件就可以了。解压后放入项目的数据集文件夹中。

查看数据集情况：

from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# 数据集
root = './data'
batch_size = 64
# 训练图像大小，所有图像将使用变换器调整为此大小
image_size = 64

# 查看数据集样本
trainset = datasets.ImageFolder(root=root)
sample = next(iter(trainset))
print(sample)
# 图片转化为tensor，查看最大最小值
totensor = transforms.ToTensor()
t = totensor(sample[0])
print(t.size())
print('min', t.min().item(), 'max', t.max().item())

可以看到图片是RGB三通道，大小不定，数值范围[0, 1]。下面按照规则将图片转换为归一化的张量数据集，并构造加载器：

# 数据集转换
transform = transforms.Compose([
    transforms.Resize(image_size),
    transforms.CenterCrop(image_size),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])
# 构造数据集
trainset = datasets.ImageFolder(root=root, transform=transform)
# 查看数据集大小
print('trainset', len(trainset))
# 构造加载器
trainloader = DataLoader(dataset=trainset, batch_size=batch_size, shuffle=True)

# 加载器输出的张量
samples, labels = next(iter(trainloader))
print(samples.size(), labels.size())

然后预览一批次的图片：

# 预览图片
import cv2
import torchvision

imgs = torchvision.utils.make_grid(samples).numpy()
# 逆归一化
imgs = imgs * 0.5 + 0.5
# 通道转置到最内维度
imgs = imgs.transpose(1, 2, 0)
# RGB2BGR
imgs = cv2.cvtColor(imgs, cv2.COLOR_RGB2BGR)
cv2.imshow('win',imgs)
cv2.waitKey(0)

可以看到我们的训练样本是各种人的正面头像。

4、模型设计

1、生成器

论文中生成器的结构如图：

生成器将潜在向量z（白噪声）转换为特定大小的图片数据，转换由一系列逆卷积操作完成，最后通过tanh函数转换成模型需要的[-1, 1]数值范围的张量。超参数的选择严格按照论文给出，值得注意的是，每一个逆卷积层后都有一个批量标准化层BN，这是DCGAN论文的关键贡献，BN层有助于解决反向转播时梯度消失或爆炸的问题。

2、判别器

判别器是正常的二元分类网络，输入图像数据，输出为真实样本的概率。不同的是，卷积层后没有添加池化层，而是采用跨步卷积的方式。DCGAN论文提到使用跨步卷积而不是池化来降低采样是一种很好的做法，因为它可以让网络学习自己的池化方法。最后Sigmoid激活函数输出真实样本的概率。

3、模型参数初始化

DCGAN论文指出，模型的权重参数应按照正态分布初始化，满足均值为0，标准差为0.02。所以应在模型实例化后按照这个原则初始化权重参数。

GANs.py

import torch.nn as nn


# 生成器网络
class Generator(nn.Module):
    def __init__(self, num_z, ngf, num_channels):
        super(Generator, self).__init__()
        self.num_z = num_z
        self.save_path = './models/dcgan_netG.pth'
        # 每一层逆卷积之后都有BN批量标准化函数，是DCGAN论文的主要贡献
        self.main = nn.Sequential(
            # 输入是Z，进入卷积
            nn.ConvTranspose2d(num_z, ngf * 8, 4, 1, 0, bias=False),    # 为什么这里padding是0，其他是1？
            nn.BatchNorm2d(ngf * 8),
            nn.ReLU(True),
            # [(ngf*8), 4,  4]
            nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 4),
            nn.ReLU(True),
            # [(ngf*4), 8,  8]
            nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 2),
            nn.ReLU(True),
            # [(ngf*2), 16,  16]
            nn.ConvTranspose2d(ngf * 2, ngf, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf),
            nn.ReLU(True),
            # [(ngf), 32,  32]
            nn.ConvTranspose2d(ngf, num_channels, 4, 2, 1, bias=False),
            nn.Tanh()   # 相比sigmoid收敛更快，输出[-1, 1]
            # [(nc), 64,  64]
        )

    def forward(self, input):
        return self.main(input)


# 判别器网络
class Discriminator(nn.Module):
    def __init__(self, ndf, num_channels):
        super(Discriminator, self).__init__()
        self.save_path = './models/dcgan_netD.pth'
        self.main = nn.Sequential(
            # 输入[(nc), 64,  64]
            nn.Conv2d(num_channels, ndf, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            # [(ndf), 32,  32]
            nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 2),
            nn.LeakyReLU(0.2, inplace=True),
            # [(ndf*2), 16,  16]
            nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 4),
            nn.LeakyReLU(0.2, inplace=True),
            # [(ndf*4), 8,  8]
            nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 8),
            nn.LeakyReLU(0.2, inplace=True),
            # [(ndf*8), 4,  4]
            nn.Conv2d(ndf * 8, 1, 4, 1, 0, bias=False),
            nn.Sigmoid()    # 输出为真的概率
        )

    def forward(self, input):
        return self.main(input)

if __name__ == '__main__':
    import torch

    x = torch.rand(1, 100, 1, 1)
    model = Generator(100, 128, 3)
    y = model(x)
    print(y.size())

    x = torch.rand(1, 3, 64, 64)
    model = Discriminator(64, 3)
    y = model(x)
    print(y.size())

5、模型训练

训练部分，官网教程按照的ganhacks最佳实践来实现论文算法。我这里做了一些修改，在训练判别器时ganhacks将真实样本和生成器样本分为不同的批次训练，分别计算损失值并各自传播梯度，然后累积梯度，再执行优化；我将两步的损失函数重新构造一个损失函数，从最终损失函数只传播一次梯度，更符合论文的算法逻辑。

1、判别器训练

损失函数目标：令- [logD(x) + log(1 - D(G(z)))]接近0，或者说最大化logD(x) + log(1 - D(G(z)))。第一步，取一批次的真实样本，通过判别器，将目标标签设为1，通过BCEloss计算-logD(x)；第二步，生成器产生相同批次大小的假样本，通过判别器，设置目标标签为0，通过BCEloss计算-log(1 - D(G(z)))；第三步，两者相加为最终损失值，再反向传播梯度，执行优化。注意，这里通过detach()切断向生成器网络传播的梯度。

2、生成器训练

损失函数目标：令 - [logD(G(z))]接近0，或者说最大化 logD(G(z))。生成器产生的假样本，通过刚刚执行过优化的判别器，这次将目标标签设置为1，损失函数计算 - [logD(G(z))]，反向传播梯度，执行优化。

实现的训练代码如下：train.py

本人没有GPU训练条件，所以没写GUP部分的代码。（CPU跑真是慢……）

import torch.nn as nn
import numpy as np
import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import torch.optim as optim
import time
from datetime import timedelta

from GANs import Generator, Discriminator


# 初始化权重
def init_weights(m):
    # print(m)
    classname = m.__class__.__name__
    # 卷积层权重设置均值为0，标准差为0.02
    if classname.find('Conv') != -1:
        nn.init.normal_(m.weight.data, 0.0, 0.02)
    # BN层权重设置均值为1，标准差为0.02
    elif classname.find('BatchNorm') != -1:
        nn.init.normal_(m.weight.data, 1.0, 0.02)
        nn.init.constant_(m.bias.data, 0)

def get_time_dif(start_time):
    """获取已使用时间"""
    end_time = time.time()
    time_dif = end_time - start_time
    return timedelta(seconds=int(round(time_dif)))

# 对抗训练
def train(netG, netD, trainloader, writer):
    start_time = time.time()
    # 训练epochs数
    num_epochs = 20
    # 论文建议学习率
    lr = 0.0002
    # Adam优化器betas
    betas = (0.5, 0.999)

    # 真假标签
    real_label = 1
    fake_label = 0

    # 指定损失函数和优化器
    loss_func = nn.BCELoss()    # 二进制交叉熵函数
    optimizerD = optim.Adam(netD.parameters(), lr=lr, betas=betas)
    optimizerG = optim.Adam(netG.parameters(), lr=lr, betas=betas)

    iters = 0   # 当前总批次
    # 最佳损失值，正无穷
    best_lossD = float('inf')
    best_lossG = float('inf')
    # 记录上次目标值增加的batch数
    last_improveD = 0
    last_improveG = 0
    most_batch = 500    # 训练停止条件
    flag = False  # 记录是否可以停止

    for epoch in range(num_epochs):
        # 数据加载器中的每个batch
        for real, _ in trainloader:
            ############################
            # 训练判别器，目标为最大化log(D(x)) + log(1 - D(G(z)))
            ############################
            netD.zero_grad()
            batch_size = real.size(0)
            # 使用真实样本训练D
            real_labels = torch.full((batch_size,), real_label)
            real_outputs = netD(real).view(-1)
            # 计算真实样本的损失值
            lossD_real = loss_func(real_outputs, real_labels)
            # lossD_real.backward()
            # 真实样本为真的平均概率
            D_x = real_outputs.mean().item()

            # 使用生成器的假样本训练D
            # 生成潜在向量z，标准正态分布
            z = torch.randn(batch_size, num_z, 1, 1)
            # 生成器生成假样本
            fake = netG(z)
            fake_labels = torch.full((batch_size,), fake_label)
            # 假样本经过判别器
            fake_outputs = netD(fake.detach()).view(-1)
            # 计算假样本的损失值
            lossD_fake = loss_func(fake_outputs, fake_labels)
            # lossD_fake.backward()
            # 假样本为真的平均概率
            D_G_z1 = fake_outputs.mean().item()
            # 损失求和，并反向传播计算梯度
            lossD = lossD_real + lossD_fake
            lossD.backward()
            # netD执行参数优化
            optimizerD.step()

            ############################
            # 训练生成器，目标为最大化log(D(G(z)))
            ############################
            netG.zero_grad()
            # 标签置为真
            labels = torch.full((batch_size,), real_label)
            # 再让假样本通过判别器，判别器刚刚执行了优化
            outputs = netD(fake).view(-1)
            # 计算假样本通过判别器损失值，并反向传播梯度
            lossG = loss_func(outputs, labels)
            lossG.backward()
            # 假样本为真的平均概率
            D_G_z2 = outputs.mean().item()
            # 执行生成器优化
            optimizerG.step()

            # 查看训练的效果
            if iters % 100 == 0:
                time_dif = get_time_dif(start_time)
                msg = '[{0}/{1}][{2:>6}],  Loss_D: {3:>5.2f},  Loss_G: {4:>5.2f}, ' \
                      'D(x): {5:>5.2} | {6:>5.2} ,  D(G(z)): {7:>5.2}, Time: {10}'
                print(msg.format(epoch+1, num_epochs, iters, lossD.item(), lossG.item(),
                                 D_x, D_G_z1, D_G_z2, time_dif))
                # 可视化训练成果
                writer.add_scalar("loss/Discriminator", lossD.item(), iters)
                writer.add_scalar("loss/Generator", lossG.item(), iters)

                # 更新最佳损失值和损失降低批次
                if lossD.item() < best_lossD:
                    best_lossD = lossD.item()
                    last_improveD = iters
                if lossG.item() < best_lossG:
                    best_lossG = lossG.item()
                    last_improveG = iters

            # 保存生成器参数，可以演进生成的假样本
            if iters % 500 == 0:
                g_path = netG.save_path + str(iters)
                torch.save(netG.state_dict(), g_path)

            # 训练停止条件
            if iters - last_improveD > most_batch and iters - last_improveG > most_batch:
                print("Training Finished ...")
                torch.save(netG.state_dict(), netG.save_path)
                torch.save(netD.state_dict(), netD.save_path)
                flag = True
                break

            iters += 1
        # 停止训练
        if flag:
            break

    if not flag:
        print("Training Finished ...")
        torch.save(netG.state_dict(), netG.save_path)
        torch.save(netD.state_dict(), netD.save_path)


if __name__ == '__main__':
    # 设置随机数种子，保证每次结果一样
    seed = 999
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True

    # 数据集
    root = './data'
    batch_size = 128
    # 训练图像缩放裁剪尺寸
    image_size = 64
    # 通道数
    num_channels = 3
    # 潜在向量z维度
    num_z = 100
    # 生成器中特征数
    num_generator_feature = 64
    # 判别器中的特征数
    num_disc_feature = 64

    # 数据集转换
    transform = transforms.Compose([
        transforms.Resize(image_size),
        transforms.CenterCrop(image_size),
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
    ])
    # 构造数据集，ImageFolder取子目录的文件作为数据集
    trainset = datasets.ImageFolder(root=root, transform=transform)
    # 查看数据集大小
    print('trainset', len(trainset))
    # 构造加载器
    trainloader = DataLoader(dataset=trainset, batch_size=batch_size, shuffle=True)

    # 加载器输出的张量
    samples, _ = next(iter(trainloader))
    print(samples.size())

    # 初始化生成器
    netG = Generator(num_z, num_generator_feature, num_channels)
    # 初始化判别器
    netD = Discriminator(num_disc_feature, num_channels)
    # 初始权重，按论文给出的建议
    netG.apply(init_weights)
    netD.apply(init_weights)

    from tensorboardX import SummaryWriter

    # 记录训练指标，可视化展示
    log_path = './runs'
    with SummaryWriter(log_dir=log_path + '/' + time.strftime('%m-%d_%H.%M', time.localtime())) as writer:
        # 开始对抗训练
        train(netG, netD, trainloader, writer)

可视化观察两个模型网络的损失值变化：

CPU跑的实在太慢了，一个epoch 20多万样本，四个epoch跑了20个小时。不过总体来看，判别器的损失值迅速降到1以下，然后不断波动，而生成器的损失值在总体向下。波动都比较大，因为两者在相互对抗，对方提升，自己就会下降。判别器从开始是占优的，因为生成器前面生成的图像肯定惨不忍睹，判别器很好区分，随着生成器水平越来越高，它的分辨真实的概率会越来越低，最终趋近50%。

6、训练成果

受限于硬件条件和时间关系，训练并不充分，不过让我们看一下生成器产生图像的变化：

visual.py

import cv2
import torchvision


def visualize(samples, name):
    # 预览图片
    imgs = torchvision.utils.make_grid(samples)
    # 通道转置到最内维度
    imgs = imgs.numpy().transpose(1, 2, 0)
    # 逆归一化
    imgs = imgs * 0.5 + 0.5
    # RGB2BGR
    imgs = cv2.cvtColor(imgs, cv2.COLOR_RGB2BGR)
    cv2.imshow(name, imgs)
    cv2.waitKey(0)


if __name__ == '__main__':
    from GANs import Generator
    import os
    import torch

    netG = Generator(num_z=100, ngf=64, num_channels=3)
    for i in range(20):
        # 加载生成器模型参数
        iters = i * 500
        save_path = netG.save_path + str(iters)
        if not os.path.exists(save_path):
            break
        netG.load_state_dict(torch.load(save_path))
        # 生成器生成假图片
        z = torch.randn(64, 100, 1, 1)
        fake = netG(z).detach()
        visualize(fake, str(iters))

生成器的目的是生成类似样本的各种人的头像。

按训练批次摘取几个阶段的图片：

可以看到，从一开始的只有噪点，到后来越来越像人类，有部分图像跟真人几乎没有区别了。

 

自从了解了GANs的应用场景，持续处在兴奋状态，这才是人工智能嘛。如果再经过几个实际项目的实践，就可以稳固深度学习的科技树了。