GAN（生成对抗网络）入门，生成图片，基于pytorch，python代码

主要任务是生成手写数字图像。
在每个训练周期结束时，生成器会生成一批假的手写数字图像。
随着训练的进行，生成的图像通常会变得越来越逼真，

代码会自下载数据集，可直接运行，代码如下。

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import torchvision.utils as vutils
import matplotlib.pyplot as plt
from PIL import Image


# 设置随机种子以便复现结果
torch.manual_seed(42)

# 超参数
batch_size = 64
image_size = 28 * 28
latent_dim = 100
epochs = 5

# 数据加载和预处理
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

dataset = datasets.MNIST(root="./data", train=True, transform=transform, download=True)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)


# 显示数据集中的一个样本图像
data_iter = iter(dataloader)
sample_data = next(data_iter)
sample_image, _ = sample_data
sample_image = sample_image.view(sample_image.size(0), 1, 28, 28)  # 将图像大小调整为[1, 28, 28]
plt.imshow(sample_image[0, 0].numpy(), cmap='gray')
plt.axis('off')
plt.title('Sample Image from the Dataset')
plt.show()

# 显示数据集中的一个样本图像
data_iter = iter(dataloader)
sample_data = next(data_iter)
sample_image, _ = sample_data
sample_image = sample_image.view(sample_image.size(0), 1, 28, 28)  # 将图像大小调整为[1, 28, 28]
plt.imshow(sample_image[1, 0].numpy(), cmap='gray')
plt.axis('off')
plt.title('Sample Image from the Dataset')
plt.show()

# 显示数据集中的一个样本图像
data_iter = iter(dataloader)
sample_data = next(data_iter)
sample_image, _ = sample_data
sample_image = sample_image.view(sample_image.size(0), 1, 28, 28)  # 将图像大小调整为[1, 28, 28]
plt.imshow(sample_image[2, 0].numpy(), cmap='gray')
plt.axis('off')
plt.title('Sample Image from the Dataset')
plt.show()
# 生成器模型
class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()
        self.main = nn.Sequential(
            nn.Linear(latent_dim, 256),
            nn.ReLU(),
            nn.Linear(256, 512),
            nn.ReLU(),
            nn.Linear(512, 1024),
            nn.ReLU(),
            nn.Linear(1024, image_size),
            nn.Tanh()
        )

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


# 判别器模型
class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.main = nn.Sequential(
            nn.Linear(image_size, 1024),
            nn.LeakyReLU(0.2),
            nn.Dropout(0.3),
            nn.Linear(1024, 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, x):
        return self.main(x)


# 将模型移到GPU上（如果可用）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
generator = Generator().to(device)
discriminator = Discriminator().to(device)

# 初始化模型、损失函数和优化器
criterion = nn.BCELoss()
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))

# 训练GAN
for epoch in range(epochs):
    for i, data in enumerate(dataloader, 0):
        real_images, _ = data
        real_images = real_images.view(real_images.size(0), -1).to(device)

        # 训练判别器
        discriminator.zero_grad()
        real_labels = torch.ones(real_images.size(0), 1).to(device)
        fake_labels = torch.zeros(real_images.size(0), 1).to(device)

        outputs_real = discriminator(real_images)
        d_loss_real = criterion(outputs_real, real_labels)
        d_loss_real.backward()

        z = torch.randn(real_images.size(0), latent_dim).to(device)
        fake_images = generator(z)
        outputs_fake = discriminator(fake_images.detach())
        d_loss_fake = criterion(outputs_fake, fake_labels)
        d_loss_fake.backward()

        d_loss = d_loss_real + d_loss_fake
        optimizer_d.step()

        # 训练生成器
        generator.zero_grad()
        outputs = discriminator(fake_images)
        g_loss = criterion(outputs, real_labels)
        g_loss.backward()
        optimizer_g.step()

        if i % 100 == 0:
            print(f"Epoch [{epoch + 1}/{epochs}] Batch {i}/{len(dataloader)} "
                  f"Loss D: {d_loss.item():.4f}, Loss G: {g_loss.item():.4f}")

    # 保存生成的图像
    with torch.no_grad():
        fake_images = generator(torch.randn(batch_size, latent_dim).to(device))
        fake_images = fake_images.view(fake_images.size(0), 1, 28, 28)
        vutils.save_image(fake_images, f"gan_generated_epoch_{epoch + 1}.png", normalize=True)
        # 显示生成的图像
        def plot_generated_images(image_path):
            img = Image.open(image_path)
            plt.imshow(img)
            plt.axis('off')
            plt.show()

        plot_generated_images(f"gan_generated_epoch_{epoch + 1}.png")
print("训练完成。生成的图像已保存为 gan_generated_epoch_[epoch_number].png。")