写一个pytorch的GAN的代码

在PyTorch中创建一个基本的生成对抗网络（GAN）通常涉及两个主要组件：生成器（Generator）和判别器（Discriminator）。这里是一个简单的实例，它使用卷积神经网络（CNN）作为基础架构。请注意，为了完整运行此代码，你需要已经安装了PyTorch库。

import torch
import torch.nn as nn
import torch.optim as optim

# 定义生成器
class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()
        self.main = nn.Sequential(
            # 添加几个卷积层和反卷积层
            nn.ConvTranspose2d(100, 256, 4, 1, 0, bias=False),
            nn.BatchNorm2d(256),
            nn.ReLU(True),
            nn.ConvTranspose2d(256, 128, 4, 2, 1, bias=False),
            nn.BatchNorm2d(128),
            nn.ReLU(True),
            nn.ConvTranspose2d(128, 64, 4, 2, 1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(True),
            nn.ConvTranspose2d(64, 1, 4, 2, 1, bias=True),
            nn.Tanh()
        )

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

# 定义判别器
class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.main = nn.Sequential(
            # 添加几个卷积层，最后连接到一个Sigmoid激活函数用于二分类
            nn.Conv2d(1, 64, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(64, 128, 4, 2, 1, bias=False),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(128, 256, 4, 2, 1, bias=False),
            nn.BatchNorm2d(256),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(256, 1, 4, 1, 0, bias=True)
        )

    def forward(self, input):
        validity = self.main(input).view(-1)
        return validity

# 初始化模型、优化器和损失函数
generator = Generator()
discriminator = Discriminator()

optimizer_G = optim.Adam(generator.parameters(), lr=0.0002, betas=(0.5, 0.999))
optimizer_D = optim.Adam(discriminator.parameters(), lr=0.0002, betas=(0.5, 0.999))

criterion = nn.BCELoss()

# GAN训练循环
for _ in range(num_epochs):  # 假设num_epochs是你想要训练的轮数
    for i, (real_images, _) in enumerate(train_dataloader):
        # 训练判别器
        discriminator.zero_grad()
        real_labels = torch.ones(real_images.size(0), 1)
        fake_images = generator(torch.randn(real_images.size(0), 100, 1, 1))  # 假数据
        fake_labels = torch.zeros(fake_images.size(0), 1)
        
        output_real = discriminator(real_images).view(-1)
        loss_real = criterion(output_real, real_labels)
        
        output_fake = discriminator(fake_images).view(-1)
        loss_fake = criterion(output_fake, fake_labels)
        
        loss_D = (loss_real + loss_fake) / 2
        loss_D.backward()
        optimizer_D.step()

        # 训练生成器
        generator.zero_grad()
        noise = torch.randn(real_images.size(0), 100, 1, 1)
        output = discriminator(fake_images)
        loss_G = criterion(output, torch.ones(output.size(0), 1))
        loss_G.backward()
        optimizer_G.step()

        print(f"Epoch [{_}/{num_epochs}], Step [{i+1}/{len(train_dataloader)}], Loss D: {loss_D.item():.4f}, Loss G: {loss_G.item():.4f}")

# 相关问题--
1. PyTorch GAN训练过程中如何调整学习率？
2. 如何保存和加载预训练的GAN模型？
3. 在GAN训练时如何防止模式崩溃？