别再为GAN训练不稳定发愁了！用PyTorch手把手实现WGAN-GP（含梯度惩罚代码）
实战指南：用PyTorch实现WGAN-GP解决GAN训练不稳定问题
如果你曾经尝试过训练原始GAN模型，大概率会遇到这样的场景：生成器输出的样本要么全是噪声，要么反复生成几张几乎相同的图片。这种训练不稳定和模式崩溃的问题让许多开发者头疼不已。2017年提出的WGAN-GP（带梯度惩罚的Wasserstein GAN）通过引入Wasserstein距离和梯度惩罚机制，显著改善了这些问题。本文将带你从零实现一个完整的WGAN-GP模型，并深入解析其稳定训练的关键技术。
1. 为什么原始GAN如此难以驯服？
在开始WGAN-GP之前，我们需要理解原始GAN的根本问题。传统GAN使用JS散度（Jensen-Shannon divergence）作为衡量生成分布与真实分布差异的指标，这带来了两个致命缺陷：
梯度消失问题：当两个分布没有重叠或重叠部分可以忽略时，JS散度会恒等于log2，导致梯度为零，生成器无法获得有效的更新信号。
模式崩溃：生成器倾向于生成有限的几种样本，而无法覆盖整个数据分布。想象一个生成数字的GAN，它可能只学会生成"1"和"7"，而完全忽略其他数字。
Wasserstein距离（又称推土机距离）的引入完美解决了这些问题。它衡量的是将一个分布"搬移"成另一个分布所需的最小"工作量"，即使两个分布没有重叠，也能提供有意义的梯度。
# Wasserstein距离的直观理解 def wasserstein_distance(P, Q): """ P: 真实分布 Q: 生成分布 返回: 将Q转化为P所需的最小"工作量" """ return min_work(P, Q) # 实际计算远比这复杂2. WGAN-GP的核心创新
WGAN-GP在原始WGAN基础上做出了关键改进：
2.1 从权重裁剪到梯度惩罚
原始WGAN使用权重裁剪（Weight Clipping）来强制判别器满足1-Lipschitz条件，但这会导致：
- 权重被限制在狭窄范围内[-c, c]
- 容易造成梯度消失或爆炸
- 需要精细调整裁剪阈值c
WGAN-GP改用梯度惩罚（Gradient Penalty）技术，直接在损失函数中添加对梯度范数的约束：
L = E[D(x)] - E[D(G(z))] + λE[(||∇D(αx + (1-α)G(z))||₂ - 1)²]其中λ是惩罚系数，通常设为10；α是从均匀分布U[0,1]中采样的随机数。
2.2 判别器架构变化
WGAN-GP的判别器（在WGAN中更准确应称为"批评器"）有几点重要区别：
- 最后一层不使用sigmoid激活，输出是无约束的标量
- 使用线性层而非BatchNorm，避免影响梯度惩罚
- 损失函数直接最大化真实样本与生成样本的评分差
class Critic(nn.Module): def __init__(self, img_shape): super().__init__() self.model = nn.Sequential( nn.Linear(int(np.prod(img_shape)), 512), nn.LeakyReLU(0.2), nn.Linear(512, 256), nn.LeakyReLU(0.2), nn.Linear(256, 1) # 无激活函数 ) def forward(self, img): img_flat = img.view(img.size(0), -1) validity = self.model(img_flat) return validity3. 完整PyTorch实现
下面我们实现一个完整的WGAN-GP模型，以MNIST数据集为例。
3.1 梯度惩罚实现
梯度惩罚是WGAN-GP最关键的组件，需要仔细实现：
def compute_gradient_penalty(critic, real_samples, fake_samples, device): """计算梯度惩罚项""" # 随机插值样本 alpha = torch.rand(real_samples.size(0), 1, 1, 1, device=device) interpolates = (alpha * real_samples + ((1 - alpha) * fake_samples)).requires_grad_(True) # 计算插值样本的判别器输出 d_interpolates = critic(interpolates) # 计算梯度 gradients = torch.autograd.grad( outputs=d_interpolates, inputs=interpolates, grad_outputs=torch.ones_like(d_interpolates, device=device), create_graph=True, retain_graph=True, only_inputs=True )[0] # 计算梯度范数偏离1的惩罚 gradients = gradients.view(gradients.size(0), -1) gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean() return gradient_penalty3.2 训练循环
WGAN-GP的训练过程与原始GAN有所不同，特别是判别器（批评器）需要训练更多次：
# 超参数设置 n_epochs = 100 batch_size = 64 lr = 0.0001 n_critic = 5 # 每次生成器更新前判别器的更新次数 lambda_gp = 10 # 初始化模型 generator = Generator() critic = Critic() optimizer_G = optim.Adam(generator.parameters(), lr=lr, betas=(0.5, 0.9)) optimizer_C = optim.Adam(critic.parameters(), lr=lr, betas=(0.5, 0.9)) for epoch in range(n_epochs): for i, (real_imgs, _) in enumerate(dataloader): # 训练判别器 optimizer_C.zero_grad() # 生成假样本 z = torch.randn(batch_size, latent_dim, device=device) fake_imgs = generator(z) # 计算判别器损失 real_validity = critic(real_imgs) fake_validity = critic(fake_imgs.detach()) gradient_penalty = compute_gradient_penalty(critic, real_imgs.data, fake_imgs.data, device) loss_C = -torch.mean(real_validity) + torch.mean(fake_validity) + lambda_gp * gradient_penalty loss_C.backward() optimizer_C.step() # 每n_critic次训练一次生成器 if i % n_critic == 0: optimizer_G.zero_grad() # 生成器希望判别器对假样本的评分越高越好 fake_validity = critic(fake_imgs) loss_G = -torch.mean(fake_validity) loss_G.backward() optimizer_G.step()4. 训练技巧与可视化
4.1 训练监控
WGAN-GP的一个优势是判别器的损失可以直接反映生成质量。我们可以监控以下指标：
- 判别器损失：应该逐渐收敛而非剧烈波动
- 梯度惩罚项：应保持在合理范围内
- 生成样本质量：定期保存生成的图片观察进展
# 训练过程中可视化 if epoch % 10 == 0: with torch.no_grad(): test_z = torch.randn(16, latent_dim, device=device) generated = generator(test_z).cpu() save_image(generated, f"results/epoch_{epoch}.png", nrow=4, normalize=True)4.2 超参数调优
经过多次实验，我发现以下设置效果较好：
| 超参数 | 推荐值 | 说明 |
|---|---|---|
| 学习率 | 0.0001 | 比常规GAN更小 |
| β1 | 0.5 | Adam优化器参数 |
| β2 | 0.9 | 比常规GAN的0.999更激进 |
| 批量大小 | 64-256 | 太大可能影响梯度惩罚效果 |
| λ | 10 | 梯度惩罚系数 |
| 判别器迭代次数 | 5 | 原始论文推荐5次 |
4.3 与其他GAN变体对比
下表展示了在MNIST数据集上不同GAN变体的表现：
| 模型 | 训练稳定性 | 模式崩溃 | 生成质量 | 训练速度 |
|---|---|---|---|---|
| 原始GAN | 差 | 严重 | 一般 | 快 |
| WGAN | 中等 | 轻微 | 较好 | 中等 |
| WGAN-GP | 好 | 很少 | 优秀 | 较慢 |
| LSGAN | 中等 | 中等 | 较好 | 快 |
5. 进阶应用与扩展
掌握了基础WGAN-GP后，你可以尝试以下进阶方向：
- 条件WGAN-GP：加入类别标签信息，实现可控生成
- 渐进式增长：从低分辨率开始逐步增加生成复杂度
- 风格混合：结合StyleGAN的思想实现多尺度生成
- 跨域转换：应用于图像到图像的转换任务
# 条件WGAN-GP示例 class ConditionalGenerator(nn.Module): def __init__(self, latent_dim, num_classes): super().__init__() self.label_emb = nn.Embedding(num_classes, latent_dim) # 其余结构与普通生成器类似 def forward(self, z, labels): c = self.label_emb(labels) x = torch.cat([z, c], 1) return self.model(x)在实际项目中，WGAN-GP特别适合以下场景：
- 需要高质量、多样化生成结果的场景
- 数据分布复杂、多模态的任务
- 对训练稳定性要求高的生产环境
- 需要定量评估生成质量的科研工作
经过多次实验验证，WGAN-GP在大多数情况下都能提供比原始GAN更稳定、更可靠的训练过程。虽然计算开销稍大，但换来的训练成功率和生成质量提升绝对是值得的。