写给程序员的机器学习入门 (十四) - 对抗生成网络 如何造假脸

这篇文章将会教你怎样用机器学习来伪造假数据,题材还是人脸,以下六张人脸里面,有两张是假的,猜猜是哪两张😎?

01

生成假人脸使用的网络是对抗生成网络 (GAN - Generative adversarial network),这个网络与之前介绍的比起来相当特殊,虽然看起来不算复杂,但训练起来极其困难,以下将从基础原理开始一直讲到具体代码,还会引入一些之前没有讲过的组件和训练方法😨。

对抗生成网络 (GAN) 的原理

所谓生成网络就是用于生成文章,音频,图片,甚至代码等数据的机器学习模型,例如我们可以给出一个需求让网络生成一份代码,如果网络足够强大,生成的代码质量足够好并且能满足需求,那码农们就要面临失业了😱。当然,目前机器学习模型可以生成的数据比较有限并且质量都很一般,码农们的饭碗还是能保住一段时间的。

生成网络和普通的模型一样,要求有输入和输出,假设我们可以传入一些条件让网络生成符合条件的图片:

02

看起来非常好用,但训练这样的模型需要一个庞大的数据集,并且得一张张图片去标记它们的属性,实现起来会累死人。这篇文章介绍的对抗生成网络属于无监督学习,可以完全不需要给数据打标签,你只需要给模型认识一些真实数据,就可以让模型输出类似真实数据的假数据。对抗生成网络分为两部分,第一部分是生成器 (Generator),第二部分是识别器 (Discriminator),生成器负责根据随机条件生成数据,识别器负责识别数据是否为真。

03

训练对抗生成网络有两大目标,这两大目标是矛盾的,这就是为什么我们叫对抗生成网络:

  • 生成器需要生成骗过识别器 (输出为真) 的数据
  • 识别器需要不被生成器骗过去 (针对生成器生成的数据输出为假,针对真实数据输出为真)

对抗生成网络的训练流程大致如下,需要循环训练生成器和识别器:

04

简单通俗一点我们可以用造假皮包为例来理解,好理解了吧🤗:

05

和现实造假皮包一样,生成器会生成越来越接近真实数据的假数据,最后会生成和真实数据一模一样的数据,但这样反而就远离我们构建生成网络的目的了(不如直接用真实数据)。使用生成网络通常是为了达到以下的目的:

  • 要求大量看上去是真的,但稍微不一样的数据
  • 要求没有版权保护的数据 (假数据没来的版权🤒)
  • 生成想要但是现实没有的数据 (需要更进一步的工作)

看以上的流程你可能会发现,因为对抗生成网络是无监督学习,不需要标签,我们只能给模型传入随机的条件来让它生成数据,模型生成出来的数据看起来可能像真的但不一定是我们想要的。如果我们想要指定具体的条件,则需要在训练完成以后分析随机条件对生成结果的影响,例如随机生成的第二个数字代表性别,第六个数字代表年龄,第八个数字代表头发的数量,这样我们就可以调整这些条件来让模型生成想要的图片。

还记得上一篇人脸识别的模型不?人脸识别的模型会把图片转换为某个长度的向量,训练完成以后这个向量的值会代表人物的属性,而这一篇是反过来,把某个长度的向量转换回图片,训练成功以后这个向量同样会代表人物的各个属性。当然,两种的向量表现形式是不同的,把人脸识别输出的向量交给对抗生成网络,生成的图片和原有的图片可能会相差很远,把人脸识别输出的向量还原回去的方法后面再研究吧🤕。

对抗生成网络的实现

反卷积层 (ConvTranspose2d)

第八篇介绍 CNN 的文章中,我们了解过卷积层运算 (Conv2d) 的实现原理,CNN 模型会利用卷积层来把图片的长宽逐渐缩小,通道数逐渐扩大,最后扁平化输出一个代表图片特征的向量:

06

07

而在对抗生成网络的生成器中,我们需要实现反向的操作,即把向量当作一个 (向量长度, 1, 1) 的图片,然后把长宽逐渐扩大,通道数 (最开始是向量长度) 逐渐缩小,最后变为 (3, 图片长度, 图片宽度) 的图片 (3 代表 RGB)。

实现反向操作需要反卷积层 (ConvTranspose2d),反卷积层简单的来说就是在参数数量相同的情况下,把输出大小的数据还原为输入大小的数据:

08

要理解反卷积层的具体运算方式,我们可以把卷积层拆解为简单的矩阵乘法:

09

可以看到卷积层计算的时候可以根据内核参数和输入大小生成一个矩阵,然后计算输入与这个矩阵的乘积来得到输出结果。

而反卷积层则会计算输入与转置 (Transpose) 后的矩阵的乘积得到输出结果:

10

可以看到卷积层与反卷积层的区别只在于是否转置计算使用的矩阵。此外,通道数量转换的计算方式也是一样的。

测试反卷积层的代码如下:

>>> import torch

# 生成测试用的矩阵
# 第一个维度代表批次,第二个维度代表通道数量,第三个维度代表长度,第四个维度代表宽度
>>> a = torch.arange(1, 5).float().reshape(1, 1, 2, 2)
>>> a
tensor([[[[1., 2.],
          [3., 4.]]]])

# 创建反卷积层
>>> convtranspose2d = torch.nn.ConvTranspose2d(1, 1, kernel_size=2, stride=2, bias=False)

# 手动指定权重 (让计算更好理解)
>>> convtranspose2d.weight = torch.nn.Parameter(torch.tensor([0.1, 0.2, 0.5, 0.8]).reshape(1, 1, 2, 2))
>>> convtranspose2d.weight
Parameter containing:
tensor([[[[0.1000, 0.2000],
          [0.5000, 0.8000]]]], requires_grad=True)

