用PyTorch和MNIST数据集，手把手教你复现CGAN生成指定数字（附完整代码）

从零实现CGAN：用PyTorch控制MNIST数字生成的实战指南

当你第一次看到计算机生成的手写数字时，是否好奇它是如何实现的？传统GAN虽然能生成逼真图像，却无法控制生成内容的具体特征。本文将带你亲手构建一个条件生成对抗网络（CGAN），实现按需生成指定数字的功能。

1. 环境准备与数据加载

在开始构建CGAN之前，我们需要配置合适的开发环境。推荐使用Python 3.8+和PyTorch 1.10+版本，这些组合在稳定性和功能支持上表现最佳。

安装核心依赖包：

pip install torch torchvision matplotlib numpy

MNIST数据集加载与预处理是项目的第一步。PyTorch的torchvision模块已经内置了MNIST的便捷接口，但我们仍需进行适当的标准化处理：

transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=(0.5,), std=(0.5,)) # 将像素值归一化到[-1,1]范围 ]) train_dataset = datasets.MNIST( root='./data', train=True, download=True, transform=transform ) train_loader = torch.utils.data.DataLoader( dataset=train_dataset, batch_size=128, shuffle=True, num_workers=4 )

关键细节说明：

• 批大小(batch_size)设置为128，这是兼顾内存使用和训练稳定性的折中选择
• 数据归一化到[-1,1]范围是为了匹配生成器输出层的tanh激活函数
• 使用4个工作进程(num_workers)可以加速数据加载

2. CGAN模型架构设计

CGAN的核心创新在于将条件信息（这里指数字标签）同时注入生成器和判别器。与基础GAN相比，这需要对网络结构进行针对性调整。

2.1 生成器网络实现

生成器接收两个输入：随机噪声z和数字标签y。我们需要将这两个输入在中间层进行融合：

class Generator(nn.Module): def __init__(self, latent_dim=100): super(Generator, self).__init__() # 噪声处理分支 self.noise_fc = nn.Sequential( nn.Linear(latent_dim, 256), nn.BatchNorm1d(256), nn.ReLU() ) # 标签处理分支 self.label_fc = nn.Sequential( nn.Linear(10, 256), # 10类MNIST数字 nn.BatchNorm1d(256), nn.ReLU() ) # 联合处理层 self.main = nn.Sequential( nn.Linear(512, 512), nn.BatchNorm1d(512), nn.ReLU(), nn.Linear(512, 1024), nn.BatchNorm1d(1024), nn.ReLU(), nn.Linear(1024, 784), nn.Tanh() # 输出范围[-1,1]匹配归一化数据 ) def forward(self, noise, labels): noise_out = self.noise_fc(noise) label_out = self.label_fc(labels) combined = torch.cat([noise_out, label_out], dim=1) return self.main(combined).view(-1, 1, 28, 28)

提示：生成器最后一层使用Tanh激活函数是为了匹配输入数据的归一化范围，这是GAN训练的常见做法。

2.2 判别器网络实现

判别器同样需要处理图像和标签两个输入，但融合方式与生成器有所不同：

class Discriminator(nn.Module): def __init__(self): super(Discriminator, self).__init__() # 图像处理分支 self.image_fc = nn.Sequential( nn.Linear(784, 1024), nn.LeakyReLU(0.2) ) # 标签处理分支 self.label_fc = nn.Sequential( nn.Linear(10, 1024), nn.LeakyReLU(0.2) ) # 联合判别层 self.main = nn.Sequential( nn.Linear(2048, 512), nn.BatchNorm1d(512), nn.LeakyReLU(0.2), nn.Linear(512, 256), nn.BatchNorm1d(256), nn.LeakyReLU(0.2), nn.Linear(256, 1), nn.Sigmoid() # 输出真假概率 ) def forward(self, images, labels): images = images.view(-1, 784) img_out = self.image_fc(images) label_out = self.label_fc(labels) combined = torch.cat([img_out, label_out], dim=1) return self.main(combined)

架构设计要点对比：

3. 训练过程实现

CGAN的训练需要精心平衡生成器和判别器的优化过程。以下是训练循环的关键实现：

3.1 初始化与损失函数

device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 初始化网络 generator = Generator().to(device) discriminator = Discriminator().to(device) # 使用Adam优化器 g_optimizer = optim.Adam(generator.parameters(), lr=0.0002, betas=(0.5, 0.999)) d_optimizer = optim.Adam(discriminator.parameters(), lr=0.0002, betas=(0.5, 0.999)) # 二元交叉熵损失 criterion = nn.BCELoss() # 固定测试噪声和标签用于可视化 fixed_noise = torch.randn(10, 100, device=device).repeat(10, 1) fixed_labels = torch.zeros(100, 10, device=device) for i in range(10): fixed_labels[i*10:(i+1)*10, i] = 1

3.2 训练循环实现

