Pytorch搭建基本的GAN模型及训练过程
Posted 半岛铁子_
tags:
篇首语：本文由小常识网(cha138.com)小编为大家整理，主要介绍了Pytorch搭建基本的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)
参考资料
[1] https://www.bilibili.com/video/BV1xm4y1X7KZ
[2] https://blog.csdn.net/hshudoudou/article/details/126922562?spm=1001.2014.3001.5502
以上是关于Pytorch搭建基本的GAN模型及训练过程的主要内容，如果未能解决你的问题，请参考以下文章
[Pytorch系列-72]：生成对抗网络GAN - 图像生成开源项目pytorch-CycleGAN-and-pix2pix - 使用预训练模型训练CycleGAN模型
将 torch.backward() 用于 GAN 生成器时，为啥 Pytorch 中的判别器损失没有变化？
[Pytorch系列-71]：生成对抗网络GAN - 图像生成开源项目pytorch-CycleGAN-and-pix2pix - 使用预训练模型训练pix2pix模型
[Pytorch系列-68]：生成对抗网络GAN - 图像生成开源项目pytorch-CycleGAN-and-pix2pix - 使用预训练模型测试CycleGAN模型