Pytorch基础知识（12）GANs生成对抗网络
在PyTorch的神经风格迁移一章中，我们学习了一种通过模仿艺术图像的风格来生成新数据的方法。在本章中，我们将介绍另一种生成新数据的方法，称为生成对抗网络(GANs)。GAN是一个通过学习数据分布来生成新数据的框架。
GAN框架由generator和discriminator两个神经网络组成，如下图所示:
在图像生成方面，当给定噪声作为输入时，生成器生成假数据，判别器将真实图像与假图像进行分类。在训练过程中，生成器和判别器会相互竞争，结果是他们的工作做得更好。生成器试图生成更好的图像来欺骗判别器，而判别器试图更好识别真假图像。
GAN仍在不断发展，每天都会出现新的应用程序。其中一些应用包括艺术图像生成、数据增强、图像到图像转换、超分辨率和视频合成。
在本章中，我们将使用PyTorch开发一个GAN来生成类似STL-10数据集的新图像。我们将遵循以下论文中提出的深度卷积GAN (DCGAN)架构。
本章将涵盖以下教程:
- 创建数据集
- 定义生成器与判别器
- 定义损失函数和优化器
- 训练模型
- 部署生成器
创建数据集
为了训练GAN，我们需要一个训练数据集。给定一个训练数据集，GAN将学习生成与训练数据集具有相同分布的新数据。例如，如果我们用猫的图像训练一个GAN，它将学会生成在我们眼中看起来真实的新猫图像。我们将使用torchvision包中的STL-10数据集。我们在多类图像分类中使用了这个数据集来完成多标签分类任务。
在本教程中，您将学习如何定义PyTorch数据集和数据加载器来训练GAN。
我们将创建内置数据集类STL-10的对象，并定义一个数据加载器如下:
#1. 导入必要的库
from torchvision import datasets
import torchvision.transforms as transforms
import os
path2data="./data"
os.makedirs(path2data, exist_ok=True)
#2. 定义数据变换
h,w=64,64
mean=(0.5,0.5,0.5)
std=(0.5,0.5,0.5)
transform=transforms.Compose([
transforms.Resize((h,w)),
transforms.CenterCrop((h,w)),
transforms.ToTensor(),
transforms.ToTensor(),
transforms.Normalize(mean,std),
])
#3. 实例化STL-10类对象
```python
train_ds=datasets.STL10(path2data, split="train", download=True, transform=transform)
print(len(train_ds))
# 5000
#4. 从数据集中获取一个样本
import torch
for x,_ in train_ds:
print(x.shape, torch.min(x), torch.max(x))
break
# torch.Size([3,64,64]) tensor(-0.8980) tensor(0.9529)
#5. 显示样本数据
from torchvision.transforms.functional import to_pil_image
import matplotlib.pylab as plt
plt.imshow(to_pil_image(0.5*x+0.5))
#6. 构建数据加载器
import torch
batch_size=32
train_dl=torch.utils.data.DataLoader(train_ds, batch_size=batch_size,shuffle=True)
#7.获得数据加载器中的批数据
for x,y in train_dl:
print(x.shape, y.shape)
break
# torch.Size([32, 3, 64, 64]) torch.Size([32])
代码解析：
在第1步中，我们导入了基本的包，并定义并创建了一个文件夹来存储下载时的数据。
在步骤2中，我们使用torchvision.transforms定义了数据变换。原始图像可能有不同的大小，因此我们使用Resize变换将图像的大小调整为64 * 64。接下来，ToTensor将图像像素缩放到[0,1]的范围。接下来，我们应用了标准化。设置标准化均值和标准值将输入标准化到[- 1,1]的范围。正如您将在定义Generator和Discriminator教程中发现的那样，Generator模型的输出是tanh函数，它生成范围为[- 1,1]的输出。
sigmoid激活函数的输出范围为[0,1]，tanh激活函数的输出范围为[-1，+1]。
在步骤3中，我们从torchvision.datasets包中实例化了STL-10类的一个对象。我们将数据地址和变换函数传递给类。该数据集有5000个数据样本。
在步骤4中，我们从数据集中获得一个样本图像，并打印其形状和最小值和最大值。正如预期的那样，提取的样本是一个PyTorch张量，形状为(3,height, width)，并归一化到[- 1,1]的范围。
在第5步中，我们显示了示例图像。注意，因为张量被归一化为[- 1,1]，我们必须为可视化目的反标准化它。
在第6步中，我们定义了一个数据加载器。批量大小设置为32。但是，您可以根据您的计算机和GPU内存进行调整。如果在训练模型时遇到内存错误，请尝试减少批处理大小。
在第7步中，我们从数据加载器中提取了一个小批量样本，并打印了它的形状。
定义生成器和判别器
GAN框架是基于两个模型的竞争，即生成器和判别器。生成器生成虚假图像，而判别器识别真假图像。这种竞争的结果是，生成器将产生更好的假图像，而判别器将变得更好地识别它们。
正如本章开头提到的，我们将使用DCGAN框架。该判别器基于DCGAN，其结构类似于基于卷积层的二值分类模型。我们之前已经开发了一个二分类模型，但是在这里，池化层被卷积层取代。
此外，该生成器的架构基于转置卷积层，将输入噪声向量上采样到期望的输出大小，如下图所示:
在本教程中，您将学习如何实现GAN框架的生成器和判别器模型。
# 我们将定义生成器和判别器模型并且初始化权重
#1. 定义Generator类：
from torch import nn
import torch.nn.functional as F
class Generator(nn.Module):
def __init__(self,params):
super(Generator, self).__init__()
nz=params["nz"]
ngf=params["ngf"]
noc=params["noc"]
self.dconv1=nn.ConvTransposed2d(nz, ngf*8, kernel_size=4, stride=1, padding=0, bias=False)
self.bn1=nn.BatchNorm2d(ngf*8)
self.dconv2=nn.ConvTransposed2d(ngf*8, ngf*4,kernel_size=4,stride=2,padding=1,bias=False)
self.bn2=nn.BatchNorm2d(ngf*4)
self.dconv3=nn.ConvTransposed2d(ngf*4, ngf*2,kernel_size=4,stride=2,padding=1,bias=False)
self.bn3=nn.BatchNorm2d(ngf*2)
self.dconv4=nn.ConvTransposed2d(ngf*2, ngf*1,kernel_size=4,stride=2,padding=1,bias=False)
self.bn4=nn.BatchNorm2d(ngf*1)
self.dconv5=nn.ConvTransposed2d(ngf, noc,kernel_size=4,stride=2,padding=1,bias=False)
def forward(self,x):
x=F.relu(self.bn1(self.dconv1(x))
x=F.relu(self.bn2(self.dconv2(x))
x=F.relu(self.bn3(self.dconv3(x))
x=F.relu(self.bn4(self.dconv4(x))
out=torch.tanh(self.dconv5(x))
return out
#2. 定义Generator类的一个对象
params_gen={
"nz":100,
"ngf":64,
"noc":3
}
model_gen = Generator(params_gen)
device=torch.device("cuda")
model_gen.to(device)
print(model_gen)
# Generator(
# (dconv1): ConvTranspose2d(100, 512, kernel_size=(4, 4),stride=(1, 1), bias=False)
# (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True,track_running_stats=True)
# (dconv2): ConvTranspose2d(512, 256, kernel_size=(4, 4),stride=(2, 2), padding=(1, 1), bias=False)
# (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True,track_running_stats=True)
# ...
# 让我们向模型传递一些虚拟输入
with torch.no_grad():
y=model_gen(torch.zeros(1,100,1,1,device=device)
print(y.shape)
# torch.Size([1, 3, 64, 64])
#3. 定义Discriminator类
class Discriminator(nn.Module):
def __init__(self,params):
super(Discriminator, self).__init__()
nic=params["nic"]
ndf=params["ndf"]
self.conv1=nn.Conv2d(nic,ndf,kernel_size=4,stride=2,padding=1,bias=False)
self.conv2=nn.Conv2d(ndf, ndf*2, kernel_size=4, stride=2, padding=1, bias=False)
self.bn2=nn.BatchNorm2d(ndf*2)
self.conv3=nn.Conv2d(ndf*2, ndf*4, kernel_size=4, stride=2, padding=1, bias=False)
self.bn3=nn.BatchNorm2d(ndf*4)
self.conv4=nn.Conv2d(ndf*4, ndf*8, kernel_size=4, stride=2, padding=1, bias=False)
self.bn4=nn.BatchNorm2d(ndf*8)
self.conv5=nn.Conv2d(ndf*8, 1, kernel_size=4, stride=1, padding=0, bias=False)
def forward(self,x):
x= F.leaky_relu(self.conv1(x), 0.2, True)
x= F.leaky_relu(self.bn2(self.conv2(x)), 0.2, inplace=True)
x= F.leaky_relu(self.bn3(self.conv3(x)), 0.2, inplace=True)
x= F.leaky_relu(self.bn4(self.conv4(x)), 0.2, inplace=True)
out=torch.sigmoid(self.conv5(x))
return out.view(-1)
#4. 定义Discriminator类的对象
params_dis={
"nic":3,
"ndf":64,
}
model_dis=Discriminator(params_dis)
model_dis.to(device)
print(model_dis)
# Discriminator(
# (conv1): Conv2d(3, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
# (conv2): Conv2d(64, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
# (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
# ...
with torch.no_grad():
y=model_dis(torch.zeros(1,3,h,w,device))
print(y.shape)
# torch.Size([1])
#5. 定义初始化模型权重的辅助函数
def initialize_weights(model):
classname=model.__class__.__name__
if classname.find('Conv') != -1:
nn.init.normal_(model.weight.data, 0.0, 0.02)
elif classname.find('BatchNorm') != -1:
nn.init.normal_(model.weight.data, 1.0, 0.02)
nn.init.constant_(model.bias.data, 0)
#6. 通过调用辅助函数初始化模型权重
model_gen.apply(initialize_weights)
model_dis.apply(initialize_weights)
代码解析：
在步骤1中，我们定义了带有两个方法的Generator类。在__init__方法中，我们定义了层。该函数有一个输入参数params，它是一个Python字典，包含以下键:
- nz:输入噪音向量的尺寸（设为100）
- ngf:生成器中的卷积滤波器数目（设置为64）
- noc:输出通道数目（对于RGB图像来说是3)
如所见，定义了5个conv-transpose（转置卷积层）。一个conv-transpose（转置卷积层）也被称为一个反卷积。它们用于将输入向量上采样到所需的输出大小。
在forward方法中，我们定义了层与层之间的连接并得到了输出。Generator的输出是一个形状为 (batch_size, 3, height, width)的张量。
在步骤2中，我们定义了一个名为model_gen的Generator类的对象。为了确保正确地创建了模型，我们向生成器模型传递了一些虚拟输入。与预期的一样，模型输出是一个形状张量[1,3,64,64]。
在步骤3中，我们定义了Discriminator类。类似地，在__init__方法中，我们定义了层，而在forward方法中，我们定义了层与层之间的连接。请注意，我们没有使用任何池化层，而是将stride参数设置为2或4来对输入大小进行下采样。另外，注意使用leaky_relu激活代替relu来减少过拟合。
在步骤4中，我们定义了Discriminator类的一个对象。为了确保正确地创建模型，我们向判别器模型传递了一些虚拟输入。这个简单的测试将帮助修复任何可能的错误，然后我们继续进行下一步。
在步骤5中，我们定义了一个辅助函数来初始化模型权重。该函数的输入是一个PyTorch模型。DCGAN的论文建议使用mean=0和std=0.02的正态分布初始化权重，就像我们在辅助函数中所做的那样。
在步骤6中，我们将initialize_weights辅助函数应用于生成器和判别器模型，以初始化它们的权重。
定义损失函数和优化器
为了让模型学习，我们需要定义一个标准。判别器模型是一种分类网络，我们可以使用二元交叉熵(BCE)损失函数。对于生成器模型，我们将其输出传递给判别器模型，然后对判别器模型的输出进行评估。因此，可以使用相同的BCE损失函数作为训练生成器模型的准则。此外，我们将使用Adam优化器更新判别器和生成器模型的参数。
在本教程中，您将学习如何定义GAN网络的损失函数和优化器。
#1. 定义BCE损失函数
loss_func=nn.BCELoss()
#2. 为生成器定义优化器
lr=2e-4
beta1=0.5
opt_dis=optim.Adam(model_dis.parameters(),lr=lr,betas=(beta1,0.999))
#3. 定义判别器的优化器
opt_gen=optim.Adam(model_gen.parameters(),lr=lr,betas=(beta1,0.999)
代码解析：
在步骤1中，我们使用torch.nn包中的BCE损失函数。正如您将在下一节中看到的，我们将在多个步骤中使用此损失函数。
在步骤2中，我们为生成器使用了torch.optim包中的Adam优化器，优化器的超参数使用了DCGAN论文中推荐的超参数。论文建议将学习速率设置为0.0002，动量项设置为beta1，以获得训练稳定性。
类似地，在步骤3中，我们为判别器使用了torch.optim包中的Adam优化器。
模型训练
对GAN的训练分为两个阶段:对判别器的训练和对生成器的训练。为此，我们将采取以下措施:
- 获得一批真实图像，标签设置为1。
- 使用生成器生成一批假的图像，标签设置为0。
- 将小批数据输入判别器，计算损失和梯度。
- 使用梯度更新判别器参数。
- 使用生成器生成一批假的图像，标签设置为1。
- 将假的小批数据输入判别器，并计算损失和梯度。
- 仅基于梯度更新生成器。
- 重复以上步骤
在本教程中，您将学习如何实现这些步骤。
# 我们将实现GAN网络的训练步骤
#1. 定义一些参数
real_label=1
fake_label=0
nz=params["nz"]
num_epochs=100
loss_history={"gen":[],
"dis":[]}
# 开始训练并计算真实样本的损失
batch_count=0
for epoch in range(num_epochs):
for xb, yb in train_dl:
ba_si=xb.size(0) # 一批共有多少数据，即batchsize
model_dis.zero_grad() # 判别器梯度清零，如果不清零，梯度会累加
xb=xb.to(device) # 送到GPU设备上
yb=torch.full((ba_si,), real_label,device=device) # 标签置为1
out_dis=model_dis(xb) # 判别器输出结果
loss_r=loss_func(out_dis,yb.float()) # 计算损失
loss_r.backward() # 计算梯度
noise=torch.randn(ba_si,nz,1,1,device=device) # 随机噪声
out_gen=model_gen(noise) # 生成器生成假数据
out_dis=model_dis(out_gen.detach()) # 将假数据送到判别器中
yb.fill_(fake_label) # 设置噪声数据标签为0
loss_f=loss_func(out_dis, yb.float()) # 计算损失
loss_f.backward() # 求梯度
loss_dis=loss_r + loss_f # 总损失
opt_dist.step() # 判别器参数更新， 包含了上面两次计算梯度的结果
# 训练生成器
model_gen.zero_grad() # 将生成器的梯度清零，如果不清零，梯度会累加
yb.fill_(real_label) # 将假数据的标签置为1
out_dis=model_dist(out_gen) # 将生成的假数据送入已经更新参数的判别器中
loss_gen=loss_func(out_dis,yb.float()) # 计算损失
loss_gen.backward() # 求梯度
opt_gen.step() # 生成器参数更新
loss_history["gen"].append(loss_gen.item())
loss_history["dis"].append(loss_dis.item())
batch_count+=1
if batch_count%100==0:
print(epoch,loss_gen.item(),loss_dis.item())
# 0 7.026479721069336 0.12888161838054657
# 1 3.8994224071502686 0.24403591454029083
# 1 12.108219146728516 1.221606731414795
# ...
#3. 绘制损失
plt.figure(figsize=(10,5))
plt.title("Loss Progress")
plt.plot(loss_history["gen"],label="Gen. Loss")
plt.plot(loss_history["dis"],label="Dis. Loss")
plt.xlabel("batch count")
plt.ylabel("Loss")
plt.legend()
plt.show()
#4. 保存权重
import os
path2models="./models/"
os.makedirs(path2models, exist_ok=True)
path2weights_gen=os.path.join(path2models, "weights_gen.pt")
path2weights_dis=os.path.join(path2models,"weights_dis.pt")
torch.save(model.state_dict(),path2weights_gen)
torch.save(model.state_dict(),path2weights_dis)
代码解析
在步骤1中，我们定义了一些参数。我们定义了real_label和fake_label，并分别将它们设置为1和0。稍后，我们将需要使用这些参数标记一个小批处理。nz参数指定生成器模型的输入噪声向量的大小。这在定义生成器和判别器教程中被设置为100。num_epochs参数指定我们希望遍历训练数据的次数。为了存储判别器和生成器模型的损失值，我们定义了loss_history字典。
在步骤2中，我们实现了训练循环。训练循环遍历真实的数据集num_epochs次。在每个epoch中，我们从train_dl中获得一批真实图像，并将其输入判别器模型，得到其输出为out_dis。注意，在这里，使用torch.full方法设置真实图像的标签为real_label。然后，计算小批真实图像的损失值为loss_r。接下来，计算loss_r对于判别器模型参数的梯度。
接下来，我们使用生成器生成了一小批假图像，并将它们提供给判别器。在将生成器的输出传递给判别器时，我们使用了.detach()方法来避免生成器模型的梯度跟踪。注意，此时，使用torch.fill_方法将假图像标记为fake_label。然后，计算小批假图像的损失值为loss_f。接下来，计算loss_f对于判别器模型参数的梯度。
最后，我们使用opt_dis.step()方法更新判别器参数。
接下来，我们对生成器模型进行训练。为此，我们将假图像传递给判别器模型并得到其输出。请注意，在这里，使用.fill_方法将假图像标记为real_label标签。乍一听这可能很奇怪，但这样做是为了迫使生成器模型生成更好看的图像。
即使生成器模型的输出是假的图像，我们在计算损失值时使用real_label作为目标值。
然后，我们计算损失值为loss_gen，计算其梯度，并使用opt_get.step()更新生成器参数。通过执行这段代码，损失值就会显示在屏幕上。
在步骤3中，我们绘制了训练过程中生成器和判别器的损失值。
在步骤4中，我们将训练后的权重存储到pickle文件中，以备将来使用。
近年来，GAN在提高生成数据的质量和维度方面取得了重大进展。作为例子，您可以参考以下论文:
StyleGAN, code
部署生成器模型
一旦我们训练了一个GAN，我们就得到了两个经过训练的模型。通常，我们会舍弃判别器模型而保留生成器模型。我们可以使用经过训练的生成器来生成新的图像。为了部署生成器模型，我们将训练好的权值加载到模型中，然后给它输入随机噪声。确保预先定义了模型类。为了避免重复，我们将不在这里定义模型类。在本教程中，您将学习如何部署生成器模型。
#1. 加载权重文件
weigths = torch.load(path2weights_gen)
model_gen.load_state_dict(weights)
#2. 设置模型为eval()模式
model_gen.eval()
#3. 将指定类型的噪声传入生成器，得到输出
with torch.no_grad():
fixed_noise=torch.randn(16, nz, 1, 1, device=device)
img_fake=model_gen(fixed_noise).detach().cpu()
print(img_fake.shape)
# torch.Size([16,3,64,64])
#4. 显示生成的图像
plt.figure(figsize=(10,10))
for ii in range(16):
plt.subplot(4,4,ii+1)
plt.imshow(to_pil_image(0.5*img_fake[ii]+0.5))
plt.axis("off")
plt.show()
代码解析：
在步骤1中，我们将训练好的权重加载到生成器模型中。在步骤2中，我们将模型设置为评估模式。在步骤3中，我们向模型中输入随机噪声向量，并接收生成的假图像。在步骤4中，我们显示了生成的伪造图像。注意，为了可视化的目的，我们必须将输出张量反标准化回其原始值。
检查生成的图像。其中一些可能看起来非常扭曲，而另一些看起来相对真实。为了改进结果，可以针对单个数据类训练模型，而不是同时训练多个类。GANs在接受单一类训练时表现更好。STL-10数据集有多个类。尝试选择一个类别训练GAN模型。此外，您可以尝试长时间地训练模型，看看它如何改变生成的图像。