Pytorch实现GAN网络

生成对抗网络（Generative Adversarial Networks, GANs）是一种深度学习模型，由Ian Goodfellow等人在2014年提出。在PyTorch中实现GAN通常涉及两个主要组件：生成器（Generator）和判别器（Discriminator）。以下是简要概述如何使用PyTorch构建基本的GAN：

1. **导入必要的库**:

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms

2. **定义生成器**:

class Generator(nn.Module):
    def __init__(self, latent_dim, output_size):
        super(Generator, self).__init__()
        # 常见的结构可能包括线性层、卷积Transpose层等
        self.net = nn.Sequential(
            nn.Linear(latent_dim, hidden_units),
            nn.ReLU(),
            nn.Linear(hidden_units, another_hidden_units),
            nn.ReLU(),
            nn.Linear(another_hidden_units, output_size)
        )

    def forward(self, input):
        return self.net(input)

3. **定义判别器**:

class Discriminator(nn.Module):
    def __init__(self, input_size):
        super(Discriminator, self).__init__()
        # 类似于生成器，但输出通常是二分类的概率
        self.net = nn.Sequential(
            nn.Linear(input_size, hidden_units),
            nn.LeakyReLU(),
            nn.Linear(hidden_units, another_hidden_units),
            nn.LeakyReLU(),
            nn.Linear(another_hidden_units, 1)  # 输出单个节点表示真实或虚假
        )

    def forward(self, input):
        return torch.sigmoid(self.net(input))

4. **训练循环**:

def train_gan(generator, discriminator, dataloader, device, epochs, lr):
    generator.train()
    discriminator.train()

    optimizer_G = optim.Adam(generator.parameters(), lr=lr)
    optimizer_D = optim.Adam(discriminator.parameters(), lr=lr)

    for epoch in range(epochs):
        for real_images in dataloader:
            # 训练判别器
            real_images = real_images.to(device)
            real_labels = torch.ones(real_images.size(0)).to(device)

            fake_images = generator(torch.randn(batch_size, latent_dim).to(device)) 
            fake_labels = torch.zeros(batch_size).to(device)

            d_loss_real = discriminator(real_images).mean()  # 判别器对真样本的损失
            d_loss_fake = discriminator(fake_images).mean()  # 判别器对假样本的损失

            d_loss = (d_loss_real - d_loss_fake).mean()
            optimizer_D.zero_grad()
            d_loss.backward()
            optimizer_D.step()

            # 训练生成器
            generator_z = torch.randn(batch_size, latent_dim).to(device)
            gen_output = generator(generator_z)

            g_loss = discriminator(gen_output).mean()
            optimizer_G.zero_grad()
            g_loss.backward()
            optimizer_G.step()

5. **保存/加载模型**:

# 保存模型
torch.save(generator.state_dict(), 'generator.pth')
torch.save(discriminator.state_dict(), 'discriminator.pth')

# 加载模型
generator.load_state_dict(torch.load('generator.pth'))
discriminator.load_state_dict(torch.load('discriminator.pth'))