# 测试反卷积层
>>> convtranspose2d(a)
tensor([[[[0.1000, 0.2000, 0.2000, 0.4000],
          [0.5000, 0.8000, 1.0000, 1.6000],
          [0.3000, 0.6000, 0.4000, 0.8000],
          [1.5000, 2.4000, 2.0000, 3.2000]]]],
       grad_fn=<SlowConvTranspose2DBackward>)

需要注意的是,不一定存在一个反卷积层可以把卷积层的输出还原到输入,这是因为卷积层的计算是不可逆的,即使存在一个可以把输出还原到输入的矩阵,这个矩阵也不一定有一个等效的反卷积层的内核参数。

生成器的实现 (Generator)

接下来我们看一下生成器的定义,原始介绍 GAN 的论文给出了生成 64x64 图片的网络,而这里给出的是生成 80x80 图片的网络,其实区别只在于一开始的输出通道数量 (论文是 4, 这里是 5)

class GenerationModel(nn.Module):
    """生成虚假数据的模型"""
    # 编码长度
    EmbeddedSize = 128

    def __init__(self):
        super().__init__()
        self.generator = nn.Sequential(
            # 128,1,1 => 512,5,5
            nn.ConvTranspose2d(128, 512, kernel_size=5, stride=1, padding=0, bias=False),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            # => 256,10,10
            nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            # => 128,20,20
            nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            # => 64,40,40
            nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            # => 3,80,80
            nn.ConvTranspose2d(64, 3, kernel_size=4, stride=2, padding=1, bias=False),
            # 限制输出在 -1 ~ 1,不使用 Hardtanh 是为了让超过范围的值可以传播给上层
            nn.Tanh())

    def forward(self, x):
        y = self.generator(x.view(x.shape[0], x.shape[1], 1, 1))
        return y

表现如下:

11

其中批次正规化 (BatchNorm) 用于控制参数值范围,防止层数过多 (后面会结合识别器训练) 导致梯度爆炸问题。

还有一个要点是生成器输出的范围会在 -1 ~ 1,也就是使用 -1 ~ 1 代表 0 ~ 255 的颜色值,这跟我们之前处理图片的时候把值除以 255 使得范围在 0 ~ 1 不一样。使用 -1 ~ 1 可以提升输出颜色的精度 (减少浮点数的精度损失)。

识别器的实现 (Discriminator)

我们再看以下识别器的定义,基本上就是前面生成器的相反流程:

