Pytorch使用MNIST數據集實現CGAN和生成指定的數字方式

更新時間：2020年01月10日 09:58:04 作者：shiheyingzhe

今天小編就為大家分享一篇Pytorch使用MNIST數據集實現CGAN和生成指定的數字方式，具有很好的參考價值，希望對大家有所幫助。一起跟隨小編過來看看吧

CGAN的全拼是Conditional Generative Adversarial Networks，條件生成對抗網絡，在初始GAN的基礎上增加了圖片的相應信息。

這里用傳統的卷積方式實現CGAN。

import torch
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision import transforms
from torch import optim
import torch.nn as nn
import matplotlib.pyplot as plt
import numpy as np
from torch.autograd import Variable
import pickle
import copy
 
import matplotlib.gridspec as gridspec
import os
 
def save_model(model, filename): #保存為CPU中可以打開的模型
 state = model.state_dict()
 x=state.copy()
 for key in x: 
  x[key] = x[key].clone().cpu()
 torch.save(x, filename)
 
def showimg(images,count):
 images=images.to('cpu')
 images=images.detach().numpy()
 images=images[[6, 12, 18, 24, 30, 36, 42, 48, 54, 60, 66, 72, 78, 84, 90, 96]]
 images=255*(0.5*images+0.5)
 images = images.astype(np.uint8)
 grid_length=int(np.ceil(np.sqrt(images.shape[0])))
 plt.figure(figsize=(4,4))
 width = images.shape[2]
 gs = gridspec.GridSpec(grid_length,grid_length,wspace=0,hspace=0)
 for i, img in enumerate(images):
  ax = plt.subplot(gs[i])
  ax.set_xticklabels([])
  ax.set_yticklabels([])
  ax.set_aspect('equal')
  plt.imshow(img.reshape(width,width),cmap = plt.cm.gray)
  plt.axis('off')
  plt.tight_layout()
#  plt.tight_layout()
 plt.savefig(r'./CGAN/images/%d.png'% count, bbox_inches='tight')
 
def loadMNIST(batch_size): #MNIST圖片的大小是28*28
 trans_img=transforms.Compose([transforms.ToTensor()])
 trainset=MNIST('./data',train=True,transform=trans_img,download=True)
 testset=MNIST('./data',train=False,transform=trans_img,download=True)
 # device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 trainloader=DataLoader(trainset,batch_size=batch_size,shuffle=True,num_workers=10)
 testloader = DataLoader(testset, batch_size=batch_size, shuffle=False, num_workers=10)
 return trainset,testset,trainloader,testloader
 
class discriminator(nn.Module):
 def __init__(self):
  super(discriminator,self).__init__()
  self.dis=nn.Sequential(
   nn.Conv2d(1,32,5,stride=1,padding=2),
   nn.LeakyReLU(0.2,True),
   nn.MaxPool2d((2,2)),
 
   nn.Conv2d(32,64,5,stride=1,padding=2),
   nn.LeakyReLU(0.2,True),
   nn.MaxPool2d((2,2))
  )
  self.fc=nn.Sequential(
   nn.Linear(7 * 7 * 64, 1024),
   nn.LeakyReLU(0.2, True),
   nn.Linear(1024, 10),
   nn.Sigmoid()
  )
 def forward(self, x):
  x=self.dis(x)
  x=x.view(x.size(0),-1)
  x=self.fc(x)
  return x
 
class generator(nn.Module):
 def __init__(self,input_size,num_feature):
  super(generator,self).__init__()
  self.fc=nn.Linear(input_size,num_feature) #1*56*56
  self.br=nn.Sequential(
   nn.BatchNorm2d(1),
   nn.ReLU(True)
  )
  self.gen=nn.Sequential(
   nn.Conv2d(1,50,3,stride=1,padding=1),
   nn.BatchNorm2d(50),
   nn.ReLU(True),
 
   nn.Conv2d(50,25,3,stride=1,padding=1),
   nn.BatchNorm2d(25),
   nn.ReLU(True),
 
   nn.Conv2d(25,1,2,stride=2),
   nn.Tanh()
  )
 def forward(self, x):
  x=self.fc(x)
  x=x.view(x.size(0),1,56,56)
  x=self.br(x)
  x=self.gen(x)
  return x
 
if __name__=="__main__":
 criterion=nn.BCELoss()
 num_img=100
 z_dimension=110
 D=discriminator()
 G=generator(z_dimension,3136) #1*56*56
 trainset, testset, trainloader, testloader = loadMNIST(num_img) # data
 D=D.cuda()
 G=G.cuda()
 d_optimizer=optim.Adam(D.parameters(),lr=0.0003)
 g_optimizer=optim.Adam(G.parameters(),lr=0.0003)
 '''
 交替訓練的方式訓練網絡
 先訓練判別器網絡D再訓練生成器網絡G
 不同網絡的訓練次數是超參數
 也可以兩個網絡訓練相同的次數，
 這樣就可以不用分別訓練兩個網絡
 '''
 count=0
 #鑒別器D的訓練,固定G的參數
 epoch = 119
 gepoch = 1
 for i in range(epoch):
  for (img, label) in trainloader:
   labels_onehot = np.zeros((num_img,10))
   labels_onehot[np.arange(num_img),label.numpy()]=1
