PyTorch深度学习框架60天进阶学习计划 - 第41天：生成对抗网络进阶（一）

PyTorch深度学习框架60天进阶学习计划 - 第41天

生成对抗网络进阶（一）：Wasserstein GAN的梯度惩罚机制与模式坍塌问题

今天我们要"对抗"一个相当有趣又有挑战性的主题——Wasserstein GAN（WGAN）的梯度惩罚机制以及条件生成与无监督生成中模式坍塌的差异。

我们的神经网络已经从最初的"小白"成长为了能创造全新内容的"艺术家"了！当我第一次看到GAN生成的假脸时，我简直惊呆了——“这不是真人吗？”。但在GAN的修炼之路上，也经常会遇到各种各样的"魔障"，而今天我们就要学习如何突破其中两大难关：梯度惩罚和模式坍塌。

第一部分：Wasserstein GAN的梯度惩罚机制

1. 标准GAN的训练困境

首先，让我们回顾一下为什么我们需要WGAN。在标准GAN（Goodfellow等人在2014年提出）中，我们面临几个关键问题：

1. 训练不稳定：判别器很容易变得过于强大，导致生成器梯度消失
2. 模式坍塌：生成器只学会产生有限种类的样本
3. 难以判断收敛：没有可靠的指标来判断训练何时应该停止
4. 超参数敏感：对学习率等超参数非常敏感

这些问题就像是GAN训练路上的"拦路虎"，让很多人望而却步。Wasserstein GAN正是为了解决这些问题而生的。

2. Wasserstein距离的引入

在标准GAN中，判别器试图最大化真实数据和生成数据之间的JS散度（Jensen-Shannon divergence）。然而，当两个分布的支撑集（support）没有显著重叠时，JS散度几乎是常数，这导致了梯度消失问题。

而Wasserstein距离（也称为Earth Mover’s Distance，推土机距离）提供了一个更平滑的度量：

W(P_r, P_g) = inf_{γ∈Π(P_r,P_g)} E_{(x,y)~γ}[||x-y||]

其中，Π(P_r,P_g)是所有可能的联合分布γ的集合，满足其边缘分布分别是P_r和P_g。

直观地说，Wasserstein距离衡量的是将一个分布"推"成另一个分布所需的最小"工作量"。

这就好比：

• JS散度像是判断两座山是否完全重叠
• Wasserstein距离则是计算将一座山的土推到另一座山所需的最小工作量

即使两座山完全分开，计算推土所需的工作量仍然是有意义的！

3. Wasserstein GAN的基本原理

WGAN的关键创新是使用Wasserstein距离而非JS散度，这带来了几个关键变化：

1. 移除判别器最后的sigmoid层（因为不再是二元分类问题）
2. 判别器（现在称为"评论家/critic"）不再区分真假，而是为每个样本分配一个"真实度"得分
3. 不使用对数损失，而是直接使用真实样本和生成样本评分之差
4. 对评论家的参数进行权重裁剪（weight clipping），确保满足1-Lipschitz约束

WGAN的目标函数如下：

min_G max_D E_{x~P_r}[D(x)] - E_{z~P_z}[D(G(z))]

其中D的参数必须保持在一个紧凑空间内（通过权重裁剪实现）。

4. 权重裁剪的局限性

原始WGAN使用权重裁剪来强制执行Lipschitz约束。具体来说，在每次参数更新后，将判别器的权重值裁剪到[-c, c]范围内：

for p in discriminator.parameters():
    p.data.clamp_(-c, c)

然而，权重裁剪存在几个问题：

1. 容量问题：可能导致模型容量降低
2. 梯度爆炸/消失：可能导致梯度爆炸或消失
3. 寻路问题：可能迫使网络选择次优路径

正如Ian Goodfellow所说：“权重裁剪就像是用大锤子来杀蚊子——有效但不优雅。”

5. 梯度惩罚（Gradient Penalty）机制

为了解决权重裁剪的问题，WGAN-GP（Gradient Penalty）被提出。梯度惩罚是一种更优雅的方式来强制Lipschitz约束。

Lipschitz约束本质上要求判别器关于输入的梯度范数不超过某个常数。在WGAN-GP中，我们通过惩罚梯度范数偏离1的行为来实现这一点：

L = E_{x~P_r}[D(x)] - E_{z~P_z}[D(G(z))] + λ * E_{x̂~P_x̂}[(||∇_x̂ D(x̂)||_2 - 1)²]

其中，x̂是真实样本和生成样本之间的随机插值点：

x̂ = εx + (1-ε)G(z)，ε~U[0,1]