class DiscriminationModel(nn.Module):
    """识别数据是否真实的模型"""

    def __init__(self):
        super().__init__()
        self.discriminator = nn.Sequential(
            # 3,80,80 => 64,40,40
            nn.Conv2d(3, 64, kernel_size=4, stride=2, padding=1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            # => 128,20,20
            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.2, inplace=True),
            # => 256,10,10
            nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(256),
            nn.LeakyReLU(0.2, inplace=True),
            # => 512,5,5
            nn.Conv2d(256, 512, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(512),
            nn.LeakyReLU(0.2, inplace=True),
            # => 1,1,1
            nn.Conv2d(512, 1, kernel_size=5, stride=1, padding=0, bias=False),
            # 扁平化
            nn.Flatten(),
            # 输出是否真实数据 (0 or 1)
            nn.Sigmoid())

    def forward(self, x):
        y = self.discriminator(x)
        return y

表现如下:

12

看到这里你可能会有几个疑问:

  • 为什么用 LeakyReLU: 这是为了防止层数叠加次数过多导致的梯度消失问题,参考第三篇,LeakyReLU 对于负数输入不会返回 0,而是返回 输入 * slope,这里的 slope 指定为 0.2
  • 为什么第一层不加批次正规化 (BatchNorm): 原有论文中提到实际测试中,如果在所有层添加批次正规化会让模型训练结果不稳定,生成器的最后一层和识别器的第一层拿掉以后效果会好一些
  • 为什么不加池化层: 添加池化层以后可逆性将会降低,例如识别器针对假数据返回接近 0 的数值时,判断哪些部分导致这个输出的依据会减少

训练生成器和识别器的方法

接下来就是训练生成器和识别器,生成器和识别器需要分别训练,训练识别器的时候不能动生成器的参数,训练生成器的时候不能动识别器的参数,使用的代码大致如下:

# 创建模型实例
generation_model = GenerationModel().to(device)
discrimination_model = DiscriminationModel().to(device)

# 创建参数调整器
# 根据生成器和识别器分别创建
optimizer_g = torch.optim.Adam(generation_model.parameters())
optimizer_d = torch.optim.Adam(discrimination_model.parameters())

# 随机生成编码
def generate_vectors(batch_size):
    vectors = torch.randn((batch_size, GenerationModel.EmbeddedSize), device=device)
    return vectors

# 开始训练过程
for epoch in range(0, 10000):
    # 枚举真实数据
    for index, batch_x in enumerate(read_batches()):
        # 生成随机编码
        training_vectors = generate_vectors(minibatch_size)
        # 生成虚假数据
        generated = generation_model(training_vectors)
        # 获取真实数据
        real = batch_x
 
        # 训练识别器 (只调整识别器的参数)
        predicted_t = discrimination_model(real)
        predicted_f = discrimination_model(generated)
        loss_d = (
            nn.functional.binary_cross_entropy(
                predicted_t, torch.ones(predicted_t.shape, device=device)) +
            nn.functional.binary_cross_entropy(
                predicted_f, torch.zeros(predicted_f.shape, device=device)))
        loss_d.backward() # 根据损失自动微分
        optimizer_d.step() # 调整识别器的参数
        optimizer_g.zero_grad() # 清空生成器参数记录的导函数值
        optimizer_d.zero_grad() # 清空识别器参数记录的导函数值

        # 训练生成器 (只调整生成器的参数)
        predicted_f = discrimination_model(generated)
        loss_g = nn.functional.binary_cross_entropy(
            predicted_f, torch.ones(predicted_f.shape, device=device))
        loss_g.backward() # 根据损失自动微分
        optimizer_g.step() # 调整生成器的参数
        optimizer_g.zero_grad() # 清空生成器参数记录的导函数值
        optimizer_d.zero_grad() # 清空识别器参数记录的导函数值

上述例子应该可以帮助你理解大致的训练流程和只训练识别器或生成器的方法,但是直接这么做效果会很差🤕,接下来我们会看看对抗生成网络的问题,并且给出优化方案,后面的完整代码会跟上述例子有一些不同。

如果对原始论文有兴趣可以参考这里,原始的对抗生成网络又称 DCGAN (Deep Convolutional GAN)。

对抗生成网络的问题

看完以上的内容你可能会觉得,嘿嘿,还是挺简单的。不🤕,虽然原理看上去挺好理解,模型本身也不复杂,但对抗生成网络是目前介绍过的模型里面训练难度最高的,这是因为对抗生成网络建立在矛盾上,没有一个明确的目标 (之前的模型目标都是针对未学习过的数据预测正确率尽可能接近 100%)。如果生成器生成 100% 可以骗过识别器的数据,那可能代表识别器根本没正常工作,或者生成器生成的数据跟真实数据 100% 相同,没实用价值;而如果识别器 100% 可以识别生成器生成的数据,那代表生成器生成的数据太垃圾,一个都骗不过。本篇介绍的例子使用了最蠢最简单的方法,把每一轮学习后生成器生成的数据输出到硬盘,然后人工鉴定生成的效果怎样🤒,同时还会每 100 轮训练记录一次模型状态,供训练完以后回滚使用 (最后一个模型状态效果不会是最好的,后面会说明)。

另一个问题是识别器和生成器不能同时训练,怎样安排训练过程对训练结果的影响非常大😮,理想的过程是:识别器稍微领先生成器,生成器跟着识别器慢慢的生成越来越精准的数据。举例来说,识别器首先会识别肤色占比较多的图片为人脸,接下来生成器会生成全部都是肤色的图片,然后识别器会识别有两个看上去是眼睛的图片为人脸,接下来生成器会加上两个看上去是眼睛的形状到图片,之后识别器会识别带有五官的图片为人脸,接下来生成器会加上剩余的五官到图片,最后识别器会识别五官和脸形状比较正常的人为人脸,生成器会尽量调整五官和人脸形状接近正常水平。而不理想的过程是识别器大幅领先生成器,例如识别器很早就达到了接近 100% 的正确率,而生成器因为找不到学习的方向正确率会一直原地踏步;另一个不理想的过程是生成器领先识别器,这时会出现识别器找不到学习的方向,生成器也找不到学习的方向而原地转的情况。实现识别器稍微领先生成器,可以增加识别器的训练次数,常见的方法是每训练 n 次识别器就训练 1 次生成器,而本文后面会介绍根据正确率动态调整识别器和生成器学习次数的方法,参考后面的代码吧。

对抗生成网络最大的问题是模式崩溃 (Mode Collapse) 问题,这个问题所有训练对抗生成网络的人都会面对,并且目前没有 100% 的方法避免😭。简单的来说就是生成器学会偷懒作弊,只会输出一到几个与真实数据几乎一模一样的虚假数据,因为生成的数据同质化非常严重,即使可以骗过识别器也没什么实用价值。发生模式崩溃以后的输出例子如下,可以看到很多人脸都非常接近:

13

为了尽量避免模式崩溃问题,以下几个改进的模型被发明了出来,这就是人民群众的智慧啊😡。

改进对抗生成网络 (WGAN)

模式崩溃问题的原因之一就是部分模型参数会随着训练固化 (达到本地最优),因为原始的对抗生成网络会让识别器输出尽可能接近 1 或者 0 的值,如果值已经是 0 或者 1 那么参数就不会被调整。WGAN (Wasserstein GAN) 的解决方式是不限制识别器输出的值范围,只要求识别器针对真实数据输出的值大于虚假数据输出的值,和要求生成器生成可以让识别器输出更大的值的数据

第一个修改是拿掉识别器最后的 Sigmoid,这样识别器输出的值就不会限制在 0 ~ 1 的范围内。

第二个修改是修改计算损失的方式:

# 计算识别器的损失,修改前
loss_d = (
    nn.functional.binary_cross_entropy(
        predicted_t, torch.ones(predicted_t.shape, device=device)) +
    nn.functional.binary_cross_entropy(
        predicted_f, torch.zeros(predicted_f.shape, device=device)))

# 计算识别器的损失,修改后
loss_d = predicted_f.mean() - predicted_t.mean()
# 计算生成器的损失,修改前
loss_g = nn.functional.binary_cross_entropy(
            predicted_f, torch.ones(predicted_f.shape, device=device))

# 计算生成器的损失,修改后
loss_g = -predicted_f.mean()

这么修改以后会出现一个问题,识别器输出的值范围会随着训练越来越大 (生成器提高虚假数据的输出值,接下来识别器提高真实数据的输出值,循环下去输出值就会越来越大😱),从而导致梯度爆炸问题。为了解决这个问题 WGAN 对识别器参数的可取范围做出了限制,也就是在调整完参数以后裁剪参数,第三个修改如下:

# 让识别器参数必须在 -0.1 ~ 0.1 之间
for p in discrimination_model.parameters():
    p.data.clamp_(-0.1, 0.1)

如果有兴趣可以参考 WGAN 的原始论文,里面一大堆数学公式可以把人吓坏😱,但主要的部分只有上面提到的三点。

改进对抗生成网络 (WGAN-GP)

WGAN 为了防止梯度爆炸问题对识别器参数的可取范围做出了限制,但这个做法比较粗暴,WGAN-GP (Wasserstein GAN Gradient Penalty) 提出了一个更优雅的方法,即限制导函数值的范围,如果导函数值偏移某个指定的值则通过损失给与模型惩罚。

具体实现如下,看起来比较复杂但做的事情只是计算识别器输入数据的导函数值,然后判断所有通道合计的导函数值的 L2 合计与常量 1 相差多少,相差越大就返回越高的损失,这样识别器模型参数自然会控制在某个水平。

def gradient_penalty(discrimination_model, real, generated):
    """控制导函数值的范围,用于防止模型参数失控 (https://arxiv.org/pdf/1704.00028.pdf)"""

    # 给批次中的每个样本分别生成不同的随机值,范围在 0 ~ 1
    batch_size = real.shape[0]
    rate = torch.randn(batch_size, 1, 1, 1)
    rate = rate.expand(batch_size, real.shape[1], real.shape[2], real.shape[3]).to(device)

    # 按随机值比例混合真样本和假样本
    mixed = (rate * real + (1 - rate) * generated)

    # 识别混合样本
    predicted_m = discrimination_model(mixed)

    # 计算 mixed 对 predicted_m 的影响,也就是 mixed => predicted_m 的微分
    # 与以下代码计算结果相同,但不会影响途中 (即模型参数) 的 grad 值
    # mixed = torch.tensor(mixed, requires_grad=True)
    # predicted_m.sum().backward()
    # grad = mixed.grad
    grad = torch.autograd.grad(
        outputs = predicted_m,
        inputs = mixed,
        grad_outputs = torch.ones(predicted_m.shape).to(device),
        create_graph=True,
        retain_graph=True)[0]

    # 让导函数值的 L2 norm (所有通道合计) 在 1 左右,如果偏离 1 则使用损失给与惩罚
    grad_penalty = ((grad.norm(2, dim=1) - 1) ** 2).mean() * 10
    return grad_penalty

然后再修改计算识别器损失的方法:

# 计算识别器的损失,修改前
loss_d = predicted_f.mean() - predicted_t.mean()

# 计算识别器的损失,修改后
loss_d = (predicted_f.mean() - predicted_t.mean() +
    gradient_penalty(discrimination_model, real, generated))

最后把识别器中的批次正规化 (BatchNorm) 删掉或者改为实例正规化 (InstanceNorm) 就完了。InstanceNorm 和 BatchNorm 的区别在于计算平均值和标准差的时候不会根据整个批次计算,而是只根据各个样本自身计算,关于 BatchNorm 的计算方式可以参考第四篇

如果有兴趣可以参考 WGAN-GP 的原始论文

完整代码

又到完整代码的时间了🤗,这份代码同时包含了原始的 GAN 模型 (DCGAN),WGAN 和 WGAN-GP 的实现,后面还会比较它们之间的效果相差多少。

使用的数据集链接如下,前一篇的人脸识别文章也用到了这个数据集:

https://www.kaggle.com/atulanandjha/lfwpeople

需要注意的是人脸图片数量越多就越容易出现模式崩溃问题,这也是对抗生成网络训练的难点之一🤒,这份代码只会随机选取 2000 张图片用于训练。

这份代码还会根据正确率动态调整生成器和识别器的训练比例,如果识别器比生成器更强则训练 1 次生成器,如果生成器比识别器更强则训练 5 次识别器,这么做可以省去手动调整训练比例的麻烦,经实验效果也不错🥳。

import os
import sys
import torch
import gzip
import itertools
import random
import numpy
import math
import json
from PIL import Image
from torch import nn
from matplotlib import pyplot
from functools import lru_cache

# 生成或识别图片的大小
IMAGE_SIZE = (80, 80)
# 训练使用的数据集路径
DATASET_DIR = "./dataset/lfwpeople/lfw_funneled"
# 模型类别, 支持 DCGAN, WGAN, WGAN-GP
MODEL_TYPE = "WGAN-GP"

# 用于启用 GPU 支持
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class GenerationModel(nn.Module):
    """生成虚假数据的模型"""
    # 编码长度
    EmbeddedSize = 128

    def __init__(self):
        super().__init__()
        self.generator = nn.Sequential(
            # 128,1,1 => 512,5,5
            nn.ConvTranspose2d(128, 512, kernel_size=5, stride=1, padding=0, bias=False),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            # => 256,10,10
            nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            # => 128,20,20
            nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            # => 64,40,40
            nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            # => 3,80,80
            nn.ConvTranspose2d(64, 3, kernel_size=4, stride=2, padding=1, bias=False),
            # 限制输出在 -1 ~ 1,不使用 Hardtanh 是为了让超过范围的值可以传播给上层
            nn.Tanh())

    def forward(self, x):
        y = self.generator(x.view(x.shape[0], x.shape[1], 1, 1))
        return y

    @staticmethod
    def calc_accuracy(predicted_f):
        """正确率计算器"""
        # 返回骗过识别器的虚假数据比例
        if MODEL_TYPE == "DCGAN":
            threshold = 0.5
        elif MODEL_TYPE in ("WGAN", "WGAN-GP"):
            threshold = DiscriminationModel.LastTrueSamplePredictedMean
        else:
            raise ValueError("unknown model type")
        return (predicted_f >= threshold).float().mean().item()

class DiscriminationModel(nn.Module):
    """识别数据是否真实的模型"""
    # 最终识别真实样本的输出平均值,WGAN 会使用这个值判断骗过识别器的虚假数据比例
    LastTrueSamplePredictedMean = 0.5

    def __init__(self):
        super().__init__()
        # 标准化函数
        def norm2d(features):
            if MODEL_TYPE == "WGAN-GP":
                # WGAN-GP 本来不需要 BatchNorm,但可以额外的加 InstanceNorm 改善效果
                # InstanceNorm 不一样的是平均值和标准差会针对批次中的各个样本分别计算
                # affine = True 表示调整量可学习 (BatchNorm2d 默认为 True)
                return nn.InstanceNorm2d(features, affine=True)
            return nn.BatchNorm2d(features)
        self.discriminator = nn.Sequential(
            # 3,80,80 => 64,40,40
            nn.Conv2d(3, 64, kernel_size=4, stride=2, padding=1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            # => 128,20,20
            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1, bias=False),
            norm2d(128),
            nn.LeakyReLU(0.2, inplace=True),
            # => 256,10,10
            nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1, bias=False),
            norm2d(256),
            nn.LeakyReLU(0.2, inplace=True),
            # => 512,5,5
            nn.Conv2d(256, 512, kernel_size=4, stride=2, padding=1, bias=False),
            norm2d(512),
            nn.LeakyReLU(0.2, inplace=True),
            # => 1,1,1
            nn.Conv2d(512, 1, kernel_size=5, stride=1, padding=0, bias=False),
            # 扁平化
            nn.Flatten())
        if MODEL_TYPE == "DCGAN":
            # 输出是否真实数据 (0 or 1)
            # WGAN 不限制输出值范围在 0 ~ 1 之间
            self.discriminator.add_module("sigmoid", nn.Sigmoid())

    def forward(self, x):
        y = self.discriminator(x)
        return y

    @staticmethod
    def calc_accuracy(predicted_f, predicted_t):
        """正确率计算器"""
        # 返回正确识别的数据比例
        if MODEL_TYPE == "DCGAN":
            return (((predicted_f <= 0.5).float().mean() + (predicted_t > 0.5).float().mean()) / 2).item()
        elif MODEL_TYPE in ("WGAN", "WGAN-GP"):
            DiscriminationModel.LastTrueSamplePredictedMean = predicted_t.mean()
            return (predicted_t > predicted_f).float().mean().item()
        else:
            raise ValueError("unknown model type")

    def gradient_penalty(self, real, generated):
        """控制导函数值的范围,用于防止模型参数失控 (https://arxiv.org/pdf/1704.00028.pdf)"""
        # 给批次中的每个样本分别生成不同的随机值,范围在 0 ~ 1
        batch_size = real.shape[0]
        rate = torch.randn(batch_size, 1, 1, 1)
        rate = rate.expand(batch_size, real.shape[1], real.shape[2], real.shape[3]).to(device)
        # 按随机值比例混合真样本和假样本
        mixed = (rate * real + (1 - rate) * generated)
        # 识别混合样本
        predicted_m = self.forward(mixed)
        # 计算 mixed 对 predicted_m 的影响,也就是 mixed => predicted_m 的微分
        # 与以下代码计算结果相同,但不会影响途中 (即模型参数) 的 grad 值
        # mixed = torch.tensor(mixed, requires_grad=True)
        # predicted_m.sum().backward()
        # grad = mixed.grad
        grad = torch.autograd.grad(
            outputs = predicted_m,
            inputs = mixed,
            grad_outputs = torch.ones(predicted_m.shape).to(device),
            create_graph=True,
            retain_graph=True)[0]
        # 让导函数值的 L2 norm (所有通道合计) 在 1 左右,如果偏离 1 则使用损失给与惩罚
        grad_penalty = ((grad.norm(2, dim=1) - 1) ** 2).mean() * 10
        return grad_penalty

