简介：这份源码资源面向深度学习入门者与图像生成爱好者，提供一套基于WGAN-GP算法生成256×256像素动漫头像的完整实现，可用于理解生成对抗网络的训练流程与梯度惩罚机制，并作为二次开发或课程实验的实操起点。压缩包共26个文件，约1.32MB，其中2个Python源文件承担生成器、判别器与训练主逻辑，11个PNG图片展示不同训练阶段生成的头像效果，6个XML与1个iml文件用于IDE项目及代码风格配置，另有readme说明与Git忽略文件辅助项目管理。目前已有306人学习下载。读者可借此掌握WGAN-GP相较传统GAN在缓解模式崩塌、提升生成稳定性方面的具体做法，观察损失函数与梯度惩罚项的实现细节，并参考目录结构快速搭建自己的训练环境，为动漫头像定制、表情包制作等场景提供可复用的代码基础。
1. 从一张 256×256 的动漫头像说起：WGAN-GP 到底解决了什么
你可能遇到过这种场景：手里攒了几千张动漫头像，想训练一个能生成新头像的模型，结果用最朴素的 GAN 跑了几轮，要么生成一堆雪花噪点，要么模式崩塌——所有输出长得一模一样。这不是你的数据有问题，而是原始 GAN 的判别器训练太激进，梯度信号不稳定。WGAN-GP 就是冲着这个痛点来的：它用 Wasserstein 距离替代原始 GAN 的 JS 散度，再叠加梯度惩罚项（Gradient Penalty），让训练过程稳定得多。这个方案的目标很明确——在 256×256 像素这个分辨率下，生成结构完整、风格统一的动漫头像。它适合谁？有基本 PyTorch 使用经验、想跑通一个完整 GAN 训练流程的工程师，以及需要批量生成头像素材的独立开发者。源码层面，核心就是生成器、判别器、梯度惩罚损失和训练循环四块，下面逐层拆开讲。
2. WGAN-GP 的核心机制与 256×256 生成器结构选型
2.1 为什么是 WGAN-GP 而不是原始 GAN 或 WGAN
原始 GAN 的判别器输出概率值，用交叉熵做损失，训练时判别器越强，生成器梯度消失越严重。WGAN 改用 Earth-Mover 距离，判别器不再输出概率而是输出一个实数分数，理论上要求判别器满足 1-Lipschitz 连续性。最初的 WGAN 用权重裁剪来强制这个条件，但裁剪阈值很难调——裁小了梯度消失，裁大了约束失效。WGAN-GP 的改进点在于：不再裁剪权重，而是在损失函数里加一个梯度惩罚项，惩罚判别器对输入梯度的范数偏离 1 的程度。
具体来说，梯度惩罚项长这样：
# 梯度惩罚核心计算 def gradient_penalty(critic, real, fake, device): batch_size, c, h, w = real.shape # 在真假样本之间随机插值 alpha = torch.rand(batch_size, 1, 1, 1).to(device) interpolated = alpha * real + (1 - alpha) * fake interpolated.requires_grad_(True) # 判别器对插值样本打分 score = critic(interpolated) # 计算梯度 gradient = torch.autograd.grad( outputs=score, inputs=interpolated, grad_outputs=torch.ones_like(score), create_graph=True, retain_graph=True, only_inputs=True )[0] # 梯度范数偏离 1 的惩罚 gradient_norm = gradient.view(batch_size, -1).norm(2, dim=1) penalty = ((gradient_norm - 1) ** 2).mean() return penalty这段代码的关键参数是alpha的采样方式——从均匀分布 U[0,1] 中采样，在真实样本和生成样本之间做线性插值。gradient_norm计算的是判别器输出对插值输入的梯度 L2 范数，惩罚项就是让这个范数尽量接近 1。create_graph=True必须开，因为惩罚项本身也要参与反向传播。retain_graph=True是为了后续还能继续用这个计算图。
判别器损失由两部分组成：真实样本分数均值减去生成样本分数均值，再加上梯度惩罚乘以惩罚系数 λ。λ 一般取 10，这是原论文推荐的默认值，实践中 5 到 20 之间都有人用，但 10 是最稳的起点。
2.2 256×256 分辨率下生成器的上采样策略
256×256 不算特别大，但也不能像 64×64 那样随便堆几层转置卷积就完事。我一般用 DCGAN 风格的生成器骨架，但针对 256 分辨率做了调整：从 8×8 的噪声向量出发，经过 5 次上采样到达 256×256。每次上采样用ConvTranspose2d或者Upsample + Conv2d的组合。
import torch.nn as nn class Generator(nn.Module): def __init__(self, z_dim=128, ngf=64): super().__init__() self.net = nn.Sequential( # 输入 z: (z_dim, 1, 1) -> (ngf*8, 4, 4) nn.ConvTranspose2d(z_dim, ngf*8, 4, 1, 0, bias=False), nn.BatchNorm2d(ngf*8), nn.ReLU(True), # (ngf*8, 4, 4) -> (ngf*4, 8, 8) nn.ConvTranspose2d(ngf*8, ngf*4, 4, 2, 1, bias=False), nn.BatchNorm2d(ngf*4), nn.ReLU(True), # (ngf*4, 8, 8) -> (ngf*2, 16, 16) nn.ConvTranspose2d(ngf*4, ngf*2, 4, 2, 1, bias=False), nn.BatchNorm2d(ngf*2), nn.ReLU(True), # (ngf*2, 16, 16) -> (ngf, 32, 32) nn.ConvTranspose2d(ngf*2, ngf, 4, 2, 1, bias=False), nn.BatchNorm2d(ngf), nn.ReLU(True), # (ngf, 32, 32) -> (ngf//2, 64, 64) nn.ConvTranspose2d(ngf, ngf//2, 4, 2, 1, bias=False), nn.BatchNorm2d(ngf//2), nn.ReLU(True), # (ngf//2, 64, 64) -> (ngf//4, 128, 128) nn.ConvTranspose2d(ngf//2, ngf//4, 4, 2, 1, bias=False), nn.BatchNorm2d(ngf//4), nn.ReLU(True), # (ngf//4, 128, 128) -> (3, 256, 256) nn.ConvTranspose2d(ngf//4, 3, 4, 2, 1, bias=False), nn.Tanh() ) def forward(self, z): return self.net(z)这里z_dim=128是潜向量维度，ngf=64是基础通道数。从 4×4 开始，每次转置卷积的kernel_size=4, stride=2, padding=1，输出尺寸翻倍。最后一层输出 3 通道 256×256，用Tanh把像素值压到 [-1, 1]。注意每一层转置卷积后面都接了BatchNorm2d，除了最后一层——最后一层不加 BN 是因为输出要直接映射到像素空间，BN 会破坏颜色分布。
判别器结构基本是生成器的镜像，但不用 BN，改用 LayerNorm 或者 InstanceNorm，因为 WGAN-GP 的梯度惩罚对每个样本独立计算，BN 会引入样本间的耦合。判别器最后输出一个标量分数，不加 Sigmoid。
2.3 训练循环里判别器和生成器的更新比例
WGAN-GP 的一个关键实践是：每更新一次生成器，判别器要更新多次（通常 5 次）。这是因为 Wasserstein 距离的估计需要判别器足够准，判别器欠拟合时生成器拿到的梯度信号是错的。
# 训练循环核心片段 for epoch in range(num_epochs): for i, real_imgs in enumerate(dataloader): real_imgs = real_imgs.to(device) batch_size = real_imgs.size(0) # ---- 训练判别器 n_critic 次 ---- for _ in range(n_critic): z = torch.randn(batch_size, z_dim, 1, 1).to(device) fake_imgs = generator(z).detach() real_score = critic(real_imgs).mean() fake_score = critic(fake_imgs).mean() gp = gradient_penalty(critic, real_imgs, fake_imgs, device) # 判别器损失：最大化 real_score - fake_score，即最小化负值 d_loss = -real_score + fake_score + lambda_gp * gp optimizer_critic.zero_grad() d_loss.backward() optimizer_critic.step() # ---- 训练生成器 1 次 ---- z = torch.randn(batch_size, z_dim, 1, 1).to(device) fake_imgs = generator(z) g_loss = -critic(fake_imgs).mean() optimizer_gen.zero_grad() g_loss.backward() optimizer_gen.step()n_critic=5是原论文的推荐值，lambda_gp=10。优化器用 Adam，学习率 1e-4，betas 设成 (0.5, 0.9)——注意第二个 beta 不要用默认的 0.999，WGAN-GP 对动量项比较敏感，0.9 更稳。生成器那边fake_imgs不需要detach()，因为要回传梯度到生成器。
3. 从零跑通训练：数据准备、参数配置与监控指标
3.1 动漫头像数据集的整理与预处理
数据来源通常是爬取或者公开数据集，但不管哪来的，统一处理成 256×256 的 RGB 图片。我一般用torchvision的ImageFolder配合自定义 transform：
from torchvision import transforms, datasets transform = transforms.Compose([ transforms.Resize((256, 256)), transforms.CenterCrop(256), transforms.ToTensor(), transforms.Normalize([0.5]*3, [0.5]*3) # 压到 [-1, 1] ]) dataset = datasets.ImageFolder(root='./anime_faces', transform=transform) dataloader = torch.utils.data.DataLoader( dataset, batch_size=32, shuffle=True, num_workers=4, drop_last=True )Normalize的均值和标准差都设 0.5，把 [0,1] 的像素值映射到 [-1,1]，和生成器最后一层Tanh的输出范围对齐。drop_last=True是为了避免最后一个 batch 尺寸不固定导致梯度惩罚计算出问题。batch_size=32在 8GB 显存上跑 256×256 基本够用，显存紧张就降到 16。
数据量方面，至少准备 5000 张以上，否则判别器很容易过拟合到训练集，生成器学不到有意义的分布。如果数据不够，可以做水平翻转增强，但不要做旋转或裁剪——动漫头像的构图通常是对称的，旋转会引入不自然的样本。
3.2 关键超参数表与显存占用估算
| 参数 | 推荐值 | 说明 |
|---|---|---|
| z_dim | 128 | 潜向量维度，太小生成多样性不足，太大训练慢 |
| ngf | 64 | 生成器基础通道数，显存不够降到 32 |
| n_critic | 5 | 判别器每轮更新次数 |
| lambda_gp | 10 | 梯度惩罚系数 |
| lr | 1e-4 | Adam 学习率 |
| betas | (0.5, 0.9) | 动量项，第二个别用 0.999 |
| batch_size | 32 | 8GB 显存下的安全值 |
| epochs | 200+ | WGAN-GP 收敛慢，别指望几十轮出结果 |
显存占用方面，256×256 分辨率、batch_size=32、ngf=64 的配置下，训练时峰值显存大约 6-7GB。如果 OOM，优先降 batch_size 到 16，其次降 ngf 到 32。不要一上来就降分辨率——256 是这个方案的核心目标，降到 128 就偏离标题了。
3.3 训练过程中该盯哪些指标
WGAN-GP 的损失曲线和原始 GAN 不一样，判别器损失不是越小越好。理想情况下，d_loss会在 0 附近波动，g_loss缓慢下降。如果d_loss持续为负且绝对值越来越大，说明判别器太强了，生成器梯度信号在变弱。这时候可以适当降低n_critic或者提高lambda_gp。
除了损失值，更直观的指标是每隔几个 epoch 保存一批生成样本，肉眼观察。我一般每 10 个 epoch 存一次fake_imgs的网格图，用torchvision.utils.save_image拼成 8×8 的网格。如果连续几个 epoch 生成的图都差不多，说明模式崩塌了，需要检查学习率是不是太高，或者判别器是不是过拟合了。
还有一个容易被忽略的指标：梯度惩罚项的实际值。如果gp远大于 1，说明判别器的梯度范数偏离 1 太多，惩罚项在主导损失，这时候训练可能不稳定。正常训练时gp应该在 0.1 到 1 之间波动。
4. 避坑与排查：WGAN-GP 训练中最容易翻车的五个地方
4.1 判别器损失变成 NaN
现象：训练几十步后d_loss突然变成 NaN，后续所有参数都变成 NaN。
原因：梯度惩罚计算时用了torch.autograd.grad，如果插值样本的梯度爆炸，惩罚项会变成无穷大。常见触发条件是学习率太高，或者alpha采样时出现了极端值。
解决：把学习率降到 1e-4 以下，检查gradient_penalty里有没有加create_graph=True。另外可以在惩罚项外面加一个torch.clamp，把gradient_norm限制在 [0, 10] 范围内，防止极端值传播。
4.2 生成器输出全是同一张脸
现象：训练到后期，生成的 64 张图看起来几乎一模一样，只是轻微色差。
原因：模式崩塌。WGAN-GP 虽然比原始 GAN 稳定，但在数据量不足或者判别器过强时仍然会出现。另一个可能原因是z_dim太小，潜空间表达能力不够。
解决：先把z_dim从 128 提到 256 试试。如果没用，检查判别器是不是更新太频繁了，把n_critic从 5 降到 3。还可以在生成器损失里加一个小的多样性正则项，但这不是标准做法，优先调结构参数。
4.3 训练 loss 正常但生成图全是噪点
现象：d_loss和g_loss都在正常范围波动，但生成的图就是雪花噪点，没有任何结构。
原因：最常见的是数据预处理和生成器输出范围没对齐。比如数据归一化用了 ImageNet 的均值和方差，但生成器最后一层是Tanh输出 [-1,1]，两者不匹配。另一个可能是判别器太弱，根本没学到东西。
解决：检查Normalize的参数是不是[0.5]*3, [0.5]*3。然后单独测试判别器：拿真实图片和随机噪声分别输入判别器，看输出的分数有没有明显差异。如果没有差异，说明判别器没训练起来，检查判别器的初始化或者学习率。
4.4 显存溢出但 batch_size 已经很小
现象：batch_size降到 8 还是 OOM，但模型参数量看起来不大。
原因：梯度惩罚计算时保留了计算图，retain_graph=True会导致中间激活值不被释放。另外如果n_critic设得很大，每次循环都在累积计算图。
解决：确保每次判别器更新后调用optimizer_critic.zero_grad()，并且在生成器更新前把判别器的计算图释放掉。可以在判别器循环里用with torch.no_grad()包住不需要梯度的部分，但注意梯度惩罚那部分不能包。如果还不行，把ngf降到 32，或者用混合精度训练。
4.5 训练了几百轮生成质量还是模糊
现象：训练了 300 个 epoch，生成的图能看出是头像，但边缘模糊、细节缺失。
原因：WGAN-GP 在 256 分辨率下收敛确实慢，几百轮不够是正常的。另一个原因是判别器容量不够，无法捕捉高频细节。
解决：先确认训练轮数——256×256 的动漫头像，我一般跑 500 到 1000 个 epoch 才看到比较清晰的结果。如果轮数够了还是模糊，把判别器的ndf从 64 提到 128，增加判别器的表达能力。还可以在生成器里加残差连接，帮助梯度传播。
5. 进阶技巧：用谱归一化加速收敛并稳定 256×256 训练
如果你已经跑通了基础版本，但觉得收敛太慢或者训练后期还是偶尔不稳定，可以试试把判别器里的 LayerNorm 换成谱归一化（Spectral Normalization）。谱归一化直接约束判别器每层的 Lipschitz 常数，和 WGAN-GP 的梯度惩罚是互补的——一个在损失层面约束，一个在权重层面约束。
from torch.nn.utils import spectral_norm class CriticSN(nn.Module): def __init__(self, ndf=64): super().__init__() self.net = nn.Sequential( spectral_norm(nn.Conv2d(3, ndf, 4, 2, 1)), nn.LeakyReLU(0.2, inplace=True), spectral_norm(nn.Conv2d(ndf, ndf*2, 4, 2, 1)), nn.LeakyReLU(0.2, inplace=True), spectral_norm(nn.Conv2d(ndf*2, ndf*4, 4, 2, 1)), nn.LeakyReLU(0.2, inplace=True), spectral_norm(nn.Conv2d(ndf*4, ndf*8, 4, 2, 1)), nn.LeakyReLU(0.2, inplace=True), spectral_norm(nn.Conv2d(ndf*8, ndf*8, 4, 2, 1)), nn.LeakyReLU(0.2, inplace=True), spectral_norm(nn.Conv2d(ndf*8, 1, 4, 1, 0)), ) def forward(self, x): return self.net(x).view(-1)spectral_norm直接包在Conv2d外面，每次前向传播时会自动对权重做谱归一化。注意用了谱归一化之后，梯度惩罚的lambda_gp可以适当降低，比如从 10 降到 5，因为权重层面的约束已经分担了一部分 Lipschitz 约束的压力。
验证谱归一化有没有生效，可以打印判别器权重的谱范数。正常情况下，每层权重的谱范数应该接近 1。如果远大于 1，说明谱归一化没起作用，检查是不是漏包了某一层。
另一个实用技巧是学习率预热。前 5 个 epoch 把学习率从 1e-5 线性升到 1e-4，让判别器先“热身”，避免一开始就产生过大的梯度惩罚。这个技巧在 256 分辨率下效果比较明显，能减少早期 NaN 的概率。
我自己的习惯是：每次开新实验，先用 500 张图跑 20 个 epoch 做 sanity check，确认 loss 曲线正常、生成图有基本结构，再换全量数据跑长训练。这样翻车成本低，不用等几百轮才发现参数配错了。希望帮到你。
本文还有配套的精品资源，点击获取
