Python案例如何用GAN生成逼真图像

本文目录导读：

1. 基础GAN架构实现
2. 训练GAN模型
3. 完整的训练脚本
4. 生成新图像
5. 高级GAN变体实现
6. 在真实数据集上训练
7. 使用预训练模型
8. 实用技巧和建议

我将为您提供一个基于GAN生成逼真图像的Python案例,使用DCGAN（深度卷积生成对抗网络）生成手写数字和自然图像。

基础GAN架构实现

1 导入必要的库

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
import matplotlib.pyplot as plt
import numpy as np
from tqdm import tqdm
import os
# 设置随机种子
torch.manual_seed(42)
np.random.seed(42)
# 检查设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

2 定义生成器网络

class Generator(nn.Module):
    def __init__(self, latent_dim=100, img_channels=1, feature_dim=64):
        super(Generator, self).__init__()
        self.main = nn.Sequential(
            # 输入：latent_dim x 1 x 1
            nn.ConvTranspose2d(latent_dim, feature_dim * 8, 4, 1, 0, bias=False),
            nn.BatchNorm2d(feature_dim * 8),
            nn.ReLU(True),
            # 状态：feature_dim*8 x 4 x 4
            nn.ConvTranspose2d(feature_dim * 8, feature_dim * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(feature_dim * 4),
            nn.ReLU(True),
            # 状态：feature_dim*4 x 8 x 8
            nn.ConvTranspose2d(feature_dim * 4, feature_dim * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(feature_dim * 2),
            nn.ReLU(True),
            # 状态：feature_dim*2 x 16 x 16
            nn.ConvTranspose2d(feature_dim * 2, feature_dim, 4, 2, 1, bias=False),
            nn.BatchNorm2d(feature_dim),
            nn.ReLU(True),
            # 状态：feature_dim x 32 x 32
            nn.ConvTranspose2d(feature_dim, img_channels, 4, 2, 1, bias=False),
            nn.Tanh()
            # 输出：img_channels x 64 x 64
        )
    def forward(self, x):
        return self.main(x)

3 定义判别器网络

class Discriminator(nn.Module):
    def __init__(self, img_channels=1, feature_dim=64):
        super(Discriminator, self).__init__()
        self.main = nn.Sequential(
            # 输入：img_channels x 64 x 64
            nn.Conv2d(img_channels, feature_dim, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            # 状态：feature_dim x 32 x 32
            nn.Conv2d(feature_dim, feature_dim * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(feature_dim * 2),
            nn.LeakyReLU(0.2, inplace=True),
            # 状态：feature_dim*2 x 16 x 16
            nn.Conv2d(feature_dim * 2, feature_dim * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(feature_dim * 4),
            nn.LeakyReLU(0.2, inplace=True),
            # 状态：feature_dim*4 x 8 x 8
            nn.Conv2d(feature_dim * 4, feature_dim * 8, 4, 2, 1, bias=False),
            nn.BatchNorm2d(feature_dim * 8),
            nn.LeakyReLU(0.2, inplace=True),
            # 状态：feature_dim*8 x 4 x 4
            nn.Conv2d(feature_dim * 8, 1, 4, 1, 0, bias=False),
            nn.Sigmoid()
        )
    def forward(self, x):
        return self.main(x).view(-1, 1)

4 权重初始化

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)

训练GAN模型

1 加载和预处理数据

def load_data(batch_size=128, image_size=64):
    # 对于MNIST数据集
    transform = transforms.Compose([
        transforms.Resize(image_size),
        transforms.ToTensor(),
        transforms.Normalize([0.5], [0.5])  # 归一化到[-1, 1]
    ])
    dataset = MNIST(root='./data', train=True, download=True, transform=transform)
    dataloader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=2
    )
    return dataloader

2 训练函数

def train_gan(generator, discriminator, dataloader, num_epochs=50, latent_dim=100):
    # 损失函数
    criterion = nn.BCELoss()
    # 优化器
    lr = 0.0002
    beta1 = 0.5
    optimizer_G = optim.Adam(generator.parameters(), lr=lr, betas=(beta1, 0.999))
    optimizer_D = optim.Adam(discriminator.parameters(), lr=lr, betas=(beta1, 0.999))
    # 记录损失
    g_losses = []
    d_losses = []
    # 固定噪声用于生成样本
    fixed_noise = torch.randn(64, latent_dim, 1, 1, device=device)
    print("开始训练...")
    for epoch in range(num_epochs):
        for i, (real_images, _) in enumerate(tqdm(dataloader, desc=f'Epoch {epoch+1}/{num_epochs}')):
            batch_size = real_images.size(0)
            # 将数据移到设备
            real_images = real_images.to(device)
            # 创建标签
            real_labels = torch.full((batch_size, 1), 1.0, device=device)
            fake_labels = torch.full((batch_size, 1), 0.0, device=device)
            # ==================== 训练判别器 ====================
            # 冻结生成器
            generator.zero_grad()
            discriminator.zero_grad()
            # 训练判别器使用真实图像
            output_real = discriminator(real_images)
            d_loss_real = criterion(output_real, real_labels)
            # 训练判别器使用生成图像
            noise = torch.randn(batch_size, latent_dim, 1, 1, device=device)
            fake_images = generator(noise)
            output_fake = discriminator(fake_images.detach())
            d_loss_fake = criterion(output_fake, fake_labels)
            # 总判别器损失
            d_loss = d_loss_real + d_loss_fake
            d_loss.backward()
            optimizer_D.step()
            # ==================== 训练生成器 ====================
            # 生成新的噪声
            noise = torch.randn(batch_size, latent_dim, 1, 1, device=device)
            fake_images = generator(noise)
            # 试图欺骗判别器
            output = discriminator(fake_images)
            g_loss = criterion(output, real_labels)
            # 反向传播和优化
            g_loss.backward()
            optimizer_G.step()
            # 记录损失
            if i % 100 == 0:
                g_losses.append(g_loss.item())
                d_losses.append(d_loss.item())
        # 每个epoch后生成样本
        with torch.no_grad():
            fake = generator(fixed_noise).detach().cpu()
        # 保存生成的图像
        if (epoch + 1) % 10 == 0:
            save_generated_images(fake, epoch+1)
            print(f'Epoch [{epoch+1}/{num_epochs}] Loss_D: {d_loss.item():.4f} Loss_G: {g_loss.item():.4f}')
    return g_losses, d_losses

3 可视化功能

def save_generated_images(images, epoch, save_dir='generated_images'):
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    # 创建网格
    grid = torchvision.utils.make_grid(images, nrow=8, normalize=True)
    # 转换为numpy数组用于显示
    grid_np = np.transpose(grid.numpy(), (1, 2, 0))
    # 保存图像
    plt.figure(figsize=(10, 10))
    plt.imshow(grid_np)
    plt.axis('off')
    plt.title(f'Generated Images - Epoch {epoch}')
    plt.savefig(f'{save_dir}/epoch_{epoch}.png')
    plt.close()
def plot_losses(g_losses, d_losses):
    plt.figure(figsize=(10, 5))
    plt.plot(g_losses, label='Generator Loss')
    plt.plot(d_losses, label='Discriminator Loss')
    plt.xlabel('Iterations')
    plt.ylabel('Loss')
    plt.legend()
    plt.title('GAN Training Losses')
    plt.savefig('training_losses.png')
    plt.show()

完整的训练脚本

def main():
    # 参数设置
    latent_dim = 100
    batch_size = 128
    image_size = 64
    num_epochs = 50
    img_channels = 1  # MNIST是灰度图像
    # 加载数据
    dataloader = load_data(batch_size, image_size)
    # 初始化模型
    generator = Generator(latent_dim, img_channels).to(device)
    discriminator = Discriminator(img_channels).to(device)
    # 应用权重初始化
    generator.apply(weights_init)
    discriminator.apply(weights_init)
    # 打印模型结构
    print("Generator结构:")
    print(generator)
    print("\nDiscriminator结构:")
    print(discriminator)
    # 计算参数量
    total_params_g = sum(p.numel() for p in generator.parameters())
    total_params_d = sum(p.numel() for p in discriminator.parameters())
    print(f"\n生成器参数量: {total_params_g:,}")
    print(f"判别器参数量: {total_params_d:,}")
    # 训练
    g_losses, d_losses = train_gan(
        generator, 
        discriminator, 
        dataloader, 
        num_epochs, 
        latent_dim
    )
    # 绘制损失曲线
    plot_losses(g_losses, d_losses)
    # 保存模型
    torch.save(generator.state_dict(), 'generator.pth')
    torch.save(discriminator.state_dict(), 'discriminator.pth')
    print("模型已保存!")
if __name__ == "__main__":
    main()

生成新图像

def generate_images(model_path='generator.pth', num_images=16, latent_dim=100, img_channels=1):
    # 加载模型
    generator = Generator(latent_dim, img_channels).to(device)
    generator.load_state_dict(torch.load(model_path))
    generator.eval()
    # 生成随机噪声
    noise = torch.randn(num_images, latent_dim, 1, 1, device=device)
    # 生成图像
    with torch.no_grad():
        generated_images = generator(noise)
    return generated_images
# 使用示例
def demo_generation():
    # 生成图像
    images = generate_images()
    # 显示图像
    grid = torchvision.utils.make_grid(images.cpu(), nrow=4, normalize=True)
    grid_np = np.transpose(grid.numpy(), (1, 2, 0))
    plt.figure(figsize=(8, 8))
    plt.imshow(grid_np)
    plt.axis('off')
    plt.title('Generated Images')
    plt.show()

高级GAN变体实现

1 条件GAN (Conditional GAN)

class ConditionalGenerator(nn.Module):
    def __init__(self, latent_dim=100, num_classes=10, img_channels=1, feature_dim=64):
        super(ConditionalGenerator, self).__init__()
        self.label_embedding = nn.Embedding(num_classes, num_classes)
        self.main = nn.Sequential(
            nn.ConvTranspose2d(latent_dim + num_classes, feature_dim * 8, 4, 1, 0, bias=False),
            nn.BatchNorm2d(feature_dim * 8),
            nn.ReLU(True),
            nn.ConvTranspose2d(feature_dim * 8, feature_dim * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(feature_dim * 4),
            nn.ReLU(True),
            nn.ConvTranspose2d(feature_dim * 4, feature_dim * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(feature_dim * 2),
            nn.ReLU(True),
            nn.ConvTranspose2d(feature_dim * 2, feature_dim, 4, 2, 1, bias=False),
            nn.BatchNorm2d(feature_dim),
            nn.ReLU(True),
            nn.ConvTranspose2d(feature_dim, img_channels, 4, 2, 1, bias=False),
            nn.Tanh()
        )
    def forward(self, noise, labels):
        # 将标签转换为嵌入向量
        label_emb = self.label_embedding(labels)
        label_emb = label_emb.view(noise.size(0), -1, 1, 1)
        label_emb = label_emb.expand(-1, -1, noise.size(2), noise.size(3))
        # 连接噪声和标签
        combined = torch.cat([noise, label_emb], dim=1)
        return self.main(combined)

2 使用WGAN-GP提升训练稳定性

class WGANGPTrainer:
    def __init__(self, generator, discriminator, lambda_gp=10, n_critic=5):
        self.generator = generator
        self.discriminator = discriminator
        self.lambda_gp = lambda_gp
        self.n_critic = n_critic
        self.criterion = nn.BCELoss()
        self.optimizer_G = optim.Adam(generator.parameters(), lr=0.0001, betas=(0.5, 0.9))
        self.optimizer_D = optim.Adam(discriminator.parameters(), lr=0.0001, betas=(0.5, 0.9))
    def compute_gradient_penalty(self, real_images, fake_images):
        batch_size = real_images.size(0)
        epsilon = torch.rand(batch_size, 1, 1, 1, device=device)
        epsilon = epsilon.expand_as(real_images)
        # 插值图像
        interpolated = epsilon * real_images + (1 - epsilon) * fake_images
        interpolated.requires_grad_(True)
        # 计算判别器输出
        d_interpolated = self.discriminator(interpolated)
        # 计算梯度
        gradients = torch.autograd.grad(
            outputs=d_interpolated,
            inputs=interpolated,
            grad_outputs=torch.ones_like(d_interpolated),
            create_graph=True,
            retain_graph=True
        )[0]
        # 计算梯度惩罚
        gradients = gradients.view(batch_size, -1)
        gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean()
        return gradient_penalty
    def train_step(self, real_images, labels=None):
        batch_size = real_images.size(0)
        latent_dim = 100
        # 训练判别器
        for _ in range(self.n_critic):
            self.discriminator.zero_grad()
            # 真实图像的损失
            d_real = self.discriminator(real_images)
            d_real_loss = -torch.mean(d_real)
            # 生成虚假图像
            noise = torch.randn(batch_size, latent_dim, 1, 1, device=device)
            fake_images = self.generator(noise)
            # 虚假图像的损失
            d_fake = self.discriminator(fake_images.detach())
            d_fake_loss = torch.mean(d_fake)
            # 梯度惩罚
            gradient_penalty = self.compute_gradient_penalty(real_images, fake_images.detach())
            # 总判别器损失
            d_loss = d_real_loss + d_fake_loss + self.lambda_gp * gradient_penalty
            d_loss.backward()
            self.optimizer_D.step()
        # 训练生成器
        self.generator.zero_grad()
        noise = torch.randn(batch_size, latent_dim, 1, 1, device=device)
        fake_images = self.generator(noise)
        g_loss = -torch.mean(self.discriminator(fake_images))
        g_loss.backward()
        self.optimizer_G.step()
        return g_loss.item(), d_loss.item()

在真实数据集上训练

def train_on_celebA():
    """
    在CelebA数据集上训练GAN生成人脸图像
    """
    # 数据集准备
    from torchvision.datasets import CelebA
    transform = transforms.Compose([
        transforms.Resize(64),
        transforms.CenterCrop(64),
        transforms.ToTensor(),
        transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
    ])
    dataset = CelebA(
        root='./data',
        split='train',
        transform=transform,
        download=True
    )
    dataloader = DataLoader(
        dataset,
        batch_size=64,
        shuffle=True,
        num_workers=2
    )
    # 初始化模型（RGB图像，3个通道）
    generator = Generator(latent_dim=100, img_channels=3).to(device)
    discriminator = Discriminator(img_channels=3).to(device)
    # 训练
    g_losses, d_losses = train_gan(
        generator,
        discriminator,
        dataloader,
        num_epochs=100,
        latent_dim=100
    )
    return generator, discriminator

使用预训练模型

def use_pretrained_stylegan():
    """
    使用预训练的StyleGAN2模型生成高分辨率图像
    """
    try:
        import dnnlib
        import legacy
        import pickle
        # 下载预训练模型
        url = 'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/ffhq.pkl'
        with dnnlib.util.open_url(url) as f:
            G = legacy.load_network_pkl(f)['G_ema']
        # 将模型移到GPU
        G = G.to(device)
        # 生成随机潜在向量
        z = torch.randn([4, G.z_dim], device=device)
        # 生成图像
        with torch.no_grad():
            img = G(z, None)
        # 保存生成的图像
        for i in range(4):
            img_np = (img[i].cpu().numpy().transpose(1, 2, 0) + 1) / 2
            plt.imsave(f'stylegan_output_{i}.png', img_np)
    except ImportError:
        print("请安装stylegan2-ada-pytorch库")

实用技巧和建议

1 训练稳定性的建议

# 添加标签平滑
def smooth_positive_labels(labels, smoothing=0.9):
    return labels * smoothing
# 添加噪声
def add_noise_to_images(images, noise_factor=0.1):
    noise = torch.randn_like(images) * noise_factor
    return images + noise
# 学习率衰减
def adjust_learning_rate(optimizer, epoch, initial_lr=0.0002):
    lr = initial_lr * (0.5 ** (epoch // 10))
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

2 图像质量评估

def calculate_fid(real_features, generated_features):
    """
    计算FID (Fréchet Inception Distance) 分数
    """
    from scipy import linalg
    # 计算均值和协方差
    mu_real = np.mean(real_features, axis=0)
    mu_gen = np.mean(generated_features, axis=0)
    cov_real = np.cov(real_features, rowvar=False)
    cov_gen = np.cov(generated_features, rowvar=False)
    # 计算FID
    diff = mu_real - mu_gen
    covmean, _ = linalg.sqrtm(cov_real.dot(cov_gen), disp=False)
    if np.iscomplexobj(covmean):
        covmean = covmean.real
    fid = diff.dot(diff) + np.trace(cov_real + cov_gen - 2 * covmean)
    return fid

这个完整的教程展示了如何使用PyTorch实现GAN生成逼真图像,从基础的DCGAN到更高级的WGAN-GP和条件GAN，您可以根据需求选择适合的架构，记得调整参数以适应您的具体数据集和计算资源。