def save_tensor(tensor, path):
    """保存 tensor 对象到文件"""
    torch.save(tensor, gzip.GzipFile(path, "wb"))

# 为了减少读取时间这里缓存了读取的 tensor 对象
# 如果内存不够应该适当减少 maxsize
@lru_cache(maxsize=200)
def load_tensor(path):
    """从文件读取 tensor 对象"""
    return torch.load(gzip.GzipFile(path, "rb"))

def image_to_tensor(img):
    """缩放并转换图片对象到 tensor 对象"""
    img = img.resize(IMAGE_SIZE) # 缩放图片,比例不一致时拉伸
    arr = numpy.asarray(img)
    t = torch.from_numpy(arr)
    t = t.transpose(0, 2) # 转换维度 H,W,C 到 C,W,H
    t = (t / 255.0) * 2 - 1 # 正规化数值使得范围在 -1 ~ 1
    return t

def tensor_to_image(t):
    """转换 tensor 对象到图片"""
    t = (t + 1) / 2 * 255.0 # 转换颜色回 0 ~ 255
    t = t.transpose(0, 2) # 转换维度 C,W,H 到 H,W,C
    t = t.int() # 转换数值到整数
    img = Image.fromarray(t.numpy().astype("uint8"), "RGB")
    return img

def prepare():
    """准备训练"""
    # 数据集转换到 tensor 以后会保存在 data 文件夹下
    if not os.path.isdir("data"):
        os.makedirs("data")

    # 查找人脸图片列表
    # 每个人最多使用 2 张图片
    image_paths = []
    for dirname in os.listdir(DATASET_DIR):
        dirpath = os.path.join(DATASET_DIR, dirname)
        if not os.path.isdir(dirpath):
            continue
        for filename in os.listdir(dirpath)[:2]:
            image_paths.append(os.path.join(DATASET_DIR, dirname, filename))
    print(f"found {len(image_paths)} images")

    # 随机打乱人脸图片列表
    random.shuffle(image_paths)

    # 限制人脸数量
    # 如果数量太多,识别器难以记住人脸的具体特征,会需要更长时间训练或直接陷入模式崩溃问题
    image_paths = image_paths[:2000]
    print(f"only use {len(image_paths)} images")

    # 保存人脸图片数据
    for batch, index in enumerate(range(0, len(image_paths), 200)):
        paths = image_paths[index:index+200]
        images = []
        for path in paths:
            img = Image.open(path)
            # 扩大人脸占比
            w, h = img.size
            img = img.crop((int(w*0.25), int(h*0.25), int(w*0.75), int(h*0.75)))
            images.append(img)
        tensors = [ image_to_tensor(img) for img in images ]
        tensor = torch.stack(tensors) # 维度: (图片数量, 3, 宽度, 高度)
        save_tensor(tensor, os.path.join("data", f"{batch}.pt"))
        print(f"saved batch {batch}")

    print("done")

