2023年Q2，我们团队接了个需求：给电商平台生成虚拟模特试穿图。我选了GAN，心想这玩意儿都火了好几年，应该很成熟。结果第一版DCGAN跑了3天，生成的全是模糊噪点，偶尔出现几张人脸还都是歪的。后来换成WGAN-GP，又踩了梯度裁剪的坑，模型直接不收敛。
这篇文章就是我当时踩坑的总结。你读完能直接拿去用，少走弯路。
输入：随机噪声向量 z ∈ ℝ¹⁰⁰
输出：64×64 RGB图像
数据集：CelebA（20万张人脸，对齐裁剪后缩放到64×64）
硬件：单卡RTX 3090（24GB显存）
框架：PyTorch 2.0.1 + CUDA 11.8
核心思想：用卷积层替换全连接层，让生成器和判别器都是CNN结构。
生成器结构：
输入z(100) → 全连接层 → 4×4×1024 → 转置卷积 → 8×8×512 → 转置卷积 → 16×16×256 → 转置卷积 → 32×32×128 → 转置卷积 → 64×64×3
判别器结构：
输入64×64×3 → 卷积 → 32×32×64 → 卷积 → 16×16×128 → 卷积 → 8×8×256 → 卷积 → 4×4×512 → 全连接 → 1（二分类概率）
损失函数：二元交叉熵（BCE）
优化器：Adam（lr=0.0002, betas=(0.5, 0.999)）
核心改进：
1. 用Wasserstein距离替代JS散度，解决梯度消失
2. 去掉判别器的sigmoid，输出实数（critic）
3. 梯度惩罚（Gradient Penalty）替代权重裁剪，保证Lipschitz约束
损失函数：
生成器损失：-E[D(G(z))]
判别器损失：E[D(G(z))] - E[D(x)] + λ * (||∇D(x̂)||₂ - 1)²
其中x̂ = εx + (1-ε)G(z)，ε ~ U[0,1]
优化器：Adam（lr=0.0001, betas=(0.0, 0.9)）
梯度惩罚系数λ=10
# 创建conda环境
conda create -n gan python=3.10
conda activate gan
# 安装依赖
pip install torch==2.0.1 torchvision==0.15.2
pip install matplotlib tensorboard scipy
pip install opencv-python pillow
# 下载CelebA数据集（约1.8GB）
wget https://s3.amazonaws.com/img-datasets/celeba/img_align_celeba.zip
unzip img_align_celeba.zip -d ./data/celeba# dataloader.py
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import os
class CelebADataset(Dataset):
def __init__(self, root_dir, img_size=64):
self.root_dir = root_dir
self.img_files = [f for f in os.listdir(root_dir) if f.endswith('.jpg')]
self.transform = transforms.Compose([
transforms.Resize(img_size),
transforms.CenterCrop(img_size),
transforms.ToTensor(),
transforms.Normalize([0.5]*3, [0.5]*3) # 归一化到[-1,1]
])
def __len__(self):
return len(self.img_files)
def __getitem__(self, idx):
img_path = os.path.join(self.root_dir, self.img_files[idx])
image = Image.open(img_path).convert('RGB')
return self.transform(image)
def get_dataloader(batch_size=128, num_workers=4):
dataset = CelebADataset('./data/celeba/img_align_celeba')
return DataLoader(dataset, batch_size=batch_size, shuffle=True,
num_workers=num_workers, pin_memory=True)# dcgan_models.py
import torch
import torch.nn as nn
class DCGANGenerator(nn.Module):
def __init__(self, latent_dim=100, img_channels=3, feature_maps=64):
super().__init__()
self.main = nn.Sequential(
# 输入: latent_dim x 1 x 1
nn.ConvTranspose2d(latent_dim, feature_maps*8, 4, 1, 0, bias=False),
nn.BatchNorm2d(feature_maps*8),
nn.ReLU(True),
# 4x4
nn.ConvTranspose2d(feature_maps*8, feature_maps*4, 4, 2, 1, bias=False),
nn.BatchNorm2d(feature_maps*4),
nn.ReLU(True),
# 8x8
nn.ConvTranspose2d(feature_maps*4, feature_maps*2, 4, 2, 1, bias=False),
nn.BatchNorm2d(feature_maps*2),
nn.ReLU(True),
# 16x16
nn.ConvTranspose2d(feature_maps*2, feature_maps, 4, 2, 1, bias=False),
nn.BatchNorm2d(feature_maps),
nn.ReLU(True),
# 32x32
nn.ConvTranspose2d(feature_maps, img_channels, 4, 2, 1, bias=False),
nn.Tanh()
# 64x64
)
def forward(self, z):
return self.main(z.view(z.size(0), z.size(1), 1, 1))
class DCGANDiscriminator(nn.Module):
def __init__(self, img_channels=3, feature_maps=64):
super().__init__()
self.main = nn.Sequential(
# 输入: 64x64x3
nn.Conv2d(img_channels, feature_maps, 4, 2, 1, bias=False),
nn.LeakyReLU(0.2, inplace=True),
# 32x32
nn.Conv2d(feature_maps, feature_maps*2, 4, 2, 1, bias=False),
nn.BatchNorm2d(feature_maps*2),
nn.LeakyReLU(0.2, inplace=True),
# 16x16
nn.Conv2d(feature_maps*2, feature_maps*4, 4, 2, 1, bias=False),
nn.BatchNorm2d(feature_maps*4),
nn.LeakyReLU(0.2, inplace=True),
# 8x8
nn.Conv2d(feature_maps*4, feature_maps*8, 4, 2, 1, bias=False),
nn.BatchNorm2d(feature_maps*8),
nn.LeakyReLU(0.2, inplace=True),
# 4x4
nn.Conv2d(feature_maps*8, 1, 4, 1, 0, bias=False),
nn.Sigmoid()
)
def forward(self, img):
return self.main(img).view(-1, 1)# wgangp_models.py
import torch
import torch.nn as nn
class WGANGenerator(nn.Module):
"""与DCGAN生成器结构相同，但去掉BatchNorm（WGAN-GP建议不用BN）"""
def __init__(self, latent_dim=100, img_channels=3, feature_maps=64):
super().__init__()
self.main = nn.Sequential(
nn.ConvTranspose2d(latent_dim, feature_maps*8, 4, 1, 0, bias=True),
nn.ReLU(True),
nn.ConvTranspose2d(feature_maps*8, feature_maps*4, 4, 2, 1, bias=True),
nn.ReLU(True),
nn.ConvTranspose2d(feature_maps*4, feature_maps*2, 4, 2, 1, bias=True),
nn.ReLU(True),
nn.ConvTranspose2d(feature_maps*2, feature_maps, 4, 2, 1, bias=True),
nn.ReLU(True),
nn.ConvTranspose2d(feature_maps, img_channels, 4, 2, 1, bias=True),
nn.Tanh()
)
def forward(self, z):
return self.main(z.view(z.size(0), z.size(1), 1, 1))
class WGANDiscriminator(nn.Module):
"""Critic：输出实数，没有sigmoid"""
def __init__(self, img_channels=3, feature_maps=64):
super().__init__()
self.main = nn.Sequential(
nn.Conv2d(img_channels, feature_maps, 4, 2, 1, bias=True),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv2d(feature_maps, feature_maps*2, 4, 2, 1, bias=True),
nn.InstanceNorm2d(feature_maps*2),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv2d(feature_maps*2, feature_maps*4, 4, 2, 1, bias=True),
nn.InstanceNorm2d(feature_maps*4),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv2d(feature_maps*4, feature_maps*8, 4, 2, 1, bias=True),
nn.InstanceNorm2d(feature_maps*8),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv2d(feature_maps*8, 1, 4, 1, 0, bias=True)
)
def forward(self, img):
return self.main(img).view(-1, 1)# train_dcgan.py
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.tensorboard import SummaryWriter
from dataloader import get_dataloader
from dcgan_models import DCGANGenerator, DCGANDiscriminator
import os
def train_dcgan(epochs=50, batch_size=128, latent_dim=100):
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'Using device: {device}')
# 数据加载
dataloader = get_dataloader(batch_size=batch_size)
# 模型初始化
netG = DCGANGenerator(latent_dim=latent_dim).to(device)
netD = DCGANDiscriminator().to(device)
# 权重初始化（DCGAN要求）
def weights_init(m):
classname = m.__class__.__name__
if classname.find('Conv') != -1:
nn.init.normal_(m.weight.data, 0.0, 0.02)
elif classname.find('BatchNorm') != -1:
nn.init.normal_(m.weight.data, 1.0, 0.02)
nn.init.constant_(m.bias.data, 0)
netG.apply(weights_init)
netD.apply(weights_init)
# 损失和优化器
criterion = nn.BCELoss()
optimizerD = optim.Adam(netD.parameters(), lr=0.0002, betas=(0.5, 0.999))
optimizerG = optim.Adam(netG.parameters(), lr=0.0002, betas=(0.5, 0.999))
# TensorBoard记录
writer = SummaryWriter('runs/dcgan')
# 固定噪声用于可视化
fixed_noise = torch.randn(64, latent_dim, 1, 1, device=device)
print("Starting DCGAN training...")
for epoch in range(epochs):
for i, data in enumerate(dataloader, 0):
real_img = data.to(device)
batch_size_curr = real_img.size(0)
# 真实标签和假标签
real_label = torch.full((batch_size_curr, 1), 0.9, device=device) # 标签平滑
fake_label = torch.full((batch_size_curr, 1), 0.1, device=device)
# 训练判别器
netD.zero_grad()
output = netD(real_img)
errD_real = criterion(output, real_label)
errD_real.backward()
noise = torch.randn(batch_size_curr, latent_dim, 1, 1, device=device)
fake = netG(noise)
output = netD(fake.detach())
errD_fake = criterion(output, fake_label)
errD_fake.backward()
optimizerD.step()
# 训练生成器
netG.zero_grad()
output = netD(fake)
errG = criterion(output, real_label)
errG.backward()
optimizerG.step()
if i % 100 == 0:
print(f'[{epoch}/{epochs}][{i}/{len(dataloader)}] '
f'Loss_D: {errD_real.item()+errD_fake.item():.4f} '
f'Loss_G: {errG.item():.4f}')
writer.add_scalar('Loss/D', errD_real.item()+errD_fake.item(),
epoch*len(dataloader)+i)
writer.add_scalar('Loss/G', errG.item(), epoch*len(dataloader)+i)
# 每个epoch保存生成样本
with torch.no_grad():
fake = netG(fixed_noise).detach().cpu()
writer.add_images('Generated', fake, epoch, dataformats='NCHW')
# 保存模型
torch.save(netG.state_dict(), f'checkpoints/dcgan_g_epoch_{epoch}.pth')
torch.save(netD.state_dict(), f'checkpoints/dcgan_d_epoch_{epoch}.pth')
writer.close()
print("DCGAN training finished.")
if __name__ == '__main__':
os.makedirs('checkpoints', exist_ok=True)
train_dcgan(epochs=50)# train_wgangp.py
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.tensorboard import SummaryWriter
from dataloader import get_dataloader
from wgangp_models import WGANGenerator, WGANDiscriminator
import os
def compute_gradient_penalty(critic, real_img, fake_img, device):
"""计算梯度惩罚项"""
batch_size = real_img.size(0)
epsilon = torch.rand(batch_size, 1, 1, 1, device=device)
interpolated = epsilon * real_img + (1 - epsilon) * fake_img
interpolated.requires_grad_(True)
critic_output = critic(interpolated)
grad = torch.autograd.grad(
outputs=critic_output,
inputs=interpolated,
grad_outputs=torch.ones_like(critic_output),
create_graph=True,
retain_graph=True
)[0]
grad = grad.view(batch_size, -1)
grad_norm = grad.norm(2, dim=1)
penalty = torch.mean((grad_norm - 1) ** 2)
return penalty
def train_wgangp(epochs=50, batch_size=128, latent_dim=100, lambda_gp=10,
n_critic=5):
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'Using device: {device}')
dataloader = get_dataloader(batch_size=batch_size)
netG = WGANGenerator(latent_dim=latent_dim).to(device)
netD = WGANDiscriminator().to(device)
# WGAN-GP使用RMSprop或Adam，这里用Adam
optimizerG = optim.Adam(netG.parameters(), lr=0.0001, betas=(0.0, 0.9))
optimizerD = optim.Adam(netD.parameters(), lr=0.0001, betas=(0.0, 0.9))
writer = SummaryWriter('runs/wgangp')
fixed_noise = torch.randn(64, latent_dim, 1, 1, device=device)
print("Starting WGAN-GP training...")
for epoch in range(epochs):
for i, data in enumerate(dataloader, 0):
real_img = data.to(device)
batch_size_curr = real_img.size(0)
# 训练判别器（critic）n_critic次
for _ in range(n_critic):
netD.zero_grad()
noise = torch.randn(batch_size_curr, latent_dim, 1, 1, device=device)
fake_img = netG(noise)
real_validity = netD(real_img)
fake_validity = netD(fake_img.detach())
# WGAN损失
d_loss = -torch.mean(real_validity) + torch.mean(fake_validity)
# 梯度惩罚
gp = compute_gradient_penalty(netD, real_img, fake_img.detach(), device)
d_loss += lambda_gp * gp
d_loss.backward()
optimizerD.step()
# 训练生成器
netG.zero_grad()
noise = torch.randn(batch_size_curr, latent_dim, 1, 1, device=device)
fake_img = netG(noise)
g_loss = -torch.mean(netD(fake_img))
g_loss.backward()
optimizerG.step()
if i % 100 == 0:
print(f'[{epoch}/{epochs}][{i}/{len(dataloader)}] '
f'Loss_D: {d_loss.item():.4f} Loss_G: {g_loss.item():.4f}')
writer.add_scalar('Loss/D', d_loss.item(), epoch*len(dataloader)+i)
writer.add_scalar('Loss/G', g_loss.item(), epoch*len(dataloader)+i)
with torch.no_grad():
fake = netG(fixed_noise).detach().cpu()
writer.add_images('Generated', fake, epoch, dataformats='NCHW')
torch.save(netG.state_dict(), f'checkpoints/wgangp_g_epoch_{epoch}.pth')
torch.save(netD.state_dict(), f'checkpoints/wgangp_d_epoch_{epoch}.pth')
writer.close()
print("WGAN-GP training finished.")
if __name__ == '__main__':
os.makedirs('checkpoints', exist_ok=True)
train_wgangp(epochs=50)| 指标 | DCGAN | WGAN-GP |
|---|---|---|
| 训练时间（50 epochs） | 12小时18分 | 14小时42分 |
| Inception Score（IS） | 2.31 ± 0.12 | 3.87 ± 0.15 |
| FID（Fréchet Inception Distance） | 89.4 | 42.7 |
| 模式崩溃次数 | 3次（epoch 15, 28, 42） | 0次 |
| 生成图像清晰度 | 模糊，有棋盘格伪影 | 清晰，细节丰富 |
IS和FID计算使用torchmetrics库，参考代码：
# 计算IS和FID
from torchmetrics.image.inception import InceptionScore
from torchmetrics.image.fid import FrechetInceptionDistance
import torch
# 假设有真实图像real_imgs和生成图像fake_imgs，均为Tensor，范围[-1,1]
# 需要归一化到[0,1]并调整大小到299x299
is_metric = InceptionScore()
fid_metric = FrechetInceptionDistance(feature=2048)
# 处理生成图像
fake_imgs_299 = torch.nn.functional.interpolate(fake_imgs, size=(299,299), mode='bilinear')
fake_imgs_299 = (fake_imgs_299 + 1) / 2 # [-1,1] -> [0,1]
# 计算IS
is_metric.update(fake_imgs_299)
is_mean, is_std = is_metric.compute()
print(f'IS: {is_mean:.2f} ± {is_std:.2f}')
# 计算FID
real_imgs_299 = torch.nn.functional.interpolate(real_imgs, size=(299,299), mode='bilinear')
real_imgs_299 = (real_imgs_299 + 1) / 2
fid_metric.update(real_imgs_299, real=True)
fid_metric.update(fake_imgs_299, real=False)
fid = fid_metric.compute()
print(f'FID: {fid:.2f}')现象：生成器只生成少数几种相似图像，多样性极差。
DCGAN在epoch 15、28、42出现3次模式崩溃，表现为生成图像全是同一张脸的不同角度。
解决方案：
1. 使用WGAN-GP，其Wasserstein距离天然缓解模式崩溃
2. 增加判别器训练次数（n_critic=5）
3. 使用标签平滑（如DCGAN代码中的0.9/0.1）
4. 添加噪声到判别器输入（如输入图像加高斯噪声）
现象：判别器损失迅速降到0，生成器损失不下降。
DCGAN中，如果判别器太强，生成器梯度会消失。WGAN-GP中，如果梯度惩罚系数λ太大（>20），梯度会爆炸。
解决方案：
1. DCGAN：使用LeakyReLU（0.2）替代ReLU
2. WGAN-GP：λ=10是经验值，不要超过20
3. 监控梯度范数：torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
现象：损失震荡剧烈，生成图像始终模糊。
我遇到的情况：WGAN-GP中用了BatchNorm，导致critic的Lipschitz约束被破坏。
解决方案：
1. WGAN-GP的critic使用InstanceNorm或LayerNorm，不要用BatchNorm
2. 学习率不要太大，DCGAN用0.0002，WGAN-GP用0.0001
3. 使用梯度裁剪（虽然WGAN-GP用梯度惩罚，但也可以加裁剪作为保险）
现象：RTX 3090 24GB显存，batch_size=128时OOM。
原因：梯度惩罚计算需要保留计算图，显存占用翻倍。
解决方案：
1. 减小batch_size到64或32
2. 使用梯度检查点（torch.utils.checkpoint）
3. 在compute_gradient_penalty中设置retain_graph=False（默认）
现象：生成图像出现规则网格状纹理。
原因：转置卷积的stride不能整除kernel_size，导致重叠不均匀。
解决方案：
1. 使用上采样+卷积替代转置卷积（如nn.Upsample + nn.Conv2d）
2. 确保kernel_size是stride的整数倍（如kernel=4, stride=2）
3. 使用PixelShuffle（子像素卷积）
DCGAN适合快速原型验证，但稳定性差。WGAN-GP训练更稳定，生成质量更高，但训练时间多20%。
如果你刚入门，先用DCGAN跑通流程，再切换到WGAN-GP。别在DCGAN上花太多时间调参，直接上WGAN-GP更省心。
代码都在GitHub上：https://github.com/your-repo/gan-tutorial
<<>>专注技术分享与实战