当前位置：首页 > news >正文

Pytorch搭建基本的GAN模型及训练过程

news 来源：原创 2024/5/14 13:35:55

文章目录

概述
代码实战
- 导包
- 数据准备
- 定义生成器
- 定义判别器
- 初始化模型、优化器及损失计算函数
- 绘图函数
- GAN的训练
- 输出
整体代码
参考资料

概述

本文通过Pytorch搭建基本的GAN模型结构，并通过 torchvision 的 MNIST 数据集进行测试。
对于GAN模型的基本结构及公式的理解可以看前一篇博客：
GAN的理论知识及公式的理解
下文的实现完全对照这一篇博客的基本理论。

代码实战

代码是基于Pytorch环境创建，需要先安装Pytorch环境
Pytorch环境搭建教程链接：
Pytorch搭建教程

导包

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt
import torchvision
from torchvision import transforms

数据准备

# 对数据做归一化 （-1， 1）
transform = transforms.Compose([
    transforms.ToTensor(),         # 将数据转换成Tensor格式，channel, high, witch,数据在（0， 1）范围内
    transforms.Normalize(0.5, 0.5) # 通过均值和方差将数据归一化到（-1， 1）之间
])

# 下载数据集
train_ds = torchvision.datasets.MNIST('data',
                                      train=True,
                                      transform=transform,
                                      download=True)
                                      
# 设置dataloader
dataloader = torch.utils.data.DataLoader(train_ds, batch_size=64, shuffle=True)

# 返回一个批次的数据
imgs, _ = next(iter(dataloader))

# imgs的大小
imgs.shape

定义生成器

# 输入是长度为 100 的 噪声（正态分布随机数）
# 输出为（1， 28， 28）的图片
# linear 1 :   100----256
# linear 2:    256----512
# linear 2:    512----28*28
# reshape:     28*28----(1, 28, 28)

class Generator(nn.Module): #创建的 Generator 类继承自 nn.Module
    def __init__(self): # 定义初始化方法
        super(Generator, self).__init__() #继承父类的属性
        self.main = nn.Sequential( #使用Sequential快速创建模型
                                  nn.Linear(100, 256),
                                  nn.ReLU(),
                                  nn.Linear(256, 512),
                                  nn.ReLU(),
                                  nn.Linear(512, 28*28),
                                  nn.Tanh()                     # 输出层使用Tanh()激活函数，使输出-1, 1之间
        )
    def forward(self, x):              # 定义前向传播 x 表示长度为100 的noise输入
        img = self.main(x)
        img = img.view(-1, 28, 28) #将img展平，转化成图片的形式，channel为1可写可不写
        return img

定义判别器

## 输入为（1， 28， 28）的图片  输出为二分类的概率值，输出使用sigmoid激活 0-1
# BCEloss计算交叉熵损失

# nn.LeakyReLU   f(x) : x>0 输出 x， 如果x<0 ,输出 a*x  a表示一个很小的斜率，比如0.1
# 判别器中一般推荐使用 LeakyReLU

class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.main = nn.Sequential(
                                  nn.Linear(28*28, 512), #输入是28*28的张量，也就是图片
                                  nn.LeakyReLU(), # 小于0的时候保存一部分梯度
                                  nn.Linear(512, 256),
                                  nn.LeakyReLU(),
                                  nn.Linear(256, 1), # 二分类问题，输出到1上
                                  nn.Sigmoid()
        )
    def forward(self, x):
        x = x.view(-1, 28*28)
        x = self.main(x)
        return x

初始化模型、优化器及损失计算函数

# 定义设备
device = 'cuda' if torch.cuda.is_available() else 'cpu'
# 初始化模型
gen = Generator().to(device)
dis = Discriminator().to(device)
# 优化器
d_optim = torch.optim.Adam(dis.parameters(), lr=0.0001)
g_optim = torch.optim.Adam(gen.parameters(), lr=0.0001)
# 损失函数
loss_fn = torch.nn.BCELoss()

绘图函数

def gen_img_plot(model, epoch, test_input):
    prediction = np.squeeze(model(test_input).detach().cpu().numpy())
    fig = plt.figure(figsize=(4, 4))
    for i in range(16):
        plt.subplot(4, 4, i+1)
        plt.imshow((prediction[i] + 1)/2) # 确保prediction[i] + 1)/2输出的结果是在0-1之间
        plt.axis('off')
    plt.show()
    
test_input = torch.randn(16, 100, device=device)

GAN的训练

# 保存每个epoch所产生的loss值
D_loss = []
G_loss = []

