PyTorch实战：手把手教你用GAN生成‘以假乱真’的MNIST数字，并打包成新Dataset

PyTorch实战：从GAN生成到Dataset封装的全流程工程指南

在深度学习项目中，数据永远是核心。但现实情况往往是：标注数据不足、样本分布不均衡、数据多样性有限。传统的数据增强方法（如旋转、裁剪）只能提供有限的多样性扩展。这时候，生成对抗网络（GAN）为我们打开了一扇新的大门——不仅能生成逼真的数据，还能将这些数据无缝集成到现有训练流程中。

本文将带你走完从GAN训练到工程落地的完整闭环。不同于大多数教程止步于模型训练，我们将重点解决"生成之后怎么办"这个实际问题：

1. 如何批量生成特定类别的样本（比如每个数字500张）
2. 如何自动保存和组织生成结果
3. 如何将这些生成数据封装成PyTorch原生的Dataset对象
4. 如何评估生成数据的质量和对模型训练的贡献

1. 环境准备与基础模型搭建

1.1 安装依赖与数据加载

首先确保你的环境已安装最新版PyTorch（建议1.8+版本）。我们将使用MNIST作为基础数据集，但方法论适用于任何图像生成任务。

import torch import torch.nn as nn import torchvision from torchvision import datasets, transforms from torch.utils.data import Dataset, DataLoader import numpy as np import os from PIL import Image import matplotlib.pyplot as plt # 基础配置 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") batch_size = 64 latent_dim = 100 num_classes = 10 img_size = 28 channels = 1

1.2 构建条件GAN模型

我们将实现一个带条件标签的DCGAN（深度卷积生成对抗网络），让生成器能够按需生成特定数字：

class Generator(nn.Module): def __init__(self): super(Generator, self).__init__() self.label_emb = nn.Embedding(num_classes, latent_dim) self.model = nn.Sequential( nn.Linear(latent_dim*2, 128*7*7), nn.LeakyReLU(0.2, inplace=True), nn.Unflatten(1, (128, 7, 7)), nn.ConvTranspose2d(128, 64, 4, 2, 1), nn.BatchNorm2d(64), nn.LeakyReLU(0.2, inplace=True), nn.ConvTranspose2d(64, channels, 4, 2, 1), nn.Tanh() ) def forward(self, noise, labels): gen_input = torch.cat((self.label_emb(labels), noise), -1) img = self.model(gen_input) return img

判别器的实现同样需要考虑类别信息：

class Discriminator(nn.Module): def __init__(self): super(Discriminator, self).__init__() self.label_emb = nn.Embedding(num_classes, img_size*img_size) self.model = nn.Sequential( nn.Conv2d(channels+1, 64, 4, 2, 1), nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(64, 128, 4, 2, 1), nn.BatchNorm2d(128), nn.LeakyReLU(0.2, inplace=True), nn.Flatten(), nn.Dropout(0.4), nn.Linear(128*7*7, 1), nn.Sigmoid() ) def forward(self, img, labels): label_emb = self.label_emb(labels).view(img.size(0), 1, img_size, img_size) d_in = torch.cat((img, label_emb), 1) validity = self.model(d_in) return validity

2. 模型训练与样本生成策略

2.1 训练循环实现

训练条件GAN需要特别注意标签信息的处理。以下是关键训练步骤：

# 初始化模型 generator = Generator().to(device) discriminator = Discriminator().to(device) # 定义损失函数和优化器 adversarial_loss = nn.BCELoss() optimizer_G = torch.optim.Adam(generator.parameters(), lr=0.0002, betas=(0.5, 0.999)) optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=0.0002, betas=(0.5, 0.999)) for epoch in range(num_epochs): for i, (imgs, labels) in enumerate(dataloader): # 真实数据准备 real_imgs = imgs.to(device) real_labels = labels.to(device) valid = torch.ones((imgs.size(0), 1)).to(device) fake = torch.zeros((imgs.size(0), 1)).to(device) # 训练生成器 optimizer_G.zero_grad() z = torch.randn(imgs.size(0), latent_dim).to(device) gen_labels = torch.randint(0, num_classes, (imgs.size(0),)).to(device) gen_imgs = generator(z, gen_labels) g_loss = adversarial_loss(discriminator(gen_imgs, gen_labels), valid) g_loss.backward() optimizer_G.step() # 训练判别器 optimizer_D.zero_grad() real_loss = adversarial_loss(discriminator(real_imgs, real_labels), valid) fake_loss = adversarial_loss(discriminator(gen_imgs.detach(), gen_labels), fake) d_loss = (real_loss + fake_loss) / 2 d_loss.backward() optimizer_D.step()

2.2 可控样本生成技术

训练完成后，我们可以按需生成特定类别的样本。以下函数可以批量生成指定类别的数字：