def train():
    """开始训练模型"""
    # 创建模型实例
    generation_model = GenerationModel().to(device)
    discrimination_model = DiscriminationModel().to(device)

    # 创建损失计算器
    ones_map = {}
    zeros_map = {}
    def loss_function_t(predicted):
        """损失计算器 (训练识别结果为 1)"""
        count = predicted.shape[0]
        ones = ones_map.get(count)
        if ones is None:
            ones = torch.ones((count, 1), device=device)
            ones_map[count] = ones
        return nn.functional.binary_cross_entropy(predicted, ones)
    def loss_function_f(predicted):
        """损失计算器 (训练识别结果为 0)"""
        count = predicted.shape[0]
        zeros = zeros_map.get(count)
        if zeros is None:
            zeros = torch.zeros((count, 1), device=device)
            zeros_map[count] = zeros
        return nn.functional.binary_cross_entropy(predicted, zeros)

    # 创建参数调整器
    if MODEL_TYPE == "DCGAN":
        optimizer_g = torch.optim.Adam(generation_model.parameters(), lr=0.0002, betas=(0.5, 0.999))
        optimizer_d = torch.optim.Adam(discrimination_model.parameters(), lr=0.0002, betas=(0.5, 0.999))
    elif MODEL_TYPE == "WGAN":
        optimizer_g = torch.optim.RMSprop(generation_model.parameters(), lr=0.00005)
        optimizer_d = torch.optim.RMSprop(discrimination_model.parameters(), lr=0.00005)
    elif MODEL_TYPE == "WGAN-GP":
        optimizer_g = torch.optim.Adam(generation_model.parameters(), lr=0.0001, betas=(0.0, 0.999))
        optimizer_d = torch.optim.Adam(discrimination_model.parameters(), lr=0.0001, betas=(0.0, 0.999))
    else:
        raise ValueError("unknown model type")

    # 记录训练集和验证集的正确率变化
    training_accuracy_g_history = []
    training_accuracy_d_history = []

    # 计算正确率的工具函数
    calc_accuracy_g = generation_model.calc_accuracy
    calc_accuracy_d = discrimination_model.calc_accuracy

    # 随机生成编码
    def generate_vectors(batch_size):
        vectors = torch.randn((batch_size, GenerationModel.EmbeddedSize), device=device)
        return vectors

    # 输出生成的图片样本
    def output_generated_samples(epoch, samples):
        dir_path = f"./generated_samples/{epoch}"
        if not os.path.isdir(dir_path):
            os.makedirs(dir_path)
        for index, sample in enumerate(samples):
            path = os.path.join(dir_path, f"{index}.png")
            tensor_to_image(sample.cpu()).save(path)

    # 读取批次的工具函数
    def read_batches():
        for batch in itertools.count():
            path = f"data/{batch}.pt"
            if not os.path.isfile(path):
                break
            x = load_tensor(path)
            yield x.to(device)

    # 开始训练过程
    validating_vectors = generate_vectors(100)
    for epoch in range(0, 10000):
        print(f"epoch: {epoch}")

        # 根据训练集训练并修改参数
        # 切换模型到训练模式
        generation_model.train()
        discrimination_model.train()
        training_accuracy_g_list = []
        training_accuracy_d_list = []
        last_accuracy_g = 0
        last_accuracy_d = 0
        minibatch_size = 20
        train_discriminator_count = 0
        for index, batch_x in enumerate(read_batches()):
            # 使用小批次训练
            training_batch_accuracy_g = 0.0
            training_batch_accuracy_d = 0.0
            minibatch_count = 0
            for begin in range(0, batch_x.shape[0], minibatch_size):
                # 测试目前生成器和识别器哪边占劣势,训练占劣势的一方
                # 最终的平衡状态是: 生成器正确率 = 1.0, 识别器正确率 = 0.5
                # 代表生成器生成的图片和真实图片基本完全一样,但不应该训练到这个程度
                training_vectors = generate_vectors(minibatch_size) # 随机向量
                generated = generation_model(training_vectors) # 根据随机向量生成的虚假数据
                real = batch_x[begin:begin+minibatch_size] # 真实数据
                predicted_t = discrimination_model(real)
                predicted_f = discrimination_model(generated)
                accuracy_g = calc_accuracy_g(predicted_f)
                accuracy_d = calc_accuracy_d(predicted_f, predicted_t)
                train_discriminator = (accuracy_g / 2) >= accuracy_d
                if train_discriminator or train_discriminator_count > 0:
                    # 训练识别器
                    if MODEL_TYPE == "DCGAN":
                        loss_d = loss_function_f(predicted_f) + loss_function_t(predicted_t)
                    elif MODEL_TYPE == "WGAN":
                        loss_d = predicted_f.mean() - predicted_t.mean()
                    elif MODEL_TYPE == "WGAN-GP":
                        loss_d = (predicted_f.mean() - predicted_t.mean() +
                            discrimination_model.gradient_penalty(real, generated))
                    else:
                        raise ValueError("unknown model type")
                    loss_d.backward()
                    optimizer_d.step()
                    # 限制识别器参数范围以防止模型参数失控 (WGAN-GP 有更好的方法)
                    # 这里的限制值比论文的值 (0.01) 更大是因为模型层数和参数量更多
                    if MODEL_TYPE == "WGAN":
                        for p in discrimination_model.parameters():
                            p.data.clamp_(-0.1, 0.1)
                    # 让识别器训练次数多于生成器
                    if train_discriminator and train_discriminator_count == 0:
                        train_discriminator_count = 5
                    train_discriminator_count -= 1
                else:
                    # 训练生成器
                    if MODEL_TYPE == "DCGAN":
                        loss_g = loss_function_t(predicted_f)
                    elif MODEL_TYPE in ("WGAN", "WGAN-GP"):
                        loss_g = -predicted_f.mean()
                    else:
                        raise ValueError("unknown model type")
                    loss_g.backward()
                    optimizer_g.step()
                optimizer_g.zero_grad()
                optimizer_d.zero_grad()
                training_batch_accuracy_g += accuracy_g
                training_batch_accuracy_d += accuracy_d
                minibatch_count += 1
            training_batch_accuracy_g /= minibatch_count
            training_batch_accuracy_d /= minibatch_count
            # 输出批次正确率
            training_accuracy_g_list.append(training_batch_accuracy_g)
            training_accuracy_d_list.append(training_batch_accuracy_d)
            print(f"epoch: {epoch}, batch: {index},",
                f"accuracy_g: {training_batch_accuracy_g}, accuracy_d: {training_batch_accuracy_d}")
        training_accuracy_g = sum(training_accuracy_g_list) / len(training_accuracy_g_list)
        training_accuracy_d = sum(training_accuracy_d_list) / len(training_accuracy_d_list)
        training_accuracy_g_history.append(training_accuracy_g)
        training_accuracy_d_history.append(training_accuracy_d)
        print(f"training accuracy_g: {training_accuracy_g}, accuracy_d: {training_accuracy_d}")

        # 保存虚假数据用于评价训练效果
        output_generated_samples(epoch, generation_model(validating_vectors))

        # 保存模型状态
        if (epoch + 1) % 10 == 0:
            save_tensor(generation_model.state_dict(), "model.generation.pt")
            save_tensor(discrimination_model.state_dict(), "model.discrimination.pt")
            if (epoch + 1) % 100 == 0:
                save_tensor(generation_model.state_dict(), f"model.generation.epoch_{epoch}.pt")
                save_tensor(discrimination_model.state_dict(), f"model.discrimination.epoch_{epoch}.pt")
            print("model saved")

    print("training finished")

    # 显示训练集的正确率变化
    pyplot.plot(training_accuracy_g_history, label="training_accuracy_g")
    pyplot.plot(training_accuracy_d_history, label="training_accuracy_d")
    pyplot.ylim(0, 1)
    pyplot.legend()
    pyplot.show()

