如果你也经历过GAN训练中那种玄学式的崩溃——生成器输出一堆噪点、loss曲线像心电图一样乱跳、换一个随机种子结果就天差地别——那么WGAN-GP和WGAN-Div这两套方案应该能帮你从泥潭里爬出来。我最近用Keras在MNIST和CIFAR-10上把这两个模型完整搭了一遍都能稳定生成图像。这篇文章就把从原理到代码再到调参的整个过程记录下来给想动手复现的人一个可以直接抄作业的参考。先说明一个前提下面所有代码基于TensorFlow 2.x内置的Keras。如果你还在用独立的Keras 2.xAPI基本通用但GradientTape那一套只存在于tf.keras里建议直接切换到TensorFlow 2.x省掉版本兼容的折腾。1. 为什么绕不开WGAN这条路线1.1 经典GAN的病灶JS散度在偷懒传统GAN里判别器输出的是一个概率用交叉熵做loss本质上是在最小化生成分布和真实分布之间的JS散度。这个度量方式有个致命问题当两个分布的重叠区域非常小的时候JS散度会趋于一个常数log2梯度几乎为0。举个例子。假设真实数据分布是二维平面上的一条细线生成器刚开始只能撒出一些离散的噪点。这两组点在绝大多数位置都不重合判别器很容易就把它们分得明明白白。此时生成器几乎收不到任何有用的梯度信号不管怎么调参数输出的还是那堆噪点。这就是大家常说的“判别器太强生成器被锤死”。实际操作中的表现就是判别器的loss快速降到接近0生成器的loss却不怎么动生成的图片越来越像噪声。很多人第一反应是调判别器的学习率但根子在于度量方式本身就给不出有效的“距离感”。1.2 Wasserstein距离换了个度量方式WGAN的核心改动是把判别器换成了“critic”输出不再是一个概率而是一个实数得分。它优化的目标也从JS散度换成Wasserstein距离也就是常说的“推土机距离”。这个距离的直觉很好理解假设真实分布是一堆土生成分布是另一堆土Wasserstein距离就是把这堆土搬到另一堆土的最小运输成本。即使两堆土完全不重叠只要它们位置不同这个成本就是一个正的有限值梯度就不会消失。这就解决了经典GAN里“距离太远就无梯度”的问题。所以WGAN的critic要做的事情不是判断“真”“假”而是给真实样本打高分给生成样本打低分。对应到代码里就是去掉最后一层的sigmoid输出一个raw score。这是WGAN与传统GAN最直观的区别。1.3 权重裁剪的副作用WGAN为了保证Wasserstein距离能算得下去要求critic的函数满足1-Lipschitz条件也就是输出对输入的改变不能过于敏感。论文作者最初用的方法是权重裁剪把critic的每一层权重强行限制在一个小范围内比如[-0.01, 0.01]。这个方法实现起来非常简单但副作用很明显。权重被强行裁剪后critic的表达能力被压缩得很厉害很多情况下它只能学到非常粗糙的判别函数。我一开始照着原始WGAN代码跑发现生成的图像总有一种“模糊感”细节上不去。后来才知道这是权重裁剪带来的典型问题critic太弱给生成器提供的梯度质量也差。WGAN-GP和WGAN-Div就是在这个基础上提出的两种改进。它们的核心思路一模一样不再裁剪权重而是通过给损失函数加一个“梯度惩罚项”直接约束critic对输入的梯度从而满足Lipschitz条件。2. 实验整体设计先定结构再谈损失2.1 数据准备与归一化我用两个数据集做验证MNIST用来快速验证收敛逻辑CIFAR-10用来验证彩色、复杂图像上的生成效果。两个数据集都缩放到[-1, 1]区间生成器最后一层用tanh激活输出自然落在[-1, 1]内正好对应输入的数据范围。这里有个容易忽略的细节如果你把图像归一化到[0, 1]生成器最后一层还在用tanh输出那生成图像的像素会被强制压到[-1, 1]看起来会“发灰”。反过来如果数据在[-1, 1]内生成器却用sigmoid输出那图像永远达不到最大亮度和对比度。所以训练前先确认好数据的范围与输出层的激活函数。2.2 生成器结构生成器沿用DCGAN那套结构先通过全连接层把随机噪声映射成一个小尺寸的特征图再用转置卷积逐步上采样到目标尺寸。我在CIFAR-10上用32×32分辨率从8×8的特征图开始连续两次stride为2的转置卷积到16×16、32×32最后一层改用普通卷积输出3个通道。def build_generator(latent_dim128): noise Input(shape(latent_dim,)) x Dense(8 * 8 * 256, use_biasFalse)(noise) x BatchNormalization()(x) x LeakyReLU(alpha0.2)(x) x Reshape((8, 8, 256))(x) x Conv2DTranspose(128, (4, 4), strides(2, 2), paddingsame, use_biasFalse)(x) x BatchNormalization()(x) x LeakyReLU(alpha0.2)(x) x Conv2DTranspose(64, (4, 4), strides(2, 2), paddingsame, use_biasFalse)(x) x BatchNormalization()(x) x LeakyReLU(alpha0.2)(x) x Conv2D(3, (3, 3), paddingsame, activationtanh)(x) return Model(noise, x)生成器里每一层都用LeakyReLUalpha取0.2。这个选择是为了防止ReLU死掉——GAN训练过程中生成器的输入分布变化很大ReLU一旦被压到负区间就可能一直输出0导致梯度断流。LeakyReLU可以保证梯度至少有一条通道能流回去。BN层在生成器里很重要能显著稳定训练。但注意调用生成器时如果处于训练阶段一定要写generator(noise, trainingTrue)让BN更新批次统计量。生成图像做推理时再用trainingFalse否则你会发现同一个噪声向量每次生成的结果不一样原因是BN在训练模式下用了当前batch的均值方差。2.3 Critic结构critic的结构与生成器镜像但有几处关键差异不使用BN不使用sigmoid最后一层直接输出一个标量。def build_critic(img_shape(32, 32, 3)): img Input(shapeimg_shape) x Conv2D(64, (4, 4), strides(2, 2), paddingsame)(img) x LeakyReLU(alpha0.2)(x) x Conv2D(128, (4, 4), strides(2, 2), paddingsame)(x) x LeakyReLU(alpha0.2)(x) x Conv2D(256, (4, 4), strides(2, 2), paddingsame)(x) x LeakyReLU(alpha0.2)(x) x Flatten()(x) x Dropout(0.2)(x) out Dense(1)(x) return Model(img, out)为什么critic不用BN梯度惩罚的计算要求在插值样本上对输入求导而BN在训练时引入批内统计量使输出不再只依赖当前单个输入相当于在批维度上产生了信息泄露。这样一来“对单个输入的梯度”就变得不再纯粹惩罚项也会受到污染。很多第一次复现WGAN-GP的人把BN留在critic里结果训练不稳定多半是这个原因。如果卷积下采样想保留更多信息也可以不加Dropout。我加了一层Dropout(0.2)是防止critic对训练集过度敏感实测对稳定性有一点帮助但不是必须的。2.4 训练超参整套实验的关键超参如下latent_dim 128batch_size 64critic更新次数n_critic 5MNIST上我试过1也能收敛优化器Adamlr 1e-4beta_1 0.5WGAN-GP的lambda 10WGAN-Div的k 1p 2起步Adam的beta_1取0.5而不是默认的0.9这里说下原因。GAN的loss曲面本身波动剧烈如果动量累积过大梯度更新方向会“刹不住车”导致训练震荡。beta_1降低以后动量衰减更快更新方向更追随当前梯度稳定性会好很多。n_critic是一个可以调的空闲参数。理论上critic需要先训练充分才能给生成器提供准确的Wasserstein梯度。但如果critic太强、生成器太弱同样会出现生成器完全跟不上的问题。我在小数据集上的经验是n_critic1到5都能跑数据集越小、越简单n_critic就可以越小。3. WGAN-GP梯度惩罚的完整实现3.1 插值样本与梯度惩罚的计算WGAN-GP的核心是构造一个“插值样本”让critic在这个样本上的梯度范数尽量接近1。所谓插值样本就是从一对真实样本和生成样本的连线上随机取一个点alpha * real (1 - alpha) * fakealpha是[0,1]间的随机数。为什么要在连线上采样因为Lipschitz条件要求函数的梯度在全空间都有界但真正计算所有位置不可行。经验上真实样本和生成样本中间过渡区域是最容易违反Lipschitz条件的地方在连线上做惩罚性价比最高。这个思想不是严格的数学证明但实验效果非常好。def gradient_penalty(critic, real, fake): alpha tf.random.uniform((tf.shape(real)[0], 1, 1, 1), 0.0, 1.0) interp alpha * real (1.0 - alpha) * fake with tf.GradientTape() as tape: tape.watch(interp) pred critic(interp, trainingTrue) grad tape.gradient(pred, interp) norm tf.sqrt(tf.reduce_sum(tf.square(grad), axis[1, 2, 3])) gp tf.reduce_mean((norm - 1.0) ** 2) return gp这里有个初学者容易出错的地方tf.GradientTape()里除了要计算critic的输出还要显式调用tape.watch(interp)否则TensorFlow不会自动记录对interp的梯度。很多报错“gradient is None”都是因为这个。3.2 损失函数与训练循环WGAN-GP中critic的损失不是交叉熵而是三部分之和对真实样本的输出取平均对生成样本的输出取平均注意方向再加上梯度惩罚。生成器的损失更简单目标是让生成样本在critic那里的得分尽量高所以取负号。with tf.GradientTape() as tape: fake_imgs generator(noise, trainingTrue) fake_imgs tf.stop_gradient(fake_imgs) d_real critic(real_imgs, trainingTrue) d_fake critic(fake_imgs, trainingTrue) gp gradient_penalty(critic, real_imgs, fake_imgs) d_loss tf.reduce_mean(d_fake) - tf.reduce_mean(d_real) lambda_gp * gp grads tape.gradient(d_loss, critic.trainable_variables) d_optimizer.apply_gradients(zip(grads, critic.trainable_variables))注意tf.stop_gradient(fake_imgs)这一行。如果不加GradientTape会把生成器的完整计算图记录下来。虽然我们只对critic参数求梯度但这部分记录会白白消耗大量显存训练速度也会变慢。加上stop_gradient之后tape只追踪从fake_imgs到d_fake这段计算内存占用小很多。生成器的更新循环with tf.GradientTape() as tape: fake_imgs generator(noise, trainingTrue) d_fake critic(fake_imgs, trainingTrue) g_loss -tf.reduce_mean(d_fake) grads tape.gradient(g_loss, generator.trainable_variables) g_optimizer.apply_gradients(zip(grads, generator.trainable_variables))整个训练流程是先固定生成器更新critic n_critic次再固定critic更新生成器1次。critic的更新目标是拉开真实样本得分和生成样本得分的差距生成器的更新目标是拉高自己样本的得分。两者在d_fake这个点上形成对抗。3.3 训练循环里的三个细节第一个细节是lambda的取值。WGAN-GP论文给出的经典值就是10这个值在大多数任务上都不需要改。lambda太小时Lipschitz约束形同虚设训练不稳定lambda太大时critic会把几乎所有注意力都放在满足梯度约束上对真实样本和生成样本的区分能力下降生成图像会变得模糊。如果调节后有类似问题优先检查它。第二个细节是critic的最后一层不要初始化偏置为0或权重过大。建议使用默认的Glorot均匀初始化如果发现训练一开始critic输出就很大可以考虑用小标准差初始化最后一层。第三个细节是d_loss的数值范围。WGAN训练中critic的loss是一个不断变化的负数不会像传统判别器那样收敛到0附近。我见过很多人看到loss是负数就以为训练出了问题其实这是正常的。WGAN-GP中d_loss稳步下降、g_loss围绕一个负值波动只能说明损失函数这一路没有崩溃真正的判断标准还是生成图像本身。4. WGAN-Div从固定约束到可调散度4.1 WGAN-Div的原理与损失形式WGAN-Div站在WGAN-GP的肩膀上进一步把“梯度惩罚”从一种工程技巧上升成了理论上的“散度”定义。它引入两个超参数k和p惩罚项写成了k * mean(|grad_norm - k|^p)。这里要交代清楚一个细节WGAN-Div论文原文里的惩罚项写法与工程代码并不完全一致。原文更接近k * E[(k - grad_norm)^p]没有先取绝对值当梯度范数大于k时惩罚项可能变号。实际测试中这种写法在部分p取值下会引入不稳定因素。所以我下面给出的代码采用开源社区里更通用的稳定版本对grad_norm - k取绝对值再求p次幂。如果你想严格对齐论文可以自己去掉绝对值对比效果。def wdiv_penalty(critic, real, fake, k1.0, p2.0): alpha tf.random.uniform((tf.shape(real)[0], 1, 1, 1), 0.0, 1.0) interp alpha * real (1.0 - alpha) * fake with tf.GradientTape() as tape: tape.watch(interp) pred critic(interp, trainingTrue) grad tape.gradient(pred, interp) norm tf.sqrt(tf.reduce_sum(tf.square(grad), axis[1, 2, 3]) 1e-12) penalty k * tf.reduce_mean(tf.abs(norm - k) ** p) return penalty加1e-12是为了防止梯度范数为0时开根号出现数值不稳。然后critic的损失函数变成d_loss tf.reduce_mean(d_fake) - tf.reduce_mean(d_real) wdiv_penalty(critic, real_imgs, fake_imgs, k1.0, p2.0)乍一看和WGAN-GP几乎一样区别只在惩罚项的形式。WGAN-GP是强制梯度范数等于1WGAN-Div是强制梯度范数等于k并且用p次幂控制惩罚力度。当k1、p2时两者非常接近可以互相替代。4.2 与WGAN-GP的代码差异把WGAN-GP换成WGAN-Div代码改动量其实非常小。生成器结构、critic结构、训练循环框架全部可以复用只替换梯度惩罚函数。这也是我推荐在同一个工程里同时保留两套实现的原因切换成本极低。WGAN-Div训练循环中生成器的loss不变还是-mean(d_fake)。唯一要注意的是训练critic时同样要用stop_gradient隔离生成器扩大提高内存效率。4.3 p值对训练的影响从实际效果看p值比k值更值得玩味。p2时惩罚项对“梯度范数偏离k”施加的是二次惩罚类似WGAN-GPp增大到4或6意味着梯度范数偏差越大受到的惩罚越重。这意味着critic会更努力地让梯度范数保持在k附近对局部梯度的约束更严格。论文里的实验偏向于用较大的p比如p6认为这让critic对空间各处的Lipschitz约束更均匀。我自己的测试结论是p2已经够用p6在某些数据集上能略微提升生成图像的多样性但代价是更容易出现数值波动尤其是训练中期。p的调节不要一上来就拉满先在p2跑通再逐档对比。k值的作用更像一个“目标梯度范数”。k越小critic输出对输入的敏感度越低生成器收到的梯度相对温和k越大critic可以更“尖锐”地响应输入变化生成器能获得更强的修正信号但也更容易震荡。我的经验是k取1或2都可以优先从1开始。5. 实测对比与调参笔记5.1 两个模型的训练曲线差异把两个模型都训练20个epoch之后我记录到一些有意思的现象。WGAN-GP的critic loss下降非常平稳基本是一条斜率稳定的曲线生成器loss则在一个负值附近小幅波动。WGAN-Div的曲线形态相似但绝对值大小和波幅会随p变化p越大critic loss的数值抖动越明显。这里必须强调不要试图用critic loss的绝对值去横向对比两个模型的优劣。WGAN的loss是有意义的“分布距离估计值”但它受梯度惩罚项的影响较大改变罚项形式后数值就没有可比性了。判断模型好坏一看出图效果二看是否出现模式崩塌三看训练是否持续平稳。5.2 生成质量与稳定性对比在MNIST上两个模型最终生成的数字都清晰可辨几乎看不出明显差别。WGAN-GP从第8个epoch左右就能看到像样的数字轮廓WGAN-Div在k1、p2时速度类似换成k1、p6后收敛略慢一点但生成样本的多样性感觉稍好一些。在CIFAR-10上两个模型的差别开始显现。WGAN-GP生成的图像整体更“干净”边缘和色块分明WGAN-Div在p较大时生成的图像偶尔能保留更多细节纹理但也更容易出现局部颜色溢出。也就是说WGAN-Div的调节空间更大但它的最优参数需要针对数据集单独摸索不像WGAN-GP那样“开箱即用”。这两条路线我用表格简单对比一下对比维度WGAN-GPWGAN-Div惩罚目标梯度范数等于1梯度范数等于k惩罚形式λ·E[(norm−1)²]k·E[关键超参λ默认10基本不用调k、p需要试验理论统一度面向工程直观从散度定义出发更一般训练稳定性很稳适合首选稳定但依赖参数选择实现复杂度低极低改动一行适用场景绝大多数GAN任务想深挖理论或寻找多样性收益5.3 关键超参调参心得学习率是这场实验里影响最大的超参没有之一。我试过把critic和生成器的学习率都调到3e-4训练很快变得颠簸调回1e-4后马上一路顺风。WGAN系列对学习率比传统GAN更敏感因为critic的梯度是要被“拿来算距离”的步子迈太大等于距离估错。Batch size也不用太大。64是性价比很高的选择。batch太大时梯度惩罚的插值样本会更密集理论上更接近理论要求但实际中显存压力和训练速度都会成问题收益却不成比例。batch太小比如16critic会过度关注当前小批量中的个别样本梯度噪声偏大。n_critic在MNIST上试过1和3收敛速度差不太多在CIFAR-10上n_critic5比1更稳。我的建议是复杂数据集用5简单数据集可以缩到1或3省下来的算力拿去多跑几个epoch更划算。5.4 一个偷懒但有效的调试技巧训练过程中一定要定期保存生成图。我习惯每隔500个step把固定的一组噪声向量丢进生成器生成一张对比图保存下来。这样跑完一遍训练把这些图按顺序连起来看就是一个动画训练过程中有没有崩溃、有没有从噪声慢慢变成图像、有没有突然模式崩塌一眼就能看清楚。只盯loss曲线是看不出这些信息的。我见过loss曲线一路下行但生成图像全部是同一个数字的情况。模式崩塌时critic依然可以把真实样本和“那一类”生成样本区分开loss看起来一切正常。6. 问题排查与避坑记录6.1 训练崩坏的几个典型原因症状可能的坑解决方向loss变成NaN学习率过高、梯度爆炸学习率降到1e-4以下检查插值样本是否有inf图像长期模糊lambda太大 / 数据与输出层不匹配lambda从10往下降或修正激活函数生成内容单一模式崩塌减小batch、降低beta_1、增大latent_dimcritic loss快速归零去掉sigmoid了吗经典GAN的判别头混入了确认critic输出层没有sigmoid图像颜色发灰数据范围与激活函数不一致数据[-1,1]配tanh数据[0,1]配sigmoid6.2 网络结构相关的坑critic里加BN是我踩过最深的坑之一。最开始为了“让训练更稳”按分类网络的习惯在critic每个卷积后加了BN结果训练时好时坏。后来查资料才发现梯度惩罚要求对单个输入样本求梯度BN却让输出依赖于整个batch的统计量等于强行让梯度“去个体化”。去掉BN后训练立刻稳定不少这也直接改变了我的网络设计习惯。生成器里BN的training参数同样是个容易出错的点。如果你在训练循环里调用generator(noise)时不带training参数Keras默认按推理模式执行BN不会更新。短期看loss还在变化但生成器内部统计量一直没动训练效果会越来越差。正确做法是生成器在训练阶段显式传trainingTrue。6.3 数据与版本相关的坑图像归一化一定要在做batch之前完成。如果你用TensorFlow的Dataset API注意map里的归一化操作不能在batch之后才做。我试过一次先batch再map归一化结果归一化作用的维度从(H, W, C)变成了(B, H, W, C)整个训练损失数值都不对修了很久才定位到。版本兼容上TensorFlow 2.4之后tf.keras和独立Keras的混用经常出问题。我建议代码里统一写from tensorflow.keras import layers别一会import keras一会import tf.keras两个库的变量各自独立加载权重和保存模型时会莫名报错。我个人在实际操作中的体会是WGAN-GP和WGAN-Div不是竞争关系而是两组可以随时切换的“梯度约束方案”。如果手里有个图像生成任务想快速出效果直接上WGAN-GP默认参数基本够用如果愿意花时间调参想观察多样性变化再切到WGAN-Div去试不同的k和p。最后再分享一个小技巧实验记录时把每个step保存的生成图统一命名成带step编号的文件训练完了用脚本拼成动画你会非常直观地看到模型是怎么从一片噪声里慢慢长出结构的——那种感觉比任何指标都让人上瘾。