当前位置：首页>編程日記>正文

GAN（生成对抗神经网络）生成MNIST 基于pytorch实现

編程日記08-29

GAN（生成对抗神经网络）生成MNIST 基于pytorch实现

运行环境

pytorch1.3.0

简介

生成对抗神经网络分为两部分：生成器

鉴别器

生成器

把一个满足高斯分布的向量映射为一个784维度的向量，期望这个784维度的向量就是我们要的图片。对于生成器的监督训练主要是靠鉴别器。

鉴别器

是一个二分类模型，输入是图片向量，输出：这张图是真图的概率。

训练数据有两种：真图、假图（生成器生成的）

真图对应的label是1、假图对应的label是0

补充

从0.4起， Variable 正式合并入Tensor类，通过Variable嵌套实现的自动微分功能已经整合进入了Tensor类中。虽然为了代码的兼容性还是可以使用Variable(tensor)这种方式进行嵌套，但是这个操作其实什么都没做。

运行方法：把程序粘贴下来保存为 xx.py ，然后最好把xx.py放在一个空的文件夹里面，python装好需要的包，就可以运行了

运行后程序会自动再生成两个文件夹，总体长这个样子

本程序数据集会自动下载，如果慢可以使用迅雷下载下来然后放到指定的地方（data\MNIST\raw）。

代码

import torch
import torch.nn as nn
import torchvision
import torch.nn.functional as F
from torchvision import datasets
from torchvision import transforms
from torchvision.utils import save_image
from torch.autograd import Variable
import os#超参数
batch_size = 128
num_epoch = 500
z_dimention = 100
d_optimizer_lr=0.0001
g_optimizer_lr=0.0001
use_gpu=Trueclass discriminator(nn.Module):  #784长度向量 -> (0,1)def __init__(self):super(discriminator, self).__init__()self.dis = nn.Sequential(nn.Linear(784, 256),nn.LeakyReLU(0.2),nn.Linear(256, 256),nn.LeakyReLU(0.2),nn.Linear(256, 1),nn.Sigmoid())def forward(self, x):x = self.dis(x)return xclass generator(nn.Module):  # 100长度向量 -> 784长度向量def __init__(self):super(generator, self).__init__()self.gen = nn.Sequential(nn.Linear(100, 256),nn.ReLU(True),nn.Linear(256, 256),nn.ReLU(True),nn.Linear(256, 784),nn.Tanh())def forward(self, x):x = self.gen(x)return x#784向量 -> 1*28*28
def to_img(x):out = 0.5 * (x + 1)out = out.clamp(0, 1)out = out.view(-1, 1, 28, 28)return outimg_transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5,), (0.5,))
])mnist = datasets.MNIST(root='./data/', train=True, transform=img_transform, download=True
)dataloader = torch.utils.data.DataLoader(dataset=mnist, batch_size=batch_size,shuffle=True,drop_last=True
)if __name__=='__main__':if not os.path.exists('./img'):os.mkdir('./img')D = discriminator()G = generator()criterion = nn.BCELoss()  # 二分类的交叉熵d_optimizer = torch.optim.Adam(D.parameters(), lr=d_optimizer_lr)g_optimizer = torch.optim.Adam(G.parameters(), lr=g_optimizer_lr)z=torch.FloatTensor(batch_size,z_dimention)if use_gpu:D = D.cuda()G = G.cuda()z = z.cuda()for epoch in range(num_epoch):for i, (img, _) in enumerate(dataloader):img = img.view(batch_size, -1)if use_gpu:real_img = img.cuda()real_label = torch.ones(batch_size).cuda()fake_label = torch.zeros(batch_size).cuda()# =================train discriminatorreal_out = D(real_img)d_loss_real = criterion(real_out, real_label)  # 真实数据对应输出 1real_scores = real_outz.data.normal_(0,1)fake_img1 = G(z)fake_out = D(fake_img1.detach())d_loss_fake = criterion(fake_out, fake_label)fake_scores = fake_outd_loss = d_loss_real + d_loss_faked_optimizer.zero_grad()d_loss.backward()d_optimizer.step()# ===============train generatorz.data.normal_(0,1)fake_img2 = G(z)output = D(fake_img2)g_loss = criterion(output, real_label)g_optimizer.zero_grad()g_loss.backward()g_optimizer.step()if epoch == 0:real_images = to_img(real_img.cpu().data)save_image(real_images, './img/real_images.png')print('Epoch[{}/{}],d_loss:{:.6f},g_loss:{:.6f} D real:{:.6f},D fake:{:.6f}'.format(epoch, num_epoch, d_loss.item(), g_loss.item(), real_scores.data.mean(), fake_scores.data.mean()))fake_images = to_img(fake_img2.cpu().data)save_image(fake_images, './img/fake_images-{}.png'.format(epoch + 1))