num_epochs = 50 for epoch in range(num_epochs): for i, (images, real_labels) in enumerate(train_loader): batch_size = images.size(0) # 准备真实数据和标签 real_images = images.to(device) real_labels_onehot = torch.zeros(batch_size, 10, device=device) real_labels_onehot.scatter_(1, real_labels.unsqueeze(1), 1) # 训练判别器 d_optimizer.zero_grad() # 真实数据损失 real_outputs = discriminator(real_images, real_labels_onehot) real_loss = criterion(real_outputs, torch.ones(batch_size, 1, device=device)) # 生成数据损失 noise = torch.randn(batch_size, 100, device=device) fake_labels = torch.randint(0, 10, (batch_size,), device=device) fake_labels_onehot = torch.zeros(batch_size, 10, device=device) fake_labels_onehot.scatter_(1, fake_labels.unsqueeze(1), 1) fake_images = generator(noise, fake_labels_onehot) fake_outputs = discriminator(fake_images.detach(), fake_labels_onehot) fake_loss = criterion(fake_outputs, torch.zeros(batch_size, 1, device=device)) d_loss = real_loss + fake_loss d_loss.backward() d_optimizer.step() # 训练生成器 g_optimizer.zero_grad() outputs = discriminator(fake_images, fake_labels_onehot) g_loss = criterion(outputs, torch.ones(batch_size, 1, device=device)) g_loss.backward() g_optimizer.step() # 每个epoch结束后输出进度和示例图像 print(f"Epoch [{epoch+1}/{num_epochs}] | d_loss: {d_loss.item():.4f} | g_loss: {g_loss.item():.4f}") # 保存生成示例 if (epoch+1) % 5 == 0: with torch.no_grad(): fake_samples = generator(fixed_noise, fixed_labels).cpu() save_image(fake_samples, f"results/epoch_{epoch+1}.png", nrow=10, normalize=True)

注意：判别器训练时，fake_images要使用detach()切断计算图，避免梯度传播到生成器；而生成器训练时需要重新计算fake_images。

3.3 学习率调整策略

在训练中后期，适当降低学习率可以帮助模型收敛到更好的解：

if epoch == 30: for param_group in g_optimizer.param_groups: param_group['lr'] /= 10 for param_group in d_optimizer.param_groups: param_group['lr'] /= 10 print("学习率降至0.00002") if epoch == 40: for param_group in g_optimizer.param_groups: param_group['lr'] /= 10 for param_group in d_optimizer.param_groups: param_group['lr'] /= 10 print("学习率降至0.000002")

4. 结果分析与调优技巧

经过完整训练后，我们可以系统地评估模型性能并探讨常见问题的解决方案。

4.1 生成结果可视化

使用以下代码生成并显示指定数字的样本：

def generate_digit(digit, num_samples=10): generator.eval() with torch.no_grad(): noise = torch.randn(num_samples, 100, device=device) labels = torch.zeros(num_samples, 10, device=device) labels[:, digit] = 1 samples = generator(noise, labels).cpu() return samples # 生成数字5的10个样本 digit_5_samples = generate_digit(5) plt.figure(figsize=(10,2)) for i in range(10): plt.subplot(1,10,i+1) plt.imshow(digit_5_samples[i][0], cmap='gray') plt.axis('off') plt.show()

4.2 常见问题与解决方案

模式崩溃(Mode Collapse)：

• 现象：生成器只产生有限几种样本，缺乏多样性
• 解决方案：
  • 增加噪声向量的维度
  • 尝试不同的网络架构
  • 使用小批量判别(minibatch discrimination)

训练不稳定：

• 现象：损失值剧烈波动或发散
• 解决方案：
  • 调整学习率(通常先尝试降低)
  • 确保判别器不要过强(可减少其训练次数)
  • 使用梯度裁剪(gradient clipping)

生成质量差：

• 现象：生成的数字难以辨认
• 解决方案：
  • 检查数据预处理是否正确
  • 增加网络容量(更多层/更大隐藏单元)
  • 延长训练时间

4.3 进阶改进方向

对于希望进一步提升模型性能的开发者，可以考虑以下扩展：

1. 使用卷积架构：将全连接网络替换为DCGAN风格的卷积网络，通常能获得更好的图像质量
2. 添加辅助分类器：在判别器中增加数字分类任务，提供更强的监督信号
3. 实现渐进式增长：从小分辨率开始训练，逐步增加分辨率，可生成更精细的图像
4. 引入自注意力机制：在生成器和判别器中加入注意力层，更好地捕捉全局依赖关系

# 卷积生成器示例 class ConvGenerator(nn.Module): def __init__(self): super().__init__() self.label_emb = nn.Embedding(10, 10) self.main = nn.Sequential( # 输入: (latent_dim + 10) x 1 x 1 nn.ConvTranspose2d(110, 512, 4, 1, 0, bias=False), nn.BatchNorm2d(512), nn.ReLU(), # 4x4 nn.ConvTranspose2d(512, 256, 4, 2, 1, bias=False), nn.BatchNorm2d(256), nn.ReLU(), # 8x8 nn.ConvTranspose2d(256, 128, 4, 2, 1, bias=False), nn.BatchNorm2d(128), nn.ReLU(), # 16x16 nn.ConvTranspose2d(128, 1, 4, 2, 1, bias=False), nn.Tanh() # 28x28 ) def forward(self, noise, labels): label_emb = self.label_emb(labels).unsqueeze(2).unsqueeze(3) combined = torch.cat([noise.view(-1,100,1,1), label_emb], dim=1) return self.main(combined)

在实际项目中，我发现调整噪声向量的维度对生成多样性影响显著。当维度从100增加到256时，生成数字的多样性明显提升，但同时也需要更长的训练时间才能达到稳定状态。另一个实用技巧是在训练初期使用较高的学习率(如0.001)，然后在第20个epoch后逐步降低，这往往能加快收敛速度而不牺牲最终质量。