跑过生成模型训练的朋友大概都经历过这种体验网络搭好了数据也准备好了训练跑到一半生成出来的图像要么是一堆重复的“鬼脸”要么干脆变成全黑或者全是噪点loss曲线像心电图一样上蹿下跳。这其实就是经典GAN最常见的两个老毛病训练不稳定、模式崩塌。后来大家发现光换网络结构治标不治本真正的根子在于目标函数——JS散度在分布不重叠的时候几乎给不了生成器任何有意义的梯度。于是WGAN系列应运而生其中WGAN-GP带梯度惩罚的Wasserstein GAN和WGAN-DIV基于散度的WGAN变体是实用性很强、也最容易用Keras落地的两个方案。这篇文章就从损失函数入手带你用KerasTensorFlow 2内置的tf.keras把WGAN-GP和WGAN-DIV完整搭起来并真正跑通图像生成任务。内容覆盖原理、生成器/判别器搭建、两种关键损失函数的实现细节、训练循环写法以及我在实际调参中踩过的坑和最终保留下来的一套稳定参数。适合有一定Keras基础、但还没吃透WGAN系列细节的同学读完可以直接动手复现也能当作日常训练时的速查手册。1. 先弄明白WGAN-GP和WGAN-DIV到底改了什么1.1 原始GAN为什么难训练经典GAN里面生成器和判别器玩的是一个极大极小博弈判别器尽量分辨真假生成器尽量骗过判别器。目标函数写出来核心是真实分布和生成分布之间的JS散度。问题在于当两个分布几乎没有重叠的时候JS散度是一个常数梯度直接消失判别器已经“完美”了生成器却收不到任何有效信息。这就是为什么很多GAN项目训练到中期生成结果突然崩掉——不是学习率太高而是目标函数本身提供不了稳定的学习信号。另外原始GAN里判别器往往训练得太好导致梯度饱和。你观察到loss变成一条平稳直线生成图像却迟迟不进步大概率就是这个问题。实践中大家都用各种trick去缓解比如标签平滑、特征匹配但这些都属于“打补丁”治标不治本。1.2 Wasserstein距离和Lipschitz约束WGAN换了一个度量方式不再用JS散度而是用Wasserstein距离也叫推土机距离、Earth Movers Distance。它衡量的是“把一个分布搬运成另一个分布需要的最小代价”这个距离是连续的并且几乎处处有梯度理论上能有效避免梯度消失。不过Wasserstein距离直接计算很困难WGAN的思路是把它转成一个带约束的优化问题要求判别器这里习惯叫critic满足1-Lipschitz条件也就是函数变化不能太“剧烈”。为了让这个条件成立原版WGAN采用了权重裁剪把每次更新后的权重强行限制在一个小范围内。这个方法简单但会带来两个副作用一是裁剪边界很难选二是权重都被推到边界上网络表达能力下降。所以后来WGAN-GP用“梯度惩罚”取代了权重裁剪。1.3 WGAN-GP与WGAN-DIV的思路差异WGAN-GP的核心是对真实样本和生成样本之间的插值点求判别器输出的梯度然后惩罚这个梯度的范数让它尽量等于1。这个做法比权重裁剪温和得多也让训练过程大幅稳定下来。你后面看到的大部分GAN训练技巧很多都能在WGAN-GP身上找到影子。WGAN-DIV则是从散度divergence的角度做文章。它不强行要求梯度范数处处等于1而是允许梯度范数在一个目标值k附近浮动通过一个柔性惩罚项来约束。换句话说WGAN-GP像是硬性规定“你每一步只能迈这么远”WGAN-DIV更像“你尽量按这个步幅走偶尔偏一点也没关系”。在实际实现中WGAN-DIV的代码和WGAN-GP非常接近惩罚项形式几乎一样只是目标范数的处理更灵活。这也是我推荐两个一起学的原因理解了WGAN-GP再切到WGAN-DIV就是改两三行代码的事情但训练表现却会呈现出不同气质。2. Keras环境准备与模型结构2.1 环境与数据集我建议直接使用TensorFlow 2.xKeras已经内置在tf.keras中不需要单独安装。版本方面2.6到2.10我都试过差异不大新版本也没必要刻意追求最高。Python环境建议3.8以上显卡不是必须的但显存8G以上会舒服很多。CPU训练MNIST也能跑就是慢一些建议先把epoch设小流程通了再换GPU。数据集先拿MNIST练手是最划算的28x28单通道网络不用太深几分钟就能看到效果。不过为了更有意思你可以把目标改成Fashion-MNIST难度稍微高一点生成结果的纹理细节更多。想要彩色图CIFAR-10、CelebA都可以但要把图像统一resize到64x64并且归一化到[-1, 1]。这一点很关键生成器的输出层用tanh天然输出[-1, 1]如果数据还在[0, 1]范围匹配就错位了。2.2 生成器搭建细节生成器的输入是一个随机噪声向量通常是128维从标准正态分布采样。第一层用Dense把噪声映射到足够大的特征图尺寸比如7x7x128然后通过BatchNormalization和LeakyReLU激活再逐步用Conv2DTranspose上采样。MNIST从7x7上采样两次就能到28x2864x64的彩色图则一般从8x8开始上采样三次到64x64。我用Keras写生成器的习惯是先写一个函数方便复用和打印结构def build_generator(latent_dim128): z keras.Input(shape(latent_dim,)) x layers.Dense(7 * 7 * 128)(z) x layers.BatchNormalization()(x) x layers.LeakyReLU(0.2)(x) x layers.Reshape((7, 7, 128))(x) x layers.Conv2DTranspose(128, (4, 4), strides2, paddingsame)(x) x layers.BatchNormalization()(x) x layers.LeakyReLU(0.2)(x) x layers.Conv2DTranspose(64, (4, 4), strides2, paddingsame)(x) x layers.BatchNormalization()(x) x layers.LeakyReLU(0.2)(x) x layers.Conv2DTranspose(1, (4, 4), strides1, paddingsame, activationtanh)(x) return keras.Model(z, x, namegenerator)注意两点。第一生成器里用BatchNormalization是常态它能让上采样过程更稳定。第二最后一层不加BN只用tanh因为tanh的输出范围决定了图像归一化范围。如果你后面想把图像放大比如64x64变成128x128本质就是多接一个Conv2DTranspose通道数按倍数递减。2.3 判别器搭建细节判别器在WGAN系列里更准确叫法应当是“批评者”critic因为它输出的是一个实数评分而不是真/假概率。网络主干用Conv2D卷积逐步下采样激活函数用LeakyReLU不用Sigmoid。这一点和普通GAN不同别搞混。我在WGAN-GP的判别器里基本不用BatchNormalization原因在于BN会把同一batch内的样本统计关联起来而梯度惩罚是逐样本计算的BN的存在会让插值点梯度的计算变得不准确。如果你的显存允许可以用LayerNormalization或InstanceNormalization替代但最简单的方案就是不用归一化层多调一下学习率也能稳住。def build_discriminator(img_shape(28, 28, 1)): x_in keras.Input(shapeimg_shape) x layers.Conv2D(64, (4, 4), strides2, paddingsame)(x_in) x layers.LeakyReLU(0.2)(x) x layers.Conv2D(128, (4, 4), strides2, paddingsame)(x) x layers.LeakyReLU(0.2)(x) x layers.Conv2D(256, (4, 4), strides2, paddingsame)(x) x layers.LeakyReLU(0.2)(x) x layers.Flatten()(x) x layers.Dense(1)(x) return keras.Model(x_in, x, namediscriminator)3. 两类关键损失函数的实现3.1 共同基础随机插值样本与梯度求取不管WGAN-GP还是WGAN-DIV都需要对“插值样本”求梯度。所谓插值样本就是在一对真实样本和生成样本之间随机取一个中间点公式是hat_x t * real_x (1 - t) * fake_x这里的t是从0到1均匀分布随机采样的一个标量再扩维到和图像一样的shape。这样做的意义是梯度惩罚不能只在真实数据上做也不能只在生成数据上做而是要在两者连线的中间区域上做因为Wasserstein距离相关约束要求分布之间的整个“路径”都有良好的梯度行为。用随机插值相当于把约束覆盖到真实分布和生成分布之间的空间远比单独约束两头更合理。在Keras里求梯度只能用tf.GradientTape包一层不能像普通loss那样直接写个函数就完事。下面这段代码就是两种WGAN变体共同的基础也是整个实现里最值得反复看的部分import tensorflow as tf def compute_gradient_penalty(critic, real_x, fake_x, lambda_gp10.0, modewgan-gp): batch_size tf.shape(real_x)[0] t tf.random.uniform((batch_size, 1, 1, 1), 0.0, 1.0) hat_x t * real_x (1.0 - t) * fake_x with tf.GradientTape() as tape: tape.watch(hat_x) hat_out critic(hat_x, trainingTrue) grads tape.gradient(hat_out, hat_x) grad_norm tf.sqrt(tf.reduce_sum(tf.square(grads), axis[1, 2, 3]) 1e-12) if mode wgan-gp: loss lambda_gp * tf.reduce_mean((grad_norm - 1.0) ** 2) else: # wgan-div 的常用简化形式 k 1.0 loss lambda_gp * tf.reduce_mean((grad_norm - k) ** 2) return loss严格来说WGAN-DIV论文里的通用写法中梯度范数项可以带指数p惩罚是梯度范数的2次方与目标值k之间的一种“软约束”。实际工程里大家落地时最常用的是上面这个简化版本把k设为1也就是让插值点附近的梯度范数尽可能落在1附近。相比WGAN-GP它不会对偏离1的梯度产生二次放大以外的额外惩罚边界行为更柔和所以有时候在复杂数据集上反而更容易收敛。3.2 WGAN-GP的梯度惩罚WGAN-GP的完整判别器损失是三部分相加真实样本评分的均值取负加生成样本评分的均值再加梯度惩罚。这里有一个容易绕晕的点对于WGAN来说判别器希望真实样本得分越高越好生成样本得分越低越好所以真实得分项是负号。写成代码就是下面这样。注意 real_pred 是critic对真实图像的打分 fake_pred 是critic对生成图像的打分真实图片的batch和生成图片的batch必须等大。def wgan_gp_loss(critic, real_x, fake_x, lambda_gp10.0): with tf.GradientTape() as tape: real_pred critic(real_x, trainingTrue) fake_pred critic(fake_x, trainingTrue) gp compute_gradient_penalty(critic, real_x, fake_x, lambda_gp, modewgan-gp) d_loss fake_pred - real_pred gp return tf.reduce_mean(d_loss), tf.reduce_mean(fake_pred - real_pred), tf.reduce_mean(gp)这段代码里隐藏了一个关键点梯度惩罚是在对外层GradientTape内部还是外部求的我这里的写法是先算real_pred和fake_pred再单独调用compute_gradient_penalty其中又自己开了GradientTape。compute_gradient_penalty内部并不追踪外层的tape因此外层tape只对critic进行一次梯度计算不会有二次梯度混叠问题。更稳妥的做法是重构训练循环时统一一个GradientTape作用域但如果初学者照着这段拆开写反而更容易理解。3.3 WGAN-DIV的散度惩罚WGAN-DIV的训练逻辑和WGAN-GP几乎一样区别只在梯度惩罚的目标范数的处理方式上。如果你把上面代码里的mode换成“wgan-div”损失值计算就自动切换过去了。不过为了让你能直观对比我单独再把WGAN-DIV的判别器损失写一遍def wgan_div_loss(critic, real_x, fake_x, lambda_gp10.0, k1.0): with tf.GradientTape() as tape: real_pred critic(real_x, trainingTrue) fake_pred critic(fake_x, trainingTrue) div_penalty compute_gradient_penalty(critic, real_x, fake_x, lambda_gp, modewgan-div) d_loss fake_pred - real_pred div_penalty return tf.reduce_mean(d_loss), tf.reduce_mean(fake_pred - real_pred), tf.reduce_mean(div_penalty)使用WGAN-DIV时我一般把lambda_gp设置在5到10之间k通常取1。想要进一步拉开和WGAN-GP的区别也可以试试把k设成0.5或2.0再观察loss变化。比如当生成器的能力明显强于判别器时适当调大k相当于放松对判别器梯度的限制让它更“敏锐”反而能给生成器更大压力。这个方向很多教程没提但你实际跑几次就会体会到k的作用不是理论摆设。3.4 训练主循环与进度保存Keras里训练GAN最灵活的方式不是直接用model.fit而是自己写train_step这样能精细控制判别器和生成器的更新次数。WGAN系列有个通用做法每个训练step里判别器先更新n_critic次生成器再更新1次。这样做的原因是判别器需要训练得足够好才能把Wasserstein距离的梯度传递给生成器。如果判别器太弱它给出的评分没参考价值生成器就会学偏。下面是一个简洁但完整的训练循环骨架用tf.GradientTape分别追踪两个网络再手动应用优化器。tf.function def train_step(real_x): batch_size tf.shape(real_x)[0] z tf.random.normal((batch_size, latent_dim)) with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape: fake_x generator(z, trainingTrue) real_pred discriminator(real_x, trainingTrue) fake_pred discriminator(fake_x, trainingTrue) d_loss, d_wasserstein, d_gp wgan_gp_loss( discriminator, real_x, fake_x, lambda_gp10.0 ) # 换成 wgan_div_loss 就是 WGAN-DIV gen_loss -tf.reduce_mean(fake_pred) gen_grads gen_tape.gradient(gen_loss, generator.trainable_variables) disc_grads disc_tape.gradient(d_loss, discriminator.trainable_variables) gen_optimizer.apply_gradients(zip(gen_grads, generator.trainable_variables)) disc_optimizer.apply_gradients(zip(disc_grads, discriminator.trainable_variables)) return d_loss, gen_loss, d_wasserstein, d_gp注意这里的gen_loss是负的fake_pred均值。因为生成器希望自己生成的图像在判别器那里得分越来越高等价于让fake_pred变大所以对fake_pred求平均后取负。你不需要对生成器做梯度惩罚生成器的更新完全靠判别器给出的梯度来反传。实际训练时每个epoch里遍历全部数据集每个batch内部再循环n_critic次。有一种常见的偷懒写法是只在train_step内部用循环但如果你用tf.function编译里面套Python循环有时会踩XLA重编译的坑。我建议直接在外层循环判别器更新次数for epoch in range(epochs): for batch in dataset: for _ in range(n_critic): # 更新判别器 n_critic 次 passthrough_batch batch d_loss, gen_loss, d_wass, d_gp train_step(passthrough_batch) # 上面train_step里已经包含了生成器更新逻辑如果想让代码更工程化可以把判别器更新和生成器更新拆成两个函数并且记录每一项loss。我通常还会每过200个batch输出一次loss每过500个batch保存一次生成图像这样能随时观察训练状态。4. 参数调节、常见失败与图像质量提升4.1 三张速查表参数、问题、调优训练GAN跟烧菜很像配方对了才会好吃。下面是我跑MNIST和CIFAR-10时最常用的一组参数直接抄问题不大。参数推荐值说明latent_dim128噪声向量维度128起步太少模式单一batch_size64太大会让梯度惩罚计算变慢太小不稳定n_critic5每更新一次生成器前先更新五次判别器lr_disc1e-4判别器学习率WGAN系列建议比普通GAN小lr_gen1e-4生成器学习率与判别器相等即可也可1e-4对2e-4lambda_gp10WGAN-GP默认WGAN-DIV可以降到5beta_10.5Adam的一阶动量GAN常用0.5而非0.9beta_20.9Adam二阶梯度的衰减因子epochs50起步MNIST 30-50轮可看效果CIFAR至少100轮训练中出现最典型的三个问题我也整理成了一张排查表。这张表是我自己踩坑时反复查阅的比看一堆论文里的证明有用得多。现象可能原因处理方向loss变成NaN学习率太高或梯度惩罚权重太大lr降到5e-5lambda_gp降到1~5生成图像全是同一张脸/同一种图案模式崩塌生成器能力弱或噪声维度太低增大latent_dim加卷积层降低判别器学习率图像一直很模糊训练epoch不够、上采样层太少增加通道数、加训练轮次、换更大数据集判别器loss迅速降为负数判别器太强生成器跟不上增大n_critic反而会加剧问题应该减小n_critic到3或1生成图像有明显棋盘格伪影Conv2DTranspose叠加不当适当加入UpSampling2D Conv2D替代转置卷积4.2 逐步排查训练失败的思路如果你发现训练完全不动第一步不是改网络结构而是检查数据归一化。很多人在加载MNIST时只用X_train / 255数据范围是[0,1]而生成器输出是[-1,1]这会导致判别器很容易找到一条“歪路”去区分真假等于告诉生成器“你输出没有任何意义”。把数据改成X_train (X_train.astype(“float32”) - 127.5) / 127.5后训练稳定性会立刻上一个台阶。第二步检查梯度惩罚是否真的生效。你可以在train_step里把d_gp单独解出来打印如果d_gp一直在0附近跳动说明梯度惩罚消失了。常见的坑是使用Keras模型时误设了trainingFalse导致GradientTape没有追踪到模型内部的可训练参数另一个坑是插值样本hat_x没有用tape.watch(hat_x)显式声明导致梯度变成None。这两个点我已经遇到好多次了。还有一点容易被忽略不要对判别器使用Sigmoid激活。WGAN的判别器输出是scalar实数不是概率。如果保留了Sigmoid输出会被压到(0,1)区间真实样本和生成样本的得分差异会变得很小梯度惩罚的实际作用空间也被压缩了。我见过有人把WGAN-GP的判别器改了Sigmoid之后loss一直不下降去掉之后立刻恢复。4.3 从MNIST到高分辨率彩色图的扩展经验MNIST能跑通之后你可以把任务升级到64x64或128x128的彩色图。这一步考验的不是模型复杂度而是训练资源。生成器要增加上采样层比如从8x8上采样到64x64需要3次stride2的转置卷积判别器要相应增加下采样层把分辨率降到更小的特征图最后再接Dense层输出一个标量。通道数建议从64起步逐层翻倍到256或512不要太贪否则显存会先爆炸。我在做128x128的人脸图片生成时发现单独靠WGAN-GP的梯度惩罚还是会有周期性的抖动。这时候有几种立竿见影的改进方法一是调整n_critic判别器比生成器弱一点时把n_critic减到3会减少生成器被“压死”的概率二是给生成器增加跳跃连接或残差块让高频细节更容易保留三是给判别器加一个简单的EMA指数滑动平均生成器虽然代码麻烦一点但出图质量会稳定很多。如果你想同时对比WGAN-DIV和WGAN-GP最公平的方式是保持网络结构、学习率、batch_size完全一致只切换损失函数。我实测下来在MNIST这个小数据集上两者差别不大WGAN-GP收敛稍快在CIFAR-10或更大分辨率上WGAN-DIV的loss抖动更小生成图像的多样性略好一点。这可能是因为WGAN-DIV对梯度约束更宽松给判别器留了更多表达空间从而不容易把生成器逼进死胡同。我个人在实际操作中的体会是不要把WGAN-GP和WGAN-DIV当成两个割裂的算法它们只是同一套“梯度正则化”思想下的两种调节尺度。你在Keras里把惩罚项的写法灵活切换就能快速对比哪种策略更适合自己的数据集。另一个值得分享的小技巧是在训练早期多保存几次模型权重因为WGAN-GP和WGAN-DIV的loss曲线并不能完全反映生成质量有时候loss还在下降图像已经变得很好了反而继续训练会让图像变差。每500个epoch存一次图、每2000个epoch存一次权重是我目前最稳妥的节奏。这个内容后续还可以继续扩展比如把梯度惩罚换成谱归一化或者引入更复杂的评估指标但基础的Keras实现和训练思路已经足够支撑你跑出能看的生成图像了。