# 训练循环
for epoch in range(20): #训练20个epoch
   d_epoch_loss = 0 # 初始损失值为0
   g_epoch_loss = 0
   # len(dataloader)返回批次数，len(dataset)返回样本数
   count = len(dataloader)
   # 对dataloader进行迭代
   for step, (img, _) in enumerate(dataloader): # enumerate加序号
       img = img.to(device) #将数据上传到设备
       size = img.size(0) # 获取每一个批次的大小
       random_noise = torch.randn(size, 100, device=device)  # 随机噪声的大小是size个
       
       d_optim.zero_grad() # 将判别器前面的梯度归0
       
       real_output = dis(img)      # 判别器输入真实的图片，real_output是对真实图片的预测结果 
       
       # 得到判别器在真实图像上的损失
       # 判别器对于真实的图片希望输出的全1的数组，将真实的输出与全1的数组进行比较
       d_real_loss = loss_fn(real_output, 
                             torch.ones_like(real_output))      
       d_real_loss.backward() # 求解梯度
       
       
       gen_img = gen(random_noise)    
       # 判别器输入生成的图片，fake_output是对生成图片的预测
       # 优化的目标是判别器，对于生成器的参数是不需要做优化的，需要进行梯度阶段，detach()会截断梯度，
       # 得到一个没有梯度的Tensor，这一点很关键
       fake_output = dis(gen_img.detach()) 
       # 得到判别器在生成图像上的损失
       d_fake_loss = loss_fn(fake_output, 
                             torch.zeros_like(fake_output))      
       d_fake_loss.backward() # 求解梯度
       
       d_loss = d_real_loss + d_fake_loss # 判别器总的损失等于两个损失之和
       d_optim.step() # 进行优化
       
       g_optim.zero_grad() # 将生成器的所有梯度归0
       fake_output = dis(gen_img) # 将生成器的图片放到判别器中，此时不做截断，因为要优化生成器
       # 生层器希望生成的图片被判定为真
       g_loss = loss_fn(fake_output, 
                        torch.ones_like(fake_output))      # 生成器的损失
       g_loss.backward() # 计算梯度
       g_optim.step() # 优化
       
       # 将损失累加到定义的数组中，这个过程不需要计算梯度
       with torch.no_grad():
           d_epoch_loss += d_loss
           g_epoch_loss += g_loss
     
   # 计算每个epoch的平均loss，仍然使用这个上下文关联器
   with torch.no_grad():
       # 计算平均的loss值
       d_epoch_loss /= count
       g_epoch_loss /= count
       # 将平均loss放入到loss数组中
       D_loss.append(d_epoch_loss.item())
       G_loss.append(g_epoch_loss.item())
       # 打印当前的epoch
       print('Epoch:', epoch)
       # 调用绘图函数
       gen_img_plot(gen, epoch, test_input)

输出

Epoch: 0
在这里插入图片描述
…(省略中间的迭代输出)

Epoch: 19
在这里插入图片描述
总共做了20次的迭代，可以看出，随着迭代次数的增加，生成的图片质量越来越好。

整体代码


import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt
import torchvision
from torchvision import transforms

# 对数据做归一化 （-1， 1）
transform = transforms.Compose([
    transforms.ToTensor(),         # 将数据转换成Tensor格式，channel, high, witch,数据在（0， 1）范围内
    transforms.Normalize(0.5, 0.5) # 通过均值和方差将数据归一化到（-1， 1）之间
])

# 下载数据集
train_ds = torchvision.datasets.MNIST('data',
                                      train=True,
                                      transform=transform,
                                      download=True)
                                      
# 设置dataloader
dataloader = torch.utils.data.DataLoader(train_ds, batch_size=64, shuffle=True)

# 返回一个批次的数据
imgs, _ = next(iter(dataloader))

# imgs的大小
imgs.shape

# 输入是长度为 100 的 噪声（正态分布随机数）
# 输出为（1， 28， 28）的图片
# linear 1 :   100----256
# linear 2:    256----512
# linear 2:    512----28*28
# reshape:     28*28----(1, 28, 28)

class Generator(nn.Module): #创建的 Generator 类继承自 nn.Module
    def __init__(self): # 定义初始化方法
        super(Generator, self).__init__() #继承父类的属性
        self.main = nn.Sequential( #使用Sequential快速创建模型
                                  nn.Linear(100, 256),
                                  nn.ReLU(),
                                  nn.Linear(256, 512),
                                  nn.ReLU(),
                                  nn.Linear(512, 28*28),
                                  nn.Tanh()                     # 输出层使用Tanh()激活函数，使输出-1, 1之间
        )
    def forward(self, x):              # 定义前向传播 x 表示长度为100 的noise输入
        img = self.main(x)
        img = img.view(-1, 28, 28) #将img展平，转化成图片的形式，channel为1可写可不写
        return img