这种方法有几个优点：

1. 保持模型容量：不会人为限制模型表达能力
2. 稳定的梯度：避免了梯度爆炸/消失问题
3. 更好的收敛性：训练更稳定，生成质量更高

6. WGAN-GP的实现细节

让我们看看如何在PyTorch中实现WGAN-GP。首先，我们需要计算梯度惩罚项：

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

# 计算梯度惩罚
def compute_gradient_penalty(D, real_samples, fake_samples, device):
    """计算WGAN-GP的梯度惩罚"""
    # 随机插值系数
    alpha = torch.rand(real_samples.size(0), 1, 1, 1).to(device)
    # 在真实样本和生成样本之间进行插值
    interpolates = (alpha * real_samples + (1 - alpha) * fake_samples).requires_grad_(True)
    # 计算插值点的判别器输出
    d_interpolates = D(interpolates)
    # 创建与d_interpolates形状相同的全1张量
    fake = torch.ones(d_interpolates.size()).to(device)
    # 计算梯度
    gradients = autograd.grad(
        outputs=d_interpolates,
        inputs=interpolates,
        grad_outputs=fake,
        create_graph=True,
        retain_graph=True,
        only_inputs=True,
    )[0]
    # 计算梯度的范数
    gradients = gradients.view(gradients.size(0), -1)
    gradient_norm = gradients.norm(2, dim=1)
    # 计算梯度惩罚 (||∇D(x̂)||_2 - 1)²
    gradient_penalty = ((gradient_norm - 1) ** 2).mean()
    return gradient_penalty

# 简单的生成器和判别器网络定义
class Generator(nn.Module):
    def __init__(self, latent_dim, img_size, channels):
        super(Generator, self).__init__()
        self.img_shape = (channels, img_size, img_size)
        
        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(latent_dim, 128, normalize=False),
            *block(128, 256),
            *block(256, 512),
            *block(512, 1024),
            nn.Linear(1024, int(np.prod(self.img_shape))),
            nn.Tanh()
        )
        
    def forward(self, z):
        img = self.model(z)
        img = img.view(img.size(0), *self.img_shape)
        return img

class Discriminator(nn.Module):
    def __init__(self, img_size, channels):
        super(Discriminator, self).__init__()
        self.img_shape = (channels, img_size, img_size)
        
        self.model = nn.Sequential(
            nn.Linear(int(np.prod(self.img_shape)), 512),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(512, 256),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(256, 1),
            # 注意：WGAN中没有sigmoid激活函数
        )
        
    def forward(self, img):
        img_flat = img.view(img.size(0), -1)
        validity = self.model(img_flat)
        return validity

# WGAN-GP训练循环（部分代码）
def train_wgan_gp(dataloader, latent_dim, n_critic, lambda_gp, 
                 generator, discriminator, g_optimizer, d_optimizer, device, n_epochs=100):
    # 训练循环
    for epoch in range(n_epochs):
        for i, (real_imgs, _) in enumerate(dataloader):
            real_imgs = real_imgs.to(device)
            
            # -----------------
            #  训练判别器
            # -----------------
            d_optimizer.zero_grad()
            
            # 采样噪声并生成假图像
            z = torch.randn(real_imgs.size(0), latent_dim).to(device)
            fake_imgs = generator(z)
            
            # 计算真实图像、生成图像的判别器输出
            real_validity = discriminator(real_imgs)
            fake_validity = discriminator(fake_imgs.detach())
            
            # 计算梯度惩罚
            gradient_penalty = compute_gradient_penalty(
                discriminator, real_imgs.data, fake_imgs.data, device
            )
            
            # WGAN-GP的判别器损失
            d_loss = -torch.mean(real_validity) + torch.mean(fake_validity) + lambda_gp * gradient_penalty
            
            d_loss.backward()
            d_optimizer.step()
            
            # 每n_critic次判别器更新后更新一次生成器
            if i % n_critic == 0:
                # -----------------
                #  训练生成器
                # -----------------
                g_optimizer.zero_grad()
                
                # 生成新的假图像
                z = torch.randn(real_imgs.size(0), latent_dim).to(device)
                fake_imgs = generator(z)
                fake_validity = discriminator(fake_imgs)
                
                # WGAN的生成器损失
                g_loss = -torch.mean(fake_validity)
                
                g_loss.backward()
                g_optimizer.step()
                
            # 每100个批次打印一次信息
            if i % 100 == 0:
                print(
                    f"[Epoch {epoch}/{n_epochs}] [Batch {i}/{len(dataloader)}] "
                    f"[D loss: {d_loss.item():.4f}] [G loss: {g_loss.item():.4f}]"
                )

