4. 基于GAN的图像数据合成

任务说明

本案例将训练一个简单的生成对抗网络（GAN），使用MNIST数据集学习手写数字的分布，并生成逼真的手写数字图像。完成本案例后，你将拥有一个可运行的GAN训练脚本，能够输出生成样本的可视化结果。

环境依赖

• Python 3.8+
• PyTorch 1.12+（含torchvision）
• Matplotlib
• NumPy

核心实现

以下代码实现了完整的GAN训练流程：生成器将100维随机噪声映射为28×28图像，判别器对真实MNIST图像与生成图像进行分类。训练采用交替优化策略，并使用BCE损失函数。

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np

# 超参数
LATENT_DIM = 100
BATCH_SIZE = 128
EPOCHS = 50
LR = 0.0002
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 数据加载
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])
train_dataset = torchvision.datasets.MNIST(
    root='./data', train=True, download=True, transform=transform
)
train_loader = torch.utils.data.DataLoader(
    train_dataset, batch_size=BATCH_SIZE, shuffle=True
)

# 生成器
class Generator(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(LATENT_DIM, 256),
            nn.ReLU(),
            nn.BatchNorm1d(256),
            nn.Linear(256, 512),
            nn.ReLU(),
            nn.BatchNorm1d(512),
            nn.Linear(512, 784),
            nn.Tanh()
        )

    def forward(self, z):
        return self.model(z).view(-1, 1, 28, 28)

# 判别器
class Discriminator(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(784, 512),
            nn.LeakyReLU(0.2),
            nn.Dropout(0.3),
            nn.Linear(512, 256),
            nn.LeakyReLU(0.2),
            nn.Dropout(0.3),
            nn.Linear(256, 1),
            nn.Sigmoid()
        )

    def forward(self, img):
        return self.model(img.view(img.size(0), -1))

# 初始化
G = Generator().to(DEVICE)
D = Discriminator().to(DEVICE)
criterion = nn.BCELoss()
opt_G = optim.Adam(G.parameters(), lr=LR, betas=(0.5, 0.999))
opt_D = optim.Adam(D.parameters(), lr=LR, betas=(0.5, 0.999))

# 训练循环
for epoch in range(EPOCHS):
    for real_imgs, _ in train_loader:
        real_imgs = real_imgs.to(DEVICE)
        batch_size = real_imgs.size(0)
        real_labels = torch.ones(batch_size, 1).to(DEVICE)
        fake_labels = torch.zeros(batch_size, 1).to(DEVICE)

        # 训练判别器
        z = torch.randn(batch_size, LATENT_DIM).to(DEVICE)
        fake_imgs = G(z)
        d_loss_real = criterion(D(real_imgs), real_labels)
        d_loss_fake = criterion(D(fake_imgs.detach()), fake_labels)
        d_loss = d_loss_real + d_loss_fake
        opt_D.zero_grad()
        d_loss.backward()
        opt_D.step()

        # 训练生成器
        z = torch.randn(batch_size, LATENT_DIM).to(DEVICE)
        fake_imgs = G(z)
        g_loss = criterion(D(fake_imgs), real_labels)
        opt_G.zero_grad()
        g_loss.backward()
        opt_G.step()

    if (epoch + 1) % 10 == 0:
        print(f'Epoch [{epoch+1}/{EPOCHS}] D_loss: {d_loss.item():.4f}, G_loss: {g_loss.item():.4f}')

# 生成样本可视化
G.eval()
with torch.no_grad():
    z = torch.randn(16, LATENT_DIM).to(DEVICE)
    samples = G(z).cpu().numpy()

fig, axes = plt.subplots(4, 4, figsize=(6, 6))
for i, ax in enumerate(axes.flat):
    ax.imshow((samples[i].squeeze() + 1) / 2, cmap='gray')
    ax.axis('off')
plt.tight_layout()
plt.savefig('generated_digits.png', dpi=150)
print("生成图像已保存为 generated_digits.png")

运行方式

pip install torch torchvision matplotlib numpy
python gan_mnist.py

验证方法

训练过程中每10个epoch会输出判别器与生成器的损失值。训练完成后，程序会在当前目录生成generated_digits.png文件。判定训练成功的标准为：

• 判别器损失稳定在0.5～1.5之间，生成器损失呈下降趋势或保持稳定
• 生成的图像中能清晰辨认出多个不同数字的轮廓，而非纯噪声或模糊色块
• 图像整体灰度分布均匀，背景与数字主体对比明显

若生成图像模糊或模式单一，可适当增加训练轮次至100轮，或调整学习率至0.0001～0.0003范围后重新训练。