用PyTorch实战WGAN：彻底解决GAN训练不稳定难题

当你第一次尝试用GAN生成图片时，可能遇到过这样的场景：精心设计的网络结构，调了无数次的超参数，但生成器要么输出一堆噪声，要么反复生成同一张脸。这就是GAN训练中著名的"模式崩塌"现象——而Wasserstein GAN（WGAN）正是为解决这个问题而生。

1. 为什么你的GAN总是训练失败？

传统GAN使用JS散度作为损失函数，这就像用一把刻度不均匀的尺子测量两个分布的距离。当真实图片和生成图片完全没有重叠时（这在训练初期很常见），JS散度会卡在log2这个固定值，导致梯度消失。想象你在爬山，但整座山都是平地——没有坡度指引方向，优化器自然找不到该往哪走。

WGAN的创新在于用Wasserstein距离（又称推土机距离）替代JS散度。这个距离度量直观理解为：把一堆土从一个形状搬成另一个形状所需的最小工作量。即使两个分布没有重叠，它也能给出有意义的距离值。这带来了三个关键优势：

• 训练更稳定 ：Critic（相当于传统GAN的判别器）的输出值与实际生成质量正相关
• 超参数更鲁棒 ：不再需要精心平衡生成器和判别器的训练节奏
• 避免模式崩塌 ：Wasserstein距离天然鼓励生成多样性

# 传统GAN与WGAN的判别器输出对比
传统GAN判别器输出：概率值(0~1) → 容易饱和导致梯度消失
WGAN Critic输出：任意实数 → 提供全局有效的梯度信号

2. WGAN的核心实现技巧

2.1 权重裁剪的魔鬼细节

WGAN通过强制Critic满足1-Lipschitz约束来实现理论保证，原始论文采用简单的权重裁剪（Weight Clipping）。虽然方法简单，但实践中存在几个关键陷阱：

常见错误1：裁剪阈值设置不当

• 阈值过大（如0.1）：约束效果弱，仍然可能出现梯度爆炸
• 阈值过小（如0.001）：网络容量严重受限，生成质量下降

提示：从0.01开始尝试，观察Critic输出的数值范围，理想情况应在[-10,10]区间波动

常见错误2：裁剪频率不足

# 错误的实现方式（只在生成器更新后裁剪）
for p in critic.parameters():
    p.data.clamp_(-0.01, 0.01)

# 正确的做法（每次Critic更新后立即裁剪）
for _ in range(critic_iters):
    # ...计算critic损失并反向传播...
    optimizer_critic.step()
    for p in critic.parameters():  # 立即裁剪！
        p.data.clamp_(-0.01, 0.01)

2.2 更先进的约束方案

虽然权重裁剪简单有效，但后续研究提出了更好的方案：

方法	优点	缺点
权重裁剪	实现简单	可能限制网络表达能力
梯度惩罚(WGAN-GP)	更精确的Lipschitz约束	计算成本增加约30%
谱归一化	理论保证强，训练稳定	实现复杂度较高

# WGAN-GP的梯度惩罚实现示例
def compute_gradient_penalty(critic, real_samples, fake_samples):
    alpha = torch.rand(real_samples.size(0), 1, 1, 1)
    interpolates = (alpha * real_samples + (1-alpha) * fake_samples).requires_grad_(True)
    d_interpolates = critic(interpolates)
    gradients = autograd.grad(
        outputs=d_interpolates,
        inputs=interpolates,
        grad_outputs=torch.ones_like(d_interpolates),
        create_graph=True
    )[0]
    gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean()
    return gradient_penalty

3. 完整PyTorch实现解析

3.1 网络架构设计要点

WGAN对网络结构没有特殊要求，但以下几个设计能显著提升效果：

1. 移除BatchNorm ：原始论文发现BN会干扰Lipschitz约束
2. 使用LeakyReLU ：给Critic保留更多梯度信息（α=0.2效果较好）
3. 避免极端参数初始化 ：使用PyTorch默认初始化即可