# 完整的WGAN-GP训练示例
def main():
    # 超参数
    latent_dim = 100
    img_size = 28
    channels = 1
    batch_size = 64
    n_epochs = 50
    n_critic = 5  # 判别器更新次数/生成器更新次数
    lambda_gp = 10  # 梯度惩罚权重
    lr = 0.0002
    b1, b2 = 0.5, 0.999  # Adam优化器的beta参数
    
    # 设备配置
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # 加载MNIST数据集
    transform = transforms.Compose([
        transforms.Resize(img_size),
        transforms.ToTensor(),
        transforms.Normalize([0.5], [0.5])
    ])
    
    mnist_dataset = datasets.MNIST(
        root="./data",
        train=True,
        download=True,
        transform=transform
    )
    
    dataloader = DataLoader(
        mnist_dataset,
        batch_size=batch_size,
        shuffle=True
    )
    
    # 初始化生成器和判别器
    generator = Generator(latent_dim, img_size, channels).to(device)
    discriminator = Discriminator(img_size, channels).to(device)
    
    # 初始化优化器
    g_optimizer = optim.Adam(generator.parameters(), lr=lr, betas=(b1, b2))
    d_optimizer = optim.Adam(discriminator.parameters(), lr=lr, betas=(b1, b2))
    
    # 训练模型
    train_wgan_gp(
        dataloader, latent_dim, n_critic, lambda_gp, 
        generator, discriminator, g_optimizer, d_optimizer, device, n_epochs
    )
    
    # 保存模型
    torch.save(generator.state_dict(), "wgan_gp_generator.pth")
    torch.save(discriminator.state_dict(), "wgan_gp_discriminator.pth")

if __name__ == "__main__":
    main()

以上代码展示了WGAN-GP的核心实现，特别是梯度惩罚的计算部分。关键步骤包括：

1. 在真实样本和生成样本之间创建随机插值点
2. 计算判别器关于这些插值点的梯度
3. 计算梯度范数
4. 对梯度范数与1的差值进行惩罚

7. WGAN-GP训练流程图

让我们通过Mermaid流程图更直观地理解WGAN-GP的训练过程：

8. WGAN与WGAN-GP的对比

让我们通过表格比较标准GAN、WGAN和WGAN-GP：

正如表格所示，WGAN-GP在大多数指标上都优于原始WGAN和标准GAN，特别是在训练稳定性和模式多样性方面。

9. WGAN-GP的超参数敏感性分析

WGAN-GP相比原始GAN大大降低了对超参数的敏感性，但仍有几个关键参数需要调整：

1. λ (lambda_gp)：梯度惩罚的权重，通常设为10
2. n_critic：每更新一次生成器，判别器更新的次数，通常为5
3. 学习率：WGAN-GP对学习率的敏感性低于原始GAN，但仍需合理设置

让我们看一下不同λ值对模型性能的影响：

10. 代码运行结果分析

运行上面的WGAN-GP代码后，我们可以观察到以下现象：

1. 判别器损失：理论上应该收敛到0附近，表示真实分布和生成分布之间的Wasserstein距离很小
2. 生成器损失：应该是一个负值，并逐渐接近0
3. 训练稳定性：与标准GAN相比，损失曲线应该更加平滑，没有剧烈波动
4. 生成质量：随着训练进行，生成图像的质量应该稳步提高

以下是典型的WGAN-GP训练损失曲线示例：

[Epoch 0/50] [Batch 0/938] [D loss: -0.9876] [G loss: 0.5432]
[Epoch 0/50] [Batch 100/938] [D loss: -0.3456] [G loss: -0.1234]
[Epoch 0/50] [Batch 200/938] [D loss: -0.2345] [G loss: -0.3456]
...
[Epoch 49/50] [Batch 900/938] [D loss: -0.0123] [G loss: -0.0234]

可以看到，判别器损失和生成器损失在训练过程中逐渐稳定，这是WGAN-GP成功训练的标志。

模式坍塌问题及解决方案

接下来，让我们转向GAN训练中的另一个关键问题：模式坍塌（Mode Collapse）。

10.1. 什么是模式坍塌？

模式坍塌是指生成器只学会产生有限种类的样本，无法覆盖真实数据分布的多样性。直观地说，就是生成器"偷懒"了，找到了几个能够"欺骗"判别器的样本，然后一直生成这些样本。

例如，在生成手写数字时，模式坍塌的模型可能只会生成看起来像"1"和"7"的数字，而忽略其他数字。

