拓冰网站定制 
简介这份源码资源面向深度学习入门者与图像生成爱好者提供一套基于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( outputsscore, inputsinterpolated, grad_outputstorch.ones_like(score), create_graphTrue, retain_graphTrue, only_inputsTrue )[0] # 梯度范数偏离 1 的惩罚 gradient_norm gradient.view(batch_size, -1).norm(2, dim1) penalty ((gradient_norm - 1) ** 2).mean() return penalty这段代码的关键参数是alpha的采样方式——从均匀分布 U[0,1] 中采样在真实样本和生成样本之间做线性插值。gradient_norm计算的是判别器输出对插值输入的梯度 L2 范数惩罚项就是让这个范数尽量接近 1。create_graphTrue必须开因为惩罚项本身也要参与反向传播。retain_graphTrue是为了后续还能继续用这个计算图。判别器损失由两部分组成真实样本分数均值减去生成样本分数均值再加上梯度惩罚乘以惩罚系数 λ。λ 一般取 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_dim128, ngf64): super().__init__() self.net nn.Sequential( # 输入 z: (z_dim, 1, 1) - (ngf*8, 4, 4) nn.ConvTranspose2d(z_dim, ngf*8, 4, 1, 0, biasFalse), nn.BatchNorm2d(ngf*8), nn.ReLU(True), # (ngf*8, 4, 4) - (ngf*4, 8, 8) nn.ConvTranspose2d(ngf*8, ngf*4, 4, 2, 1, biasFalse), nn.BatchNorm2d(ngf*4), nn.ReLU(True), # (ngf*4, 8, 8) - (ngf*2, 16, 16) nn.ConvTranspose2d(ngf*4, ngf*2, 4, 2, 1, biasFalse), nn.BatchNorm2d(ngf*2), nn.ReLU(True), # (ngf*2, 16, 16) - (ngf, 32, 32) nn.ConvTranspose2d(ngf*2, ngf, 4, 2, 1, biasFalse), nn.BatchNorm2d(ngf), nn.ReLU(True), # (ngf, 32, 32) - (ngf//2, 64, 64) nn.ConvTranspose2d(ngf, ngf//2, 4, 2, 1, biasFalse), nn.BatchNorm2d(ngf//2), nn.ReLU(True), # (ngf//2, 64, 64) - (ngf//4, 128, 128) nn.ConvTranspose2d(ngf//2, ngf//4, 4, 2, 1, biasFalse), nn.BatchNorm2d(ngf//4), nn.ReLU(True), # (ngf//4, 128, 128) - (3, 256, 256) nn.ConvTranspose2d(ngf//4, 3, 4, 2, 1, biasFalse), nn.Tanh() ) def forward(self, z): return self.net(z)这里z_dim128是潜向量维度ngf64是基础通道数。从 4×4 开始每次转置卷积的kernel_size4, stride2, padding1输出尺寸翻倍。最后一层输出 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_critic5是原论文的推荐值lambda_gp10。优化器用 Adam学习率 1e-4betas 设成 (0.5, 0.9)——注意第二个 beta 不要用默认的 0.999WGAN-GP 对动量项比较敏感0.9 更稳。生成器那边fake_imgs不需要detach()因为要回传梯度到生成器。3. 从零跑通训练数据准备、参数配置与监控指标3.1 动漫头像数据集的整理与预处理数据来源通常是爬取或者公开数据集但不管哪来的统一处理成 256×256 的 RGB 图片。我一般用torchvision的ImageFolder配合自定义 transformfrom 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, transformtransform) dataloader torch.utils.data.DataLoader( dataset, batch_size32, shuffleTrue, num_workers4, drop_lastTrue )Normalize的均值和标准差都设 0.5把 [0,1] 的像素值映射到 [-1,1]和生成器最后一层Tanh的输出范围对齐。drop_lastTrue是为了避免最后一个 batch 尺寸不固定导致梯度惩罚计算出问题。batch_size32在 8GB 显存上跑 256×256 基本够用显存紧张就降到 16。数据量方面至少准备 5000 张以上否则判别器很容易过拟合到训练集生成器学不到有意义的分布。如果数据不够可以做水平翻转增强但不要做旋转或裁剪——动漫头像的构图通常是对称的旋转会引入不自然的样本。3.2 关键超参数表与显存占用估算参数推荐值说明z_dim128潜向量维度太小生成多样性不足太大训练慢ngf64生成器基础通道数显存不够降到 32n_critic5判别器每轮更新次数lambda_gp10梯度惩罚系数lr1e-4Adam 学习率betas(0.5, 0.9)动量项第二个别用 0.999batch_size328GB 显存下的安全值epochs200WGAN-GP 收敛慢别指望几十轮出结果显存占用方面256×256 分辨率、batch_size32、ngf64 的配置下训练时峰值显存大约 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_graphTrue。另外可以在惩罚项外面加一个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_graphTrue会导致中间激活值不被释放。另外如果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, ndf64): super().__init__() self.net nn.Sequential( spectral_norm(nn.Conv2d(3, ndf, 4, 2, 1)), nn.LeakyReLU(0.2, inplaceTrue), spectral_norm(nn.Conv2d(ndf, ndf*2, 4, 2, 1)), nn.LeakyReLU(0.2, inplaceTrue), spectral_norm(nn.Conv2d(ndf*2, ndf*4, 4, 2, 1)), nn.LeakyReLU(0.2, inplaceTrue), spectral_norm(nn.Conv2d(ndf*4, ndf*8, 4, 2, 1)), nn.LeakyReLU(0.2, inplaceTrue), spectral_norm(nn.Conv2d(ndf*8, ndf*8, 4, 2, 1)), nn.LeakyReLU(0.2, inplaceTrue), 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 曲线正常、生成图有基本结构再换全量数据跑长训练。这样翻车成本低不用等几百轮才发现参数配错了。希望帮到你。本文还有配套的精品资源点击获取 WGAN-GP 实战：从零训练 256×256 动漫头像生成模型
简介这份源码资源面向深度学习入门者与图像生成爱好者提供一套基于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( outputsscore, inputsinterpolated, grad_outputstorch.ones_like(score), create_graphTrue, retain_graphTrue, only_inputsTrue )[0] # 梯度范数偏离 1 的惩罚 gradient_norm gradient.view(batch_size, -1).norm(2, dim1) penalty ((gradient_norm - 1) ** 2).mean() return penalty这段代码的关键参数是alpha的采样方式——从均匀分布 U[0,1] 中采样在真实样本和生成样本之间做线性插值。gradient_norm计算的是判别器输出对插值输入的梯度 L2 范数惩罚项就是让这个范数尽量接近 1。create_graphTrue必须开因为惩罚项本身也要参与反向传播。retain_graphTrue是为了后续还能继续用这个计算图。判别器损失由两部分组成真实样本分数均值减去生成样本分数均值再加上梯度惩罚乘以惩罚系数 λ。λ 一般取 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_dim128, ngf64): super().__init__() self.net nn.Sequential( # 输入 z: (z_dim, 1, 1) - (ngf*8, 4, 4) nn.ConvTranspose2d(z_dim, ngf*8, 4, 1, 0, biasFalse), nn.BatchNorm2d(ngf*8), nn.ReLU(True), # (ngf*8, 4, 4) - (ngf*4, 8, 8) nn.ConvTranspose2d(ngf*8, ngf*4, 4, 2, 1, biasFalse), nn.BatchNorm2d(ngf*4), nn.ReLU(True), # (ngf*4, 8, 8) - (ngf*2, 16, 16) nn.ConvTranspose2d(ngf*4, ngf*2, 4, 2, 1, biasFalse), nn.BatchNorm2d(ngf*2), nn.ReLU(True), # (ngf*2, 16, 16) - (ngf, 32, 32) nn.ConvTranspose2d(ngf*2, ngf, 4, 2, 1, biasFalse), nn.BatchNorm2d(ngf), nn.ReLU(True), # (ngf, 32, 32) - (ngf//2, 64, 64) nn.ConvTranspose2d(ngf, ngf//2, 4, 2, 1, biasFalse), nn.BatchNorm2d(ngf//2), nn.ReLU(True), # (ngf//2, 64, 64) - (ngf//4, 128, 128) nn.ConvTranspose2d(ngf//2, ngf//4, 4, 2, 1, biasFalse), nn.BatchNorm2d(ngf//4), nn.ReLU(True), # (ngf//4, 128, 128) - (3, 256, 256) nn.ConvTranspose2d(ngf//4, 3, 4, 2, 1, biasFalse), nn.Tanh() ) def forward(self, z): return self.net(z)这里z_dim128是潜向量维度ngf64是基础通道数。从 4×4 开始每次转置卷积的kernel_size4, stride2, padding1输出尺寸翻倍。最后一层输出 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_critic5是原论文的推荐值lambda_gp10。优化器用 Adam学习率 1e-4betas 设成 (0.5, 0.9)——注意第二个 beta 不要用默认的 0.999WGAN-GP 对动量项比较敏感0.9 更稳。生成器那边fake_imgs不需要detach()因为要回传梯度到生成器。3. 从零跑通训练数据准备、参数配置与监控指标3.1 动漫头像数据集的整理与预处理数据来源通常是爬取或者公开数据集但不管哪来的统一处理成 256×256 的 RGB 图片。我一般用torchvision的ImageFolder配合自定义 transformfrom 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, transformtransform) dataloader torch.utils.data.DataLoader( dataset, batch_size32, shuffleTrue, num_workers4, drop_lastTrue )Normalize的均值和标准差都设 0.5把 [0,1] 的像素值映射到 [-1,1]和生成器最后一层Tanh的输出范围对齐。drop_lastTrue是为了避免最后一个 batch 尺寸不固定导致梯度惩罚计算出问题。batch_size32在 8GB 显存上跑 256×256 基本够用显存紧张就降到 16。数据量方面至少准备 5000 张以上否则判别器很容易过拟合到训练集生成器学不到有意义的分布。如果数据不够可以做水平翻转增强但不要做旋转或裁剪——动漫头像的构图通常是对称的旋转会引入不自然的样本。3.2 关键超参数表与显存占用估算参数推荐值说明z_dim128潜向量维度太小生成多样性不足太大训练慢ngf64生成器基础通道数显存不够降到 32n_critic5判别器每轮更新次数lambda_gp10梯度惩罚系数lr1e-4Adam 学习率betas(0.5, 0.9)动量项第二个别用 0.999batch_size328GB 显存下的安全值epochs200WGAN-GP 收敛慢别指望几十轮出结果显存占用方面256×256 分辨率、batch_size32、ngf64 的配置下训练时峰值显存大约 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_graphTrue。另外可以在惩罚项外面加一个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_graphTrue会导致中间激活值不被释放。另外如果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, ndf64): super().__init__() self.net nn.Sequential( spectral_norm(nn.Conv2d(3, ndf, 4, 2, 1)), nn.LeakyReLU(0.2, inplaceTrue), spectral_norm(nn.Conv2d(ndf, ndf*2, 4, 2, 1)), nn.LeakyReLU(0.2, inplaceTrue), spectral_norm(nn.Conv2d(ndf*2, ndf*4, 4, 2, 1)), nn.LeakyReLU(0.2, inplaceTrue), spectral_norm(nn.Conv2d(ndf*4, ndf*8, 4, 2, 1)), nn.LeakyReLU(0.2, inplaceTrue), spectral_norm(nn.Conv2d(ndf*8, ndf*8, 4, 2, 1)), nn.LeakyReLU(0.2, inplaceTrue), 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 曲线正常、生成图有基本结构再换全量数据跑长训练。这样翻车成本低不用等几百轮才发现参数配错了。希望帮到你。本文还有配套的精品资源点击获取