from http.server import HTTPServer, BaseHTTPRequestHandler
from urllib.parse import parse_qs
from io import BytesIO
class RequestHandler(BaseHTTPRequestHandler):
    """用于测试生成图片的简单服务器"""
    # 模型状态的路径,这里使用看起来效果最好的记录
    MODEL_STATE_PATH = "model.generation.epoch_2999.pt"
    Model = None

    @staticmethod
    def get_model():
        if RequestHandler.Model is None:
            # 创建模型实例,加载训练好的状态,然后切换到验证模式
            model = GenerationModel().to(device)
            model.load_state_dict(load_tensor(RequestHandler.MODEL_STATE_PATH))
            model.eval()
            RequestHandler.Model = model
        return RequestHandler.Model

    def do_GET(self):
        parts = self.path.partition("?")
        if parts[0] == "/":
            self.send_response(200)
            self.send_header("Content-type", "text/html")
            self.end_headers()
            with open("gan_eval.html", "rb") as f:
                self.wfile.write(f.read())
        elif parts[0] == "/generate":
            # 根据传入的参数生成图片
            params = parse_qs(parts[-1])
            vector = (torch.tensor([float(x) for x in params["values"][0].split(",")])
                .reshape(1, GenerationModel.EmbeddedSize)
                .to(device))
            generated = RequestHandler.get_model()(vector)[0]
            img = tensor_to_image(generated.cpu())
            bytes_io = BytesIO()
            img.save(bytes_io, format="PNG")
            # 返回图片
            self.send_response(200)
            self.send_header("Content-type", "image/png")
            self.end_headers()
            self.wfile.write(bytes_io.getvalue())
        else:
            self.send_response(404)
            self.end_headers()
            self.wfile.write(b"Not Found")

