gan网络pytorch代码

原创

mob649e815ddfb8 2024-10-13 05:03:17 ©著作权

©著作权归作者所有：来自51CTO博客作者mob649e815ddfb8的原创作品，请联系作者获取转载授权，否则将追究法律责任

如何使用PyTorch实现GAN网络

生成对抗网络（GAN）是一种深度学习模型，常用于生成数据。实现GAN的一般流程包括数据准备、模型构建、损失函数定义、训练过程以及结果展示。以下是实现GAN的步骤：

步骤	描述
1. 数据准备	加载并预处理数据集
2. 模型构建	创建生成器和判别器模型
3. 定义损失	选择适当的损失函数和优化器
4. 模型训练	使用循环结构训练模型
5. 展示结果	可视化生成的数据

接下来，我们将逐步实现这些步骤，每一步附上代码和注释。

1. 数据准备

首先，我们需要加载并预处理数据。这里以MNIST手写数字数据集为例：

import torch
import torchvision.transforms as transforms
from torchvision import datasets

# 定义数据转换方法，将图片转换为Tensor并归一化
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))  # 归一化操作
])

# 下载MNIST数据集
dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)

# 加载数据集
dataloader = torch.utils.data.DataLoader(dataset, batch_size=64, shuffle=True)

2. 模型构建

我们需要定义生成器（Generator）和判别器（Discriminator）的网络结构。

import torch.nn as nn

class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()
        self.model = nn.Sequential(
            nn.Linear(100, 256),  # 输入100维噪声
            nn.ReLU(),
            nn.Linear(256, 512),
            nn.ReLU(),
            nn.Linear(512, 1024),
            nn.ReLU(),
            nn.Linear(1024, 784),  # 输出784维（28*28图像）
            nn.Tanh()  # 输出图像范围[-1, 1]
        )
    
    def forward(self, z):
        return self.model(z)


class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.model = nn.Sequential(
            nn.Linear(784, 512),  # 输入784维（28*28图像）
            nn.LeakyReLU(0.2),
            nn.Linear(512, 256),
            nn.LeakyReLU(0.2),
            nn.Linear(256, 1),  # 输出真假判别
            nn.Sigmoid()  # 输出范围[0, 1]
        )
    
    def forward(self, x):
        return self.model(x)

3. 定义损失

我们需要选择损失函数和优化器：

import torch.optim as optim

# 实例化生成器和判别器
generator = Generator()
discriminator = Discriminator()

# 损失函数
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))

4. 模型训练

在此步骤中，我们将训练模型，通过对抗训练来不断改进生成器和判别器。

num_epochs = 50
for epoch in range(num_epochs):
    for i, (images, _) in enumerate(dataloader):
        # 训练判别器
        real_images = images.view(images.size(0), -1)  # 将图片展平成一维
        real_labels = torch.ones(images.size(0), 1)  # 真实标签
        fake_labels = torch.zeros(images.size(0), 1)  # 假标签

        # 判别器反馈
        outputs = discriminator(real_images)
        d_loss_real = criterion(outputs, real_labels)

        z = torch.randn(images.size(0), 100)  # 随机生成噪声
        fake_images = generator(z)  # 生成假图像
        outputs = discriminator(fake_images.detach())  # 不训练生成器，计算假图像损失
        d_loss_fake = criterion(outputs, fake_labels)

        d_loss = d_loss_real + d_loss_fake  # 判别器总损失
        optimizer_D.zero_grad()
        d_loss.backward()
        optimizer_D.step()

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

    print(f'Epoch [{epoch}/{num_epochs}], d_loss: {d_loss.item()}, g_loss: {g_loss.item()}')

5. 展示结果

可以将生成的图像进行可视化：

import matplotlib.pyplot as plt

# 生成假图像
z = torch.randn(16, 100)
fake_images = generator(z).view(-1, 1, 28, 28).detach().numpy()

# 可视化生成的图像
fig, axes = plt.subplots(4, 4, figsize=(8, 8))
for ax, img in zip(axes.ravel(), fake_images):
    ax.imshow(img[0], cmap='gray')
    ax.axis('off')  # 不显示坐标轴
plt.show()