山东大学软件学院ai导论实验之生成对抗网络-EW帮帮网

实验目的

基于Pytorch搭建一个生成对抗网络，使用MNIST数据集。

实验代码

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import numpy as np
import matplotlib.pyplot as plt
import os

# 设置环境变量
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

# 创建保存生成图像的文件夹
output_path = r"xxxxxxxxxxxxxxxxxx"
os.makedirs(output_path, exist_ok=True)


# 生成器网络
class Generator(nn.Module):
    def __init__(self, latent_dim):
        super(Generator, self).__init__()
        self.network = nn.Sequential(
            nn.Linear(latent_dim, 256),
            nn.LeakyReLU(0.2),
            nn.Linear(256, 512),
            nn.LeakyReLU(0.2),
            nn.Linear(512, 1024),
            nn.LeakyReLU(0.2),
            nn.Linear(1024, 784),
            nn.Tanh()
        )

    def forward(self, z):
        img = self.network(z)
        return img.view(img.size(0), 1, 28, 28)


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

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


def generate_and_save_images(generator, test_input, epoch, img_path):
    with torch.no_grad():
        generated_images = generator(test_input).cpu().numpy()

    fig, axes = plt.subplots(4, 4, figsize=(4, 4))
    for i, ax in enumerate(axes.flat):
        # 将图像从形状 (1, 28, 28) 转换为 (28, 28)，去除通道维度
        ax.imshow(np.squeeze(generated_images[i]), cmap='gray')
        ax.axis('off')

    img_filename = os.path.join(img_path, f"generated_epoch_{epoch}.png")
    plt.tight_layout()
    plt.savefig(img_filename)
    plt.close()


# 设置设备（使用GPU或CPU）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 超参数
lr = 0.0001
batch_size = 128
latent_dim = 100
epochs = 2000

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

train_dataset = datasets.MNIST(root='./MNIST_data', train=True, transform=transform, download=True)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

# 测试数据：随机噪声作为输入
test_data = torch.randn(batch_size, latent_dim).to(device)

# 初始化生成器和判别器，并定义损失函数和优化器
generator = Generator(latent_dim).to(device)
discriminator = Discriminator().to(device)
adversarial_loss = nn.BCELoss()
optimizer_G = optim.Adam(generator.parameters(), lr=lr)
optimizer_D = optim.Adam(discriminator.parameters(), lr=lr)

# 记录损失
D_losses = []
G_losses = []

# 训练过程
for epoch in range(epochs):
    for i, (imgs, _) in enumerate(train_loader):
        real_imgs = imgs.to(device)
        batch_size = real_imgs.size(0)

        # 判别器训练
        z = torch.randn(batch_size, latent_dim).to(device)
        fake_imgs = generator(z)

        real_labels = torch.ones(batch_size, 1).to(device)
        fake_labels = torch.zeros(batch_size, 1).to(device)

        # 计算损失
        real_loss = adversarial_loss(discriminator(real_imgs), real_labels)
        fake_loss = adversarial_loss(discriminator(fake_imgs.detach()), fake_labels)
        d_loss = (real_loss + fake_loss) / 2

        optimizer_D.zero_grad()
        d_loss.backward()
        optimizer_D.step()

        # 生成器训练
        z = torch.randn(batch_size, latent_dim).to(device)
        fake_imgs = generator(z)
        g_loss = adversarial_loss(discriminator(fake_imgs), real_labels)

        optimizer_G.zero_grad()
        g_loss.backward()
        optimizer_G.step()

        # 记录损失
        D_losses.append(d_loss.item())
        G_losses.append(g_loss.item())

        # 打印每2000个步骤的迭代信息
        if (epoch * len(train_loader) + i) % 2000 == 0:
            print(f"Iter: {epoch * len(train_loader) + i}")
            print(f"D_loss: {d_loss.item():.4f}")
            print(f"G_loss: {g_loss.item():.4f}")
    # 每个epoch保存生成的图像
    generate_and_save_images(generator, test_data, epoch, output_path)

    # 保存生成器和判别器的模型
    torch.save(generator.state_dict(), "Generator_mnist.pth")
    torch.save(discriminator.state_dict(), "Discriminator_mnist.pth")

# 绘制损失曲线
plt.figure(figsize=(10, 5))
plt.plot(D_losses, label='Discriminator Loss')
plt.plot(G_losses, label='Generator Loss')
plt.xlabel('Iterations')
plt.ylabel('Loss')
plt.legend()
plt.title('Loss Curve')
plt.savefig('loss_curve.png')  # 保存图像
plt.show()  # 显示图像