1. 程式人生 > 程式設計 >Pytorch使用MNIST資料集實現CGAN和生成指定的數字方式

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,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,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,2))
  )
  self.fc=nn.Sequential(
   nn.Linear(7 * 7 * 64,1024),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,padding=1),nn.BatchNorm2d(50),nn.ReLU(True),nn.Conv2d(50,25,nn.BatchNorm2d(25),nn.Conv2d(25,1,2,stride=2),nn.Tanh()
  )
 def forward(self,x):
  x=self.fc(x)
  x=x.view(x.size(0),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,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,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(),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')
_,1)#返回值有兩個,第一個是按列的最大值,第二個是相應最大值的列標號
x=x.numpy()
print(x[[6,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,2))
  )
  self.fc = nn.Sequential(
   nn.Linear(7 * 7 * 64,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,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,nn.Tanh()
  )
 
 def forward(self,x):
  x = self.fc(x)
  x = x.view(x.size(0),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,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,img in enumerate(x):
  ax = plt.subplot(gs[i])
  ax.set_xticklabels([])
  ax.set_yticklabels([])
  ax.set_aspect('equal')
  plt.imshow(img.reshape(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()

生成的結果是:

以上這篇Pytorch使用MNIST資料集實現CGAN和生成指定的數字方式就是小編分享給大家的全部內容了,希望能給大家一個參考,也希望大家多多支援我們。