class Critic(nn.Module):
    def __init__(self, img_channels=3):
        super().__init__()
        self.model = nn.Sequential(
            nn.Conv2d(img_channels, 64, 4, 2, 1),
            nn.LeakyReLU(0.2),
            nn.Conv2d(64, 128, 4, 2, 1),
            nn.LeakyReLU(0.2),
            nn.Conv2d(128, 256, 4, 2, 1),
            nn.LeakyReLU(0.2),
            nn.Conv2d(256, 512, 4, 2, 1),
            nn.LeakyReLU(0.2),
            nn.Conv2d(512, 1, 4, 1, 0),
            nn.Flatten()  # 输出标量值而非特征图
        )

    def forward(self, x):
        return self.model(x)

3.2 训练流程的关键控制

WGAN要求Critic比Generator训练得更充分，典型设置是5:1的比例：

# 训练循环骨架
for epoch in range(epochs):
    for real_imgs, _ in dataloader:
        # 训练Critic多次
        for _ in range(critic_iters):
            noise = torch.randn(batch_size, z_dim)
            fake_imgs = generator(noise)
            
            # 计算Wasserstein距离
            critic_real = critic(real_imgs).mean()
            critic_fake = critic(fake_imgs.detach()).mean()
            w_distance = critic_real - critic_fake
            
            # 加梯度惩罚时需额外计算
            loss_critic = -w_distance + lambda_gp * gradient_penalty
            
            optimizer_critic.zero_grad()
            loss_critic.backward()
            optimizer_critic.step()
            
            # 如果是原始WGAN则进行权重裁剪
            if use_clipping:
                for p in critic.parameters():
                    p.data.clamp_(-clip_value, clip_value)
        
        # 训练Generator一次
        fake_imgs = generator(noise)
        loss_gen = -critic(fake_imgs).mean()
        
        optimizer_gen.zero_grad()
        loss_gen.backward()
        optimizer_gen.step()

4. 实战调试技巧与可视化监控

4.1 训练过程诊断

WGAN的训练状态可以通过几个指标判断：

• Critic输出值 ：应该在零附近波动，绝对值过大说明约束失效
• 生成样本多样性 ：定期用固定噪声向量生成图片观察变化
• Wasserstein距离 ：理想情况下应随时间逐渐减小

注意：如果Critic损失快速趋近零，可能是著名的"模式坍塌"前兆，此时应降低学习率或增加梯度惩罚系数

4.2 学习率与优化器选择

不同于传统GAN使用Adam，WGAN配合RMSProp往往表现更好：

# 优化器配置建议
optimizer_C = torch.optim.RMSprop(critic.parameters(), lr=5e-5)
optimizer_G = torch.optim.RMSprop(generator.parameters(), lr=5e-5)

# 学习率调度策略
scheduler_C = torch.optim.lr_scheduler.ExponentialLR(optimizer_C, gamma=0.99)
scheduler_G = torch.optim.lr_scheduler.ExponentialLR(optimizer_G, gamma=0.99)

关键参数经验值 ：

• 初始学习率：5e-5 ~ 1e-4
• 批大小：64~128（太大可能影响梯度惩罚效果）
• Critic迭代次数：3~5次/每生成器迭代

在实际项目中，我习惯用TensorBoard记录以下指标：

writer.add_scalar('Loss/critic', loss_critic.item(), global_step)
writer.add_scalar('Loss/generator', loss_gen.item(), global_step)
writer.add_scalar('Metrics/w_distance', w_distance.item(), global_step)
writer.add_histogram('Critic/real_output', critic_real, global_step)
writer.add_histogram('Critic/fake_output', critic_fake, global_step)

WGAN虽然解决了传统GAN的许多痛点，但要获得最佳效果仍需要耐心调试。建议从CelebA等中等复杂度数据集开始，待pipeline稳定后再挑战更高分辨率生成任务。记住，当生成器开始产出有意义的样本时，Critic的损失可能���暂时上升——这是正常现象，说明对抗训练正在起效。