#    img=img.view(num_img,-1)
#    img=np.concatenate((img.numpy(),labels_onehot))
#    img=torch.from_numpy(img)
   img=Variable(img).cuda()
   real_label=Variable(torch.from_numpy(labels_onehot).float()).cuda()#真實label為1
   fake_label=Variable(torch.zeros(num_img,10)).cuda()#假的label為0
 
   #compute loss of real_img
   real_out=D(img) #真實圖片送入判別器D輸出0~1
   d_loss_real=criterion(real_out,real_label)#得到loss
   real_scores=real_out#真實圖片放入判別器輸出越接近1越好
 
   #compute loss of fake_img
   z=Variable(torch.randn(num_img,z_dimension)).cuda()#隨機生成向量
   fake_img=G(z)#將向量放入生成網絡G生成一張圖片
   fake_out=D(fake_img)#判別器判斷假的圖片
   d_loss_fake=criterion(fake_out,fake_label)#假的圖片的loss
   fake_scores=fake_out#假的圖片放入判別器輸出越接近0越好
 
   #D bp and optimize
   d_loss=d_loss_real+d_loss_fake
   d_optimizer.zero_grad() #判別器D的梯度歸零
   d_loss.backward() #反向傳播
   d_optimizer.step() #更新判別器D參數
 
   #生成器G的訓練compute loss of fake_img
   for j in range(gepoch):
    z =torch.randn(num_img, 100) # 隨機生成向量
    z=np.concatenate((z.numpy(),labels_onehot),axis=1)
    z=Variable(torch.from_numpy(z).float()).cuda()
    fake_img = G(z) # 將向量放入生成網絡G生成一張圖片
    output = D(fake_img) # 經過判別器得到結果
    g_loss = criterion(output, real_label)#得到假的圖片與真實標簽的loss
    #bp and optimize
    g_optimizer.zero_grad() #生成器G的梯度歸零
    g_loss.backward() #反向傳播
    g_optimizer.step()#更新生成器G參數
    temp=real_label
  if (i%10==0) and (i!=0):
   print(i)
   torch.save(G.state_dict(),r'./CGAN/Generator_cuda_%d.pkl'%i)
   torch.save(D.state_dict(), r'./CGAN/Discriminator_cuda_%d.pkl' % i)
   save_model(G, r'./CGAN/Generator_cpu_%d.pkl'%i) #保存為CPU中可以打開的模型
   save_model(D, r'./CGAN/Discriminator_cpu_%d.pkl'%i) #保存為CPU中可以打開的模型
  print('Epoch [{}/{}], d_loss: {:.6f}, g_loss: {:.6f} '
     'D real: {:.6f}, D fake: {:.6f}'.format(
    i, epoch, d_loss.data[0], g_loss.data[0],
    real_scores.data.mean(), fake_scores.data.mean()))
  temp=temp.to('cpu')
  _,x=torch.max(temp,1)
  x=x.numpy()
  print(x[[6, 12, 18, 24, 30, 36, 42, 48, 54, 60, 66, 72, 78, 84, 90, 96]])
  showimg(fake_img,count)
  plt.show()
  count += 1

和基礎GAN Pytorch使用MNIST數據集實現基礎GAN 里面的卷積版網絡比較起來，這里修改的主要是這幾個地方：

生成網絡的輸入值增加了真實圖片的類標簽，生成網絡的初始向量z_dimension之前用的是100維，由于MNIST有10類，Onehot以后一張圖片的類標簽是10維，所以將類標簽放在后面z_dimension=100+10=110維；

訓練生成器的時候，由于生成網絡的輸入向量z_dimension=110維，而且是100維隨機向量和10維真實圖片標簽拼接，需要做相應的拼接操作；

z =torch.randn(num_img, 100) # 隨機生成向量
z=np.concatenate((z.numpy(),labels_onehot),axis=1)
z=Variable(torch.from_numpy(z).float()).cuda()

由于計算Loss和生成網絡的輸入向量都需要用到真實圖片的類標簽，需要重新生成real_label，對label進行onehot。其中real_label就是真實圖片的標簽，當num_img=100時，real_label的維度是（100,10）；

labels_onehot = np.zeros((num_img,10))
labels_onehot[np.arange(num_img),label.numpy()]=1
img=Variable(img).cuda()
real_label=Variable(torch.from_numpy(labels_onehot).float()).cuda()#真實label為1
fake_label=Variable(torch.zeros(num_img,10)).cuda()#假的label為0

real_label的維度是（100,10），計算Loss的時候也要有對應的維度，判別網絡的輸出也不再是標量，而是要修改為10維；

nn.Linear(1024, 10)

在輸出圖片的同時輸出期望的類標簽。