## 输入为（1， 28， 28）的图片  输出为二分类的概率值，输出使用sigmoid激活 0-1
# BCEloss计算交叉熵损失

# nn.LeakyReLU   f(x) : x>0 输出 x， 如果x<0 ,输出 a*x  a表示一个很小的斜率，比如0.1
# 判别器中一般推荐使用 LeakyReLU

class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.main = nn.Sequential(
                                  nn.Linear(28*28, 512), #输入是28*28的张量，也就是图片
                                  nn.LeakyReLU(), # 小于0的时候保存一部分梯度
                                  nn.Linear(512, 256),
                                  nn.LeakyReLU(),
                                  nn.Linear(256, 1), # 二分类问题，输出到1上
                                  nn.Sigmoid()
        )
    def forward(self, x):
        x = x.view(-1, 28*28)
        x = self.main(x)
        return x

# 定义设备
device = 'cuda' if torch.cuda.is_available() else 'cpu'
# 初始化模型
gen = Generator().to(device)
dis = Discriminator().to(device)
# 优化器
d_optim = torch.optim.Adam(dis.parameters(), lr=0.0001)
g_optim = torch.optim.Adam(gen.parameters(), lr=0.0001)
# 损失函数
loss_fn = torch.nn.BCELoss()

def gen_img_plot(model, epoch, test_input):
    prediction = np.squeeze(model(test_input).detach().cpu().numpy())
    fig = plt.figure(figsize=(4, 4))
    for i in range(16):
        plt.subplot(4, 4, i+1)
        plt.imshow((prediction[i] + 1)/2) # 确保prediction[i] + 1)/2输出的结果是在0-1之间
        plt.axis('off')
    plt.show()
    
test_input = torch.randn(16, 100, device=device)

 # 保存每个epoch所产生的loss值
D_loss = []
G_loss = []

# 训练循环
for epoch in range(20): #训练20个epoch
    d_epoch_loss = 0 # 初始损失值为0
    g_epoch_loss = 0
    # len(dataloader)返回批次数，len(dataset)返回样本数
    count = len(dataloader)
    # 对dataloader进行迭代
    for step, (img, _) in enumerate(dataloader): # enumerate加序号
        img = img.to(device) #将数据上传到设备
        size = img.size(0) # 获取每一个批次的大小
        random_noise = torch.randn(size, 100, device=device)  # 随机噪声的大小是size个
        
        d_optim.zero_grad() # 将判别器前面的梯度归0
        
        real_output = dis(img)      # 判别器输入真实的图片，real_output是对真实图片的预测结果 
        
        # 得到判别器在真实图像上的损失
        # 判别器对于真实的图片希望输出的全1的数组，将真实的输出与全1的数组进行比较
        d_real_loss = loss_fn(real_output, 
                              torch.ones_like(real_output))      
        d_real_loss.backward() # 求解梯度
        
        
        gen_img = gen(random_noise)    
        # 判别器输入生成的图片，fake_output是对生成图片的预测
        # 优化的目标是判别器，对于生成器的参数是不需要做优化的，需要进行梯度阶段，detach()会截断梯度，
        # 得到一个没有梯度的Tensor，这一点很关键
        fake_output = dis(gen_img.detach()) 
        # 得到判别器在生成图像上的损失
        d_fake_loss = loss_fn(fake_output, 
                              torch.zeros_like(fake_output))      
        d_fake_loss.backward() # 求解梯度
        
        d_loss = d_real_loss + d_fake_loss # 判别器总的损失等于两个损失之和
        d_optim.step() # 进行优化
        
        g_optim.zero_grad() # 将生成器的所有梯度归0
        fake_output = dis(gen_img) # 将生成器的图片放到判别器中，此时不做截断，因为要优化生成器
        # 生层器希望生成的图片被判定为真
        g_loss = loss_fn(fake_output, 
                         torch.ones_like(fake_output))      # 生成器的损失
        g_loss.backward() # 计算梯度
        g_optim.step() # 优化
        
        # 将损失累加到定义的数组中，这个过程不需要计算梯度
        with torch.no_grad():
            d_epoch_loss += d_loss
            g_epoch_loss += g_loss
      
    # 计算每个epoch的平均loss，仍然使用这个上下文关联器
    with torch.no_grad():
        # 计算平均的loss值
        d_epoch_loss /= count
        g_epoch_loss /= count
        # 将平均loss放入到loss数组中
        D_loss.append(d_epoch_loss.item())
        G_loss.append(g_epoch_loss.item())
        # 打印当前的epoch
        print('Epoch:', epoch)
        # 调用绘图函数
        gen_img_plot(gen, epoch, test_input)