图像生成之条件生成对抗网络（CGAN）

简要介绍

​ 原文：Conditional Generative Adversarial Nets

​ 这篇十年前的论文提出了多模态的思想，在我看来真的很有意义，所以我专门去写这篇文章。

​ 在原文中，作者指出扩大神经网络模型去适应大量的预测输出类别是很有挑战性的，所以作者提到，利用来自其他模态的附加信息去解决这个问题: 例如，通过使用自然语言语料库来学习标签的向量表示，因为几何关系具有语义意义，可以自然地对训练期间未见过的标签进行预测概括。研究表明，即使是简单的从图像特征空间到词表示空间的线性映射也能提高分类性能。

​ CGAN问世以前的大部分工作都集中在学习从输入到输出的一对一映射上。但是，许多有趣的问题更自然地被认为是概率性的一对多映射。比如说一种图像可以由多种同义的术语表示，即图像可能有多种标签映射得到。所以作者提出使用条件概率生成模型，将输入作为条件变量，将一对多映射实例化为条件预测分布。

图2 CGAN模型体系结构图示

目标函数

​ 如果生成器和判别器都以一些额外的信息y为条件，那么生成对抗网络可以扩展到条件模型。y可以是任何类型的辅助信息，例如类标签或来自其他模态的数据。可以通过将y作为附加输入层同时输入鉴别器和生成器来执行调节。
$$\underset{G}{\min }\underset{D}{\max }V(D,G)=E_{x\sim p_{\text{data}}(x)}[\log (D(x|y))]+E_{z\sim p_z(z)}[\log (1-D(G(x|y))]$$

​ 无论是生成器还是判别器，其概率表达都变成了条件概率公式，也就是对于不同的输入条件$y$，可以得到不同的目标函数。很多任务都是通过这种生成去实现。$y$可以是任何类型的辅助信息，例如类标签或来自其他模态的数据。CGAN相对于一般的GAN在输入上引入了额外的条件信息，通过这些信息指导生成器生成特定的图像，从而实现更加精确和可控的图像生成。

​ 虽然目标函数的公式改变了，但GAN架构的输出、目标函数的具体实现方式等都没有发生变化，模型最终输出的依然是一个概率值，只不过在数学上它的性质由普通概率变化为了条件概率。由于在神经网络中，条件概率的条件就表示输入，例如： $p(b|a)$，表示网络的输入为a，输出为b。因此，在编程实现时，对于GAN的网络架构我们也不需要做太多的修改，只需要在输入时加上真实标签、词向量等即可。原理与公式推导也与GAN一模一样，详见另一篇我的文章：图像生成之生成对抗网络（Generative adversarial networks，GAN）

训练代码

import argparse
import os
import numpy as np
import math

import torchvision.transforms as transforms
from torchvision.utils import save_image

from torch.utils.data import DataLoader
from torchvision import datasets
from torch.autograd import Variable

import torch.nn as nn
import torch.nn.functional as F
import torch

# 创建保存生成图片的目录
os.makedirs("images", exist_ok=True)

# 解析命令行参数
parser = argparse.ArgumentParser()
parser.add_argument("--n_epochs", type=int, default=200, help="训练的轮数")
parser.add_argument("--batch_size", type=int, default=64, help="每批的样本数量")
parser.add_argument("--lr", type=float, default=0.0002, help="Adam优化器的学习率")
parser.add_argument("--b1", type=float, default=0.5, help="Adam优化器的梯度一阶矩的衰减系数")
parser.add_argument("--b2", type=float, default=0.999, help="Adam优化器的梯度二阶矩的衰减系数")
parser.add_argument("--n_cpu", type=int, default=8, help="用于生成批次的CPU线程数")
parser.add_argument("--latent_dim", type=int, default=100, help="潜在空间的维度")
parser.add_argument("--n_classes", type=int, default=10, help="数据集的类别数")
parser.add_argument("--img_size", type=int, default=32, help="每个图像维度的大小")
parser.add_argument("--channels", type=int, default=1, help="图像的通道数")
parser.add_argument("--sample_interval", type=int, default=400, help="生成图像样本的间隔")
opt = parser.parse_args()
print(opt)

# 定义图像形状
img_shape = (opt.channels, opt.img_size, opt.img_size)

cuda = True if torch.cuda.is_available() else False

# 定义生成器模型
class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()

        self.label_emb = nn.Embedding(opt.n_classes, opt.n_classes)

        def block(in_feat, out_feat, normalize=True):
            layers = [nn.Linear(in_feat, out_feat)]
            if normalize:
                layers.append(nn.BatchNorm1d(out_feat, 0.8))
            layers.append(nn.LeakyReLU(0.2, inplace=True))
            return layers

        self.model = nn.Sequential(
            *block(opt.latent_dim + opt.n_classes, 128, normalize=False),
            *block(128, 256),
            *block(256, 512),
            *block(512, 1024),
            nn.Linear(1024, int(np.prod(img_shape))),
            nn.Tanh()
        )

    def forward(self, noise, labels):
        # 拼接标签嵌入和噪声以生成输入
        gen_input = torch.cat((self.label_emb(labels), noise), -1)
        img = self.model(gen_input)
        img = img.view(img.size(0), *img_shape)
        return img


# 定义鉴别器模型
class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()

        self.label_embedding = nn.Embedding(opt.n_classes, opt.n_classes)

        self.model = nn.Sequential(
            nn.Linear(opt.n_classes + int(np.prod(img_shape)), 512),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(512, 512),
            nn.Dropout(0.4),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(512, 512),
            nn.Dropout(0.4),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(512, 1),
        )

    def forward(self, img, labels):
        # 拼接标签嵌入和图像以生成输入
        d_in = torch.cat((img.view(img.size(0), -1), self.label_embedding(labels)), -1)
        validity = self.model(d_in)
        return validity


# 损失函数
adversarial_loss = torch.nn.MSELoss()

# 初始化生成器和鉴别器
generator = Generator()
discriminator = Discriminator()

if cuda:
    generator.cuda()
    discriminator.cuda()
    adversarial_loss.cuda()

# 配置dataloader
os.makedirs("../../data/mnist", exist_ok=True)
dataloader = torch.utils.data.DataLoader(
    datasets.MNIST(
        "../../data/mnist",
        train=True,
        download=True,
        transform=transforms.Compose(
            [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]
        ),
    ),
    batch_size=opt.batch_size,
    shuffle=True,
)

# 优化器
optimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))
optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))

FloatTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor
LongTensor = torch.cuda.LongTensor if cuda else torch.LongTensor

# 保存生成的图像函数
def sample_image(n_row, batches_done):
    """Saves a grid of generated digits ranging from 0 to n_classes"""
    # 生成噪声
    z = Variable(FloatTensor(np.random.normal(0, 1, (n_row ** 2, opt.latent_dim))))
    # 获取标签从0到n_classes的n行
    labels = np.array([num for _ in range(n_row) for num in range(n_row)])
    labels = Variable(LongTensor(labels))
    gen_imgs = generator(z, labels)
    save_image(gen_imgs.data, "images/%d.png" % batches_done, nrow=n_row, normalize=True)

# ----------
#  训练
# ----------

for epoch in range(opt.n_epochs):
    for i, (imgs, labels) in enumerate(dataloader):

        batch_size = imgs.shape[0]

        # 对抗网络的真实标签和生成标签
        valid = Variable(FloatTensor(batch_size, 1).fill_(1.0), requires_grad=False)
        fake = Variable(FloatTensor(batch_size, 1).fill_(0.0), requires_grad=False)

        # 输入
        real_imgs = Variable(imgs.type(FloatTensor))
        labels = Variable(labels.type(LongTensor))

        # -----------------
        #  训练生成器
        # -----------------

        optimizer_G.zero_grad()

        # 采样噪声和标签作为生成器的输入
                z = Variable(FloatTensor(np.random.normal(0, 1, (batch_size, opt.latent_dim))))
        gen_labels = Variable(LongTensor(np.random.randint(0, opt.n_classes, batch_size)))

        # 生成一批图像
        gen_imgs = generator(z, gen_labels)

        # 生成器的损失，度量其欺骗鉴别器的能力
        validity = discriminator(gen_imgs, gen_labels)
        g_loss = adversarial_loss(validity, valid)

        g_loss.backward()
        optimizer_G.step()

        # ---------------------
        #  训练鉴别器
        # ---------------------

        optimizer_D.zero_grad()

        # 对真实图像的损失
        validity_real = discriminator(real_imgs, labels)
        d_real_loss = adversarial_loss(validity_real, valid)

        # 对生成图像的损失
        validity_fake = discriminator(gen_imgs.detach(), gen_labels)
        d_fake_loss = adversarial_loss(validity_fake, fake)

        # 总体鉴别器损失
        d_loss = (d_real_loss + d_fake_loss) / 2

        d_loss.backward()
        optimizer_D.step()

        print(
            "[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]"
            % (epoch, opt.n_epochs, i, len(dataloader), d_loss.item(), g_loss.item())
        )

        batches_done = epoch * len(dataloader) + i
        if batches_done % opt.sample_interval == 0:
            sample_image(n_row=10, batches_done=batches_done)