temp=temp.to('cpu')
_,x=torch.max(temp,1)#返回值有兩個，第一個是按列的最大值，第二個是相應最大值的列標號
x=x.numpy()
print(x[[6, 12, 18, 24, 30, 36, 42, 48, 54, 60, 66, 72, 78, 84, 90, 96]])

epoch等于0、25、50、75、100時訓練的結果：

可以看到訓練到后面圖像反而變模糊可能是訓練過擬合

用模型生成指定的數字：

在訓練的過程中保存了訓練好的模型，根據輸出圖片的清晰度，用清晰度較高的模型，使用隨機向量和10維類標簽來指定生成的數字。

import torch
import torch.nn as nn
import pickle
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
 
num_img=9
class discriminator(nn.Module):
 def __init__(self):
  super(discriminator, self).__init__()
  self.dis = nn.Sequential(
   nn.Conv2d(1, 32, 5, stride=1, padding=2),
   nn.LeakyReLU(0.2, True),
   nn.MaxPool2d((2, 2)),
 
   nn.Conv2d(32, 64, 5, stride=1, padding=2),
   nn.LeakyReLU(0.2, True),
   nn.MaxPool2d((2, 2))
  )
  self.fc = nn.Sequential(
   nn.Linear(7 * 7 * 64, 1024),
   nn.LeakyReLU(0.2, True),
   nn.Linear(1024, 10),
   nn.Sigmoid()
  )
 
 def forward(self, x):
  x = self.dis(x)
  x = x.view(x.size(0), -1)
  x = self.fc(x)
  return x
 
 
class generator(nn.Module):
 def __init__(self, input_size, num_feature):
  super(generator, self).__init__()
  self.fc = nn.Linear(input_size, num_feature) # 1*56*56
  self.br = nn.Sequential(
   nn.BatchNorm2d(1),
   nn.ReLU(True)
  )
  self.gen = nn.Sequential(
   nn.Conv2d(1, 50, 3, stride=1, padding=1),
   nn.BatchNorm2d(50),
   nn.ReLU(True),
 
   nn.Conv2d(50, 25, 3, stride=1, padding=1),
   nn.BatchNorm2d(25),
   nn.ReLU(True),
 
   nn.Conv2d(25, 1, 2, stride=2),
   nn.Tanh()
  )
 
 def forward(self, x):
  x = self.fc(x)
  x = x.view(x.size(0), 1, 56, 56)
  x = self.br(x)
  x = self.gen(x)
  return x
 
 
def show(images):
 images = images.detach().numpy()
 images = 255 * (0.5 * images + 0.5)
 images = images.astype(np.uint8)
 plt.figure(figsize=(4, 4))
 width = images.shape[2]
 gs = gridspec.GridSpec(1, num_img, wspace=0, hspace=0)
 for i, img in enumerate(images):
  ax = plt.subplot(gs[i])
  ax.set_xticklabels([])
  ax.set_yticklabels([])
  ax.set_aspect('equal')
  plt.imshow(img.reshape(width, width), cmap=plt.cm.gray)
  plt.axis('off')
  plt.tight_layout()
 plt.tight_layout()
 # plt.savefig(r'drive/深度學習/DCGAN/images/%d.png' % count, bbox_inches='tight')
 return width
 
def show_all(images_all):
 x=images_all[0]
 for i in range(1,len(images_all),1):
  x=np.concatenate((x,images_all[i]),0)
 print(x.shape)
 x = 255 * (0.5 * x + 0.5)
 x = x.astype(np.uint8)
 plt.figure(figsize=(9, 10))
 width = x.shape[2]
 gs = gridspec.GridSpec(10, num_img, wspace=0, hspace=0)
 for i, img in enumerate(x):
  ax = plt.subplot(gs[i])
  ax.set_xticklabels([])
  ax.set_yticklabels([])
  ax.set_aspect('equal')
  plt.imshow(img.reshape(width, width), cmap=plt.cm.gray)
  plt.axis('off')
  plt.tight_layout()
 
 
 # 導入相應的模型
z_dimension = 110
D = discriminator()
G = generator(z_dimension, 3136) # 1*56*56
D.load_state_dict(torch.load(r'./CGAN/Discriminator.pkl'))
G.load_state_dict(torch.load(r'./CGAN/Generator.pkl'))
# 依次生成0到9
lis=[]
for i in range(10):
 z = torch.randn((num_img, 100)) # 隨機生成向量
 x=np.zeros((num_img,10))
 x[:,i]=1
 z = np.concatenate((z.numpy(), x),1)
 z = torch.from_numpy(z).float()
 fake_img = G(z) # 將向量放入生成網絡G生成一張圖片
 lis.append(fake_img.detach().numpy())
 output = D(fake_img) # 經過判別器得到結果
 show(fake_img)
 plt.savefig('./CGAN/generator/%d.png' % i, bbox_inches='tight')
 
show_all(lis)
plt.savefig('./CGAN/generator/all.png', bbox_inches='tight')
plt.show()

生成的結果是：