Pytorch GAN实战 MINIST手写数字识别分布解析

前言、准备工作

本案例需要导入的包, 没有下载的通过pip install来下载

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

GAN网络模型结果

一、数据集导入

MINIST数据集就不过多介绍了, 这里主要是做了个Normalize的操作,将0~1之间的数据转换为-1~1之间的

这里的transform中的操作是为了将输入的图像转为-1~1之间的数值, 是因为在GAN中生成器的最后一层往往会使用nn.Tanh()效果比较好, 而nn.Tanh()返回的数据就是-1~1之间, 为了使得生成器生成的图像数据能和真是的图像数据都在同一个范围内,便于比较,因此在加载数据集的transform操作中将其归一化到-1~1之间

transform = transforms.Compose([transforms.ToTensor(),  # 0~1, [channel,h,w]transforms.Normalize(0.5, 0.5)  # 0~1 => -1~1
])
train_ds = torchvision.datasets.MNIST(r'D:\Source\Datasets',train=True,transform=transform,download=False)
loader = DataLoader(train_ds, batch_size=64, shuffle=True)
x, _ = iter(loader).__next__()
print(x.shape)

这里可以随便打印一下MINIST中的数字

二、生成器构建

本案例中生成器使用长度为 100的噪声(正态分布随机数)作为输入
(1, 28, 28)的图片作为输出

其中各个层的维度变换作用如下所示, 并在期间交替使用激活函数, 这里要注意的是最后一层的激活函数要使用tanh()

linear 1: 100->256
linear 2: 256->512
linear 3: 512->28*28
reshape: 28-28->(1, 28, 28)

class Generator(nn.Module):def __init__(self):super(Generator, self).__init__()self.main = nn.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的噪声输入img = self.main(x)img = img.view(-1, 1, 28, 28)return img

‘’‘判别器’’’

三、判别器构建

判别器输入为图片(1, 28, 28)

输出为二分类的概率值, 使用sigmoid激活0~1

用BCEloss计算交叉熵损失

nn.LeakyReLU: x<0时返回α*x x>0时返回x

在判别器中推荐用LeakyReLU的原因是, 让负值产生梯度, 利于更新

这里判别器其实与之前的图像识别网络类似, 不过多解释了

class Discriminator(nn.Module):def __init__(self):super(Discriminator, self).__init__()self.main = nn.Sequential(nn.Linear(28 * 28, 256),nn.LeakyReLU(0.2),nn.Linear(256, 256),nn.LeakyReLU(0.2),nn.Linear(256, 1),nn.Sigmoid(),)def forward(self, x):# x为28*28的图片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 = optim.Adam(dis.parameters(), lr=1e-4)
g_optim = optim.Adam(gen.parameters(), lr=1e-4)
loss_fn = nn.BCELoss()'''绘图函数'''
def gen_img_plot(model, test_input):predict = np.squeeze(model(test_input).detach().cpu().numpy())plt.figure(figsize=(4, 4))for i in range(16):plt.subplot(4, 4, i + 1)plt.imshow((predict[i] + 1) / 2)plt.axis('off')plt.show()
# 随机生成16batch,100长的噪声作为Generator的输入
test_input = torch.randn(16, 100).to(device)

五、网络训练

GAN的网络训练还是比较特殊的, 我们可以将其分为两个阶段

训练判别器

        d_optim.zero_grad()# 在判别器训练真实图片real_output = dis(img)  # 判别器输入真实的图片, 希望real_output是1d_real_loss = loss_fn(real_output, torch.ones_like(real_output))  # 得到判别器在真实数据上的损失d_real_loss.backward()# 在判别器训练生成的假图片gen_img = gen(random_noise)# 此时优化的对象是判别器, 要把生成器的梯度截断fake_output = dis(gen_img.detach())  # 判别器输入生成的图片 fake_output对生成的图片预测d_fake_loss = loss_fn(fake_output, torch.zeros_like(fake_output))  # 得到判别器在生成数据上的损失d_fake_loss.backward()d_loss = d_real_loss + d_fake_lossd_optim.step()

这边有两个步骤

判别器训练真实图片
- 对于输入数据为真实图片, 判别器需要识别的结尾是1, 因此我们的loss是与torch.ones_like相比较
判别器训练生成的假图片
- 这里需要注意的是在训练生成器生成的图片过程中我们要截断生成器的反向传播过程, 使用detach来达到目的
- 对于输入的数据为Generator生成的图片, 判别器要做的是尽量将它判断为0, 也就是与torch.zeros_like相比较

最终两个loss加起来便是判别器总的loss

训练生成器

        g_optim.zero_grad()fake_output = dis(gen_img)  # 希望其骗过disg_loss = loss_fn(fake_output, torch.ones_like(fake_output))  # 得到生成器的损失g_loss.backward()g_optim.step()

生成器的目的就是骗过判别器, 也就是让自己生成的图片在输入判别器之后尽可能的接近1, 因此与torch.ones_like相比较

总体训练网络代码如下

'''GAN的训练'''
D_loss = []
G_loss = []
# 训练循环
for epoch in range(50):d_epoch_loss = 0g_epoch_loss = 0count = len(loader)for step, (img, _) in enumerate(loader):img = img.to(device)size = img.size(0)random_noise = torch.randn(size, 100).to(device)d_optim.zero_grad()# 在判别器训练真实图片real_output = dis(img)  # 判别器输入真实的图片, 希望real_output是1d_real_loss = loss_fn(real_output, torch.ones_like(real_output))  # 得到判别器在真实数据上的损失d_real_loss.backward()# 在判别器训练生成的假图片gen_img = gen(random_noise)# 此时优化的对象是判别器, 要把生成器的梯度截断fake_output = dis(gen_img.detach())  # 判别器输入生成的图片 fake_output对生成的图片预测d_fake_loss = loss_fn(fake_output, torch.zeros_like(fake_output))  # 得到判别器在生成数据上的损失d_fake_loss.backward()d_loss = d_real_loss + d_fake_lossd_optim.step()# 对生成器进行优化g_optim.zero_grad()fake_output = dis(gen_img)  # 希望其骗过disg_loss = loss_fn(fake_output, torch.ones_like(fake_output))  # 得到生成器的损失g_loss.backward()g_optim.step()with torch.no_grad():d_epoch_loss += d_lossg_epoch_loss += g_losswith torch.no_grad():d_epoch_loss /= countg_epoch_loss /= countD_loss.append(d_epoch_loss)G_loss.append(g_epoch_loss)print(f'Epoch: {epoch}, d_epoch_loss: {d_epoch_loss}, g_epoch_loss:{g_epoch_loss}')gen_img_plot(gen, test_input)

六、训练结果

可以看到随着epoch的增加, 生成的效果越来越好