def generate_samples(generator, num_samples, target_label, save_dir=None): """生成指定类别的样本并可选保存""" generator.eval() z = torch.randn(num_samples, latent_dim).to(device) labels = torch.full((num_samples,), target_label, dtype=torch.long).to(device) with torch.no_grad(): gen_imgs = generator(z, labels) # 将生成的张量转换为图像 gen_imgs = 0.5 * gen_imgs + 0.5 # 从[-1,1]转换到[0,1] gen_imgs = gen_imgs.cpu().numpy() if save_dir: os.makedirs(save_dir, exist_ok=True) for i in range(num_samples): img = (gen_imgs[i].transpose(1, 2, 0) * 255).astype(np.uint8) img = Image.fromarray(img.squeeze()) img.save(os.path.join(save_dir, f"{target_label}_{i}.png")) return gen_imgs

提示：生成样本时建议使用generator.eval()模式，并配合torch.no_grad()上下文管理器，这样可以减少内存消耗并提高生成速度。

3. 生成数据的工程化处理

3.1 自动化数据流水线

为了实现大规模数据生成，我们需要建立一个自动化流程。以下脚本可以生成所有数字类别的平衡数据集：

def generate_full_dataset(generator, samples_per_class, output_dir): """生成平衡的MNIST风格数据集""" for label in range(num_classes): print(f"Generating {samples_per_class} samples for digit {label}") generate_samples( generator, samples_per_class, label, os.path.join(output_dir, str(label)) ) # 创建标签文件 with open(os.path.join(output_dir, "labels.csv"), "w") as f: for label in range(num_classes): for i in range(samples_per_class): f.write(f"{label}/{label}_{i}.png,{label}\n")

执行这个函数将创建一个结构化的数据集目录：

generated_mnist/ ├── 0/ │ ├── 0_0.png │ ├── 0_1.png │ └── ... ├── 1/ │ ├── 1_0.png │ └── ... ├── ... └── labels.csv

3.2 数据质量评估指标

在将生成数据用于训练前，建议进行质量评估。常用的评估指标包括：

对于MNIST这样的简单数据集，我们可以实现一个轻量级的评估方法：

def evaluate_generated_data(generator, test_loader): """使用预训练分类器评估生成数据质量""" classifier = torch.load("pretrained_mnist_classifier.pth").to(device) classifier.eval() all_labels = [] all_preds = [] for _ in range(100): # 评估100个批次 z = torch.randn(batch_size, latent_dim).to(device) labels = torch.randint(0, num_classes, (batch_size,)).to(device) with torch.no_grad(): gen_imgs = generator(z, labels) preds = classifier(gen_imgs).argmax(dim=1) all_labels.append(labels.cpu()) all_preds.append(preds.cpu()) accuracy = (torch.cat(all_preds) == torch.cat(all_labels)).float().mean() print(f"Classifier accuracy on generated data: {accuracy.item():.2%}") return accuracy

4. 创建PyTorch Dataset类

4.1 自定义Dataset实现

为了让生成的数据能够无缝接入现有训练流程，我们需要实现一个标准的Dataset类：

class GeneratedMNIST(Dataset): def __init__(self, root_dir, transform=None): """ 参数: root_dir (string): 包含生成数据的目录 transform (callable, optional): 应用于样本的可选变换 """ self.root_dir = root_dir self.transform = transform # 加载标签文件 self.samples = [] with open(os.path.join(root_dir, "labels.csv"), "r") as f: for line in f: img_path, label = line.strip().split(",") self.samples.append((img_path, int(label))) def __len__(self): return len(self.samples) def __getitem__(self, idx): img_path, label = self.samples[idx] img = Image.open(os.path.join(self.root_dir, img_path)) if self.transform: img = self.transform(img) return img, label

4.2 数据加载与增强

现在我们可以像使用标准MNIST数据集一样使用生成的数据：

# 定义数据变换 transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,)) ]) # 创建数据集实例 generated_dataset = GeneratedMNIST( root_dir="generated_mnist", transform=transform ) # 创建数据加载器 generated_loader = DataLoader( generated_dataset, batch_size=64, shuffle=True, num_workers=4 ) # 也可以混合真实和生成数据 real_dataset = datasets.MNIST( root="data", train=True, download=True, transform=transform ) mixed_dataset = torch.utils.data.ConcatDataset([real_dataset, generated_dataset]) mixed_loader = DataLoader(mixed_dataset, batch_size=64, shuffle=True)

4.3 实际应用效果对比

为了验证生成数据的价值，我们可以进行一个简单的对比实验：

从实验结果可以看出，混合使用真实和生成数据可以获得最佳平衡——既提高了模型性能，又减少了过拟合风险。