10.2. 模式坍塌的原因

模式坍塌主要有以下几个原因：

1. 生成器优化目标的局限性：标准GAN的生成器只关注"欺骗"判别器，而不直接关注多样性
2. 判别器能力不足：如果判别器无法区分不同的真实样本模式，生成器就没有动力生成多样化样本
3. 训练不平衡：判别器和生成器之间的能力不平衡可能导致坍塌
4. 优化过程中的动态：交替优化过程可能导致振荡或收敛到局部最优解

10.3. 无监督生成中的模式坍塌

在无监督生成（如标准GAN）中，模式坍塌问题尤为严重。因为没有额外信息指导生成器覆盖不同模式，生成器很容易找到"最简单"的方式来欺骗判别器。

例如，假设我们正在生成人脸图像。无监督GAN可能会发现生成某种特定类型的面部特征（比如微笑的白人男性）最容易欺骗判别器，因此会重复生成这类图像，而忽略其他种族、性别或表情的多样性。

10.4. 条件生成中的模式坍塌

条件生成对抗网络（Conditional GAN）通过引入额外的条件信息（如类别标签）来指导生成过程。这种额外信息可以帮助减轻模式坍塌问题，但并不能完全解决它。

在条件生成中，模式坍塌通常表现为每个条件类别内部的多样性不足。例如，在条件生成手写数字的任务中，虽然模型可能能够生成所有10个数字类别，但每个类别内部的多样性（如不同的书写风格）可能很有限。

10.5. 条件生成与无监督生成的模式坍塌对比

让我们通过表格比较条件生成与无监督生成在模式坍塌方面的差异：

10.6. 条件GAN与无监督GAN的对比流程图

如上图所示，条件GAN的关键区别在于将条件标签作为生成器和判别器的额外输入。这种方式可以有效缓解模式坍塌问题，因为它强制生成器学习针对不同条件的不同模式。

10.7. WGAN-GP对模式坍塌的改善

前面我们详细讨论了WGAN-GP的梯度惩罚机制，它不仅提高了训练稳定性，还有助于减轻模式坍塌问题。让我们看看为什么WGAN-GP能够改善模式坍塌：

1. 更平滑的梯度：梯度惩罚确保了判别器的梯度不会消失或爆炸，为生成器提供更稳定、信息更丰富的梯度信号
2. 更好的距离度量：Wasserstein距离比JS散度更适合度量不重叠分布间的距离，鼓励生成器探索真实数据分布的全部模式
3. 平衡的训练动态：通过梯度惩罚，判别器能力不至于过强，生成器有足够的机会学习多样的模式
4. 改进的优化过程：通过避免判别器过拟合，WGAN-GP能够减少优化过程中的振荡

一项实验研究表明，在同样的条件下，WGAN-GP比标准GAN能够生成更多样化的样本，模式覆盖率也更高。

10.8. 混合方法：条件WGAN-GP

结合条件生成和WGAN-GP的优势，我们可以构建条件WGAN-GP来更有效地解决模式坍塌问题。下面是实现条件WGAN-GP的关键代码片段：

import torch
import torch.nn as nn
import torch.autograd as autograd

# 条件WGAN-GP的梯度惩罚计算
def compute_gradient_penalty(D, real_samples, fake_samples, labels, device):
    """计算条件WGAN-GP的梯度惩罚"""
    # 批次大小
    batch_size = real_samples.size(0)
    # 随机插值系数
    alpha = torch.rand(batch_size, 1, 1, 1).to(device)
    # 在真实样本和生成样本之间进行插值
    interpolates = (alpha * real_samples + (1 - alpha) * fake_samples).requires_grad_(True)
    # 计算插值点的判别器输出 (注意这里传入标签)
    d_interpolates = D(interpolates, labels)
    # 创建与d_interpolates形状相同的全1张量
    fake = torch.ones(d_interpolates.size()).to(device)
    # 计算梯度
    gradients = autograd.grad(
        outputs=d_interpolates,
        inputs=interpolates,
        grad_outputs=fake,
        create_graph=True,
        retain_graph=True,
        only_inputs=True,
    )[0]
    # 计算梯度的范数
    gradients = gradients.view(gradients.size(0), -1)
    gradient_norm = gradients.norm(2, dim=1)
    # 计算梯度惩罚 (||∇D(x̂)||_2 - 1)²
    gradient_penalty = ((gradient_norm - 1) ** 2).mean()
    return gradient_penalty