def eval_model():
    """使用训练好的模型生成图片"""
    server = HTTPServer(("localhost", 8666), RequestHandler)
    print("Please access http://localhost:8666")
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        pass
    server.server_close()
    exit()

def main():
    """主函数"""
    if len(sys.argv) < 2:
        print(f"Please run: {sys.argv[0]} prepare|train|eval")
        exit()

    # 给随机数生成器分配一个初始值,使得每次运行都可以生成相同的随机数
    # 这是为了让过程可重现,你也可以选择不这样做
    random.seed(0)
    torch.random.manual_seed(0)

    # 根据命令行参数选择操作
    operation = sys.argv[1]
    if operation == "prepare":
        prepare()
    elif operation == "train":
        train()
    elif operation == "eval":
        eval_model()
    else:
        raise ValueError(f"Unsupported operation: {operation}")

if __name__ == "__main__":
    main()

保存代码到 gan.py,然后执行以下命令即可开始训练:

python3 gan.py prepare
python3 gan.py train

同样训练 2000 轮以后,DCGAN, WGAN, WGAN-GP 输出的样本如下:

DCGAN

samples_dcgan

WGAN

samples_wgan

WGAN-GP

samples_wgan

可以看到 WGAN-GP 受模式崩溃问题影响最少,并且效果也更好😤。

WGAN-GP 训练到 3000 次以后输出的样本如下:

samples_wgan_gp_3000

WGAN-GP 训练到 10000 次以后输出的样本如下:

samples_wgan_gp_10000.png

随着训练次数增多,WGAN-GP 一样无法避免模式崩溃问题,这就是为什么以上代码会记录每一轮训练后输出的样本,并在每 100 轮训练以后保存单独的模型状态,这样训练结束以后我们可以通过评价输出的样本找到效果最好的批次,然后使用该批次的模型状态。

上述的例子效果最好的状态是训练 3000 次以后的状态。

你可能发现输出的样本中夹杂了一些畸形🥴,这是因为生成器没有覆盖到输入的向量空间,最主要的原因是随机输入中包含了很多接近 0 的值,避免这个问题简单的做法是生成随机输入时限制值必须小于或大于某个值。原则上给反卷积层设置 Bias 也可以避免这个问题,但会更容易陷入模式崩溃问题。

使用训练好的模型生成人脸就比较简单了:

generation_model = GenerationModel().to(device)
model.load_state_dict(load_tensor("model.generation.epoch_2999.pt"))
model.eval()

# 随机生成 100 张人脸
vector = torch.randn((100, GenerationModel.EmbeddedSize), device=device)
samples = model(vector)
for index, sample in enumerate(samples):
    img = tensor_to_image(sample.cpu())
    img.save(f"{index}.png")

额外的,我做了一个可以动态调整参数捏脸的网页,html 代码如下:

<!DOCTYPE html>
<html lang="cn">
  <head>
    <meta charset="utf-8">
    <title>测试人脸生成</title>
    <style>
        html, body {
            width: 100%;
            height: 100%;
            margin: 0px;
        }
        .left-pane {
            width: 50%;
            height: 100%;
            border-right: 1px solid #000;
        }
        .right-pane {
            position: fixed;
            left: 70%;
            top: 35%;
            width: 25%;
        }
        .sliders {
            padding: 8px;
        }
        .slider-container {
            display: inline-block;
            min-width: 25%;
        }
        #image {
            left: 25%;
            top: 25%;
            width: 50%;
            height: 50%;
        }
    </style>
  </head>
  <body>
    <div class="left-pane">
        <div class="sliders">
        </div>
    </div>
    <div class="right-pane">
       <p><img id="target" src="data:image/png;base64," alt="image" /></p>
       <p><button class="set-random">随机生成</button></p>
    </div>
  </body>
  <script>
    (function() {
        // 滑动条改变后的处理
        var onChanged = function() {
            var sliderInputs = document.querySelectorAll(".slider");
            var values = [];
            sliderInputs.forEach(function(s) {
                values.push(s.value);
            });
            var image = document.querySelector("#target");
            image.setAttribute("src", "/generate?values=" + values.join(","));
        };

        // 点击随机生成时的处理
        var setRandomButton = document.querySelector(".set-random");
        setRandomButton.onclick = function() {
            var sliderInputs = document.querySelectorAll(".slider");
            sliderInputs.forEach(function(s) { s.value = Math.random() * 2 - 1; });
            onChanged();
        };

        // 添加滑动条
        var sliders = document.querySelector(".sliders");
        for (var n = 0; n < 128; ++n) {
            var container = document.createElement("div");
            container.setAttribute("class", "slider-container");
            var span = document.createElement("span");
            span.innerText = n;
            container.appendChild(span);
            var slider = document.createElement("input");
            slider.setAttribute("type", "range")
            slider.setAttribute("class", "slider");
            slider.setAttribute("min", "-1");
            slider.setAttribute("max", "1");
            slider.setAttribute("step", "0.01");
            slider.value = 0;
            slider.onchange = onChanged;
            slider.oninput = onChanged;
            container.appendChild(slider);
            sliders.appendChild(container);
        }
    })();
  </script>
</html>

保存到 gan_eval.html 以后执行以下命令即可启动服务器:

python3 gan.py eval

浏览器打开 http://localhost:8666 以后会显示以下界面,点击随机生成按钮可以随机生成人脸,拉动左边的参数条可以动态调整参数:

eval_html

一些捏脸的网站会分析各个参数的含义,看看哪些参数代表肤色,那些参数代表表情,哪些参数代表脱发程度,我比较懒就只给出各个参数的序号了🤒。

写在最后

又摸完一个新的模型了,跟到这篇的人也越来越少了,估计这个系列再写一两篇就会结束 (VAE, 强化学习)。

Written on April 22, 2021