# 条件WGAN-GP训练循环的一部分
def train_conditional_wgan_gp_step(real_imgs, labels, latent_dim, n_classes, lambda_gp,
                                  generator, discriminator, g_optimizer, d_optimizer, device):
    batch_size = real_imgs.size(0)
    
    # -----------------
    #  训练判别器
    # -----------------
    d_optimizer.zero_grad()
    
    # 采样噪声和标签
    z = torch.randn(batch_size, latent_dim).to(device)
    gen_labels = torch.randint(0, n_classes, (batch_size,)).to(device)
    
    # 生成假图像
    fake_imgs = generator(z, gen_labels)
    
    # 计算真实图像、生成图像的判别器输出
    real_validity = discriminator(real_imgs, labels)
    fake_validity = discriminator(fake_imgs.detach(), gen_labels)
    
    # 计算梯度惩罚
    gradient_penalty = compute_gradient_penalty(
        discriminator, real_imgs.data, fake_imgs.data, labels, device
    )
    
    # WGAN-GP的判别器损失
    d_loss = -torch.mean(real_validity) + torch.mean(fake_validity) + lambda_gp * gradient_penalty
    
    d_loss.backward()
    d_optimizer.step()
    
    # -----------------
    #  训练生成器 (每n_critic次判别器更新后)
    # -----------------
    
    g_optimizer.zero_grad()
    
    # 生成新的假图像
    z = torch.randn(batch_size, latent_dim).to(device)
    gen_labels = torch.randint(0, n_classes, (batch_size,)).to(device)
    fake_imgs = generator(z, gen_labels)
    fake_validity = discriminator(fake_imgs, gen_labels)
    
    # WGAN的生成器损失
    g_loss = -torch.mean(fake_validity)
    
    g_loss.backward()
    g_optimizer.step()
    
    return d_loss.item(), g_loss.item()

条件WGAN-GP结合了两种方法的优势：

1. 条件生成通过标签信息确保覆盖全部类别
2. WGAN-GP的梯度惩罚机制提高训练稳定性
3. Wasserstein距离帮助生成器学习多样的模式
4. 条件和梯度惩罚共同作用，显著减轻模式坍塌

10.9. 评估模式坍塌的方法

如何客观地评估模式坍塌的严重程度呢？以下是一些常用方法：

1. 多样性指标：

   • Inception Score (IS)：评估生成图像的质量和多样性
   • Fréchet Inception Distance (FID)：度量真实分布和生成分布之间的相似性
   • 多样性得分 (LPIPS)：评估生成样本间的感知差异

2. 覆盖率指标：

   • 支撑模式数：生成模型能够产生的不同模式数量
   • 生成分布的熵：更高的熵表示更多样的分布
   • 类别覆盖率：在条件生成环境中，评估覆盖不同类别的能力

3. 可视化方法：

   • t-SNE或UMAP降维：观察生成样本在特征空间中的分布
   • 样本网格：为不同条件/噪声生成样本并排列为网格查看多样性

10.10. 无监督与条件生成的模式坍塌实例分析

以下是一个无监督GAN与条件GAN在MNIST数据集上的模式坍塌对比：

总结：梯度惩罚与模式坍塌的关系

在本文的第一部分中，我们深入探讨了Wasserstein GAN的梯度惩罚机制以及模式坍塌问题。关键要点包括：

1. WGAN-GP的梯度惩罚机制是对原始WGAN中权重裁剪的改进，通过惩罚判别器梯度范数偏离1的行为，更优雅地实现Lipschitz约束，提高训练稳定性。

2. 模式坍塌是GAN训练中的常见问题，表现为生成器只产生有限种类的样本，无法覆盖真实数据分布的多样性。

3. 无监督生成中的模式坍塌通常更严重，可能完全缺失某些类别的样本，而条件生成通过引入标签信息，能够在一定程度上缓解这个问题，至少确保覆盖所有类别。

4. WGAN-GP通过改进的距离度量和梯度机制，能够帮助生成器学习多样的模式，减轻模式坍塌问题。

5. 条件WGAN-GP结合了条件生成和WGAN-GP的优势，是解决模式坍塌的有效方法。

通过对比表格和流程图，我们清晰地看到了各种方法在处理模式坍塌问题上的效果差异。了解这些机制和差异，对于设计和训练高质量的生成模型至关重要。

清华大学全五版的《DeepSeek教程》完整的文档需要的朋友，关注我私信：deepseek 即可获得。

怎么样今天的内容还满意吗？再次感谢朋友们的观看，关注GZH：凡人的AI工具箱，回复666，送您价值199的AI大礼包。最后，祝您早日实现财务自由，还请给个赞，谢谢！