通過PyTorch用DCGAN生成動漫頭像

數據集

數據集我們用AnimeFaces數據集,共5萬多張動漫頭像。

https://pan.baidu.com/s/1eSifHcA 提取碼:g5qa

要把所有的圖片保存於data/face/目錄下,後邊用ImageFolder就能直接讀取到。

模型

模型我們選擇DCGAN。

#coding:utf-8

from torch import nn

class NetG(nn.Module):
	"""
	生成器定義
	"""
	def __init__(self,opt):
		super(NetG,self).__init__()
		ngf = opt.ngf # 生成器feature map數 

		self.main = nn.Sequential(
			# 輸入是一個nz維度的噪聲,我們可以認爲它是一個1*1×nz的feature map
			nn.ConvTranspose2d(opt.nz,ngf*8,4,1,0,bias=False),
			nn.BatchNorm2d(ngf*8),
			nn.ReLU(True),
			# 這一步的輸出形狀:(ngf*8)*4*4

			nn.ConvTranspose2d(ngf*8,ngf*4,4,2,1,bias=False),
			nn.BatchNorm2d(ngf*4),
			nn.ReLU(True),
			# 這一步的輸出形狀:(ngf*4)*8*8

			nn.ConvTranspose2d(ngf*4,ngf*2,4,2,1,bias=False),
			nn.BatchNorm2d(ngf*2),
			nn.ReLU(True),
			# 這一步的輸出形狀:(ngf*2)*16*16

			nn.ConvTranspose2d(ngf*2,ngf,4,2,1,bias=False),
			nn.BatchNorm2d(ngf),
			nn.ReLU(True),
			# 這一步的輸出形狀:(ngf*1)*32*32

			nn.ConvTranspose2d(ngf,3,5,3,1,bias=False),
			nn.Tanh(), # 輸出範圍-1~1,故而採用Tanh
			# 最後的輸出形狀:3*96*96
			)

	def forward(self,x):
		return self.main(x)

'''
這裏需要注意上卷積ConvTransposed2d的使用。當kernel_size爲4,stride爲2,padding爲1時,根據公式輸出尺寸剛好變成輸入的兩倍。
最後一層採用kernel_size爲5,stride爲3,padding爲1,是爲了將32*32上採樣到96*96,這正好是我們輸入圖片的尺寸。
最後一層用Tanh將輸出圖片的像素歸一化到-1~1,如果希望歸一化到0~1則需要使用Sigmoid。

'''

class NetD(nn.Module):
	"""
	判別器定義
	"""
	def __init__(self,opt):
		super(NetD,self).__init__()
		ndf = opt.ndf # 判別器feature map數

		self.main = nn.Sequential(
			# 輸入 3*96*96
			nn.Conv2d(3,ndf,5,3,1,bias=False),
			nn.LeakyReLU(0.2,inplace=True),
			# 輸出 (ndf*1)*32*32

			nn.Conv2d(ndf,ndf*2,4,2,1,bias=False),
			nn.BatchNorm2d(ndf*2),
			nn.LeakyReLU(0.2,inplace=True),
			# 輸出 (ndf*2)*16*16

			nn.Conv2d(ndf*2,ndf*4,4,2,1,bias=False),
			nn.BatchNorm2d(ndf*4),
			nn.LeakyReLU(0.2,inplace=True),
			# 輸出 (ndf*4)*8*8

			nn.Conv2d(ndf*4,ndf*8,4,2,1,bias=False),
			nn.BatchNorm2d(ndf*8),
			nn.LeakyReLU(0.2,inplace=True),
			# 輸出 (ndf*8)*4*4

			nn.Conv2d(ndf*8,1,4,1,0,bias=False),
			nn.Sigmoid() # 輸出一個數(概率)
			)

	def forward(self,x):
		return self.main(x).view(-1)

'''
判別器和生成器的網絡結構差不多是對稱的。
這裏需要注意的是生成器的激活函數用的是ReLU,而判別器使用的是LeakyReLU,二者並無本質區別,這裏的選擇更多是經驗總結。
每一個樣本經過判別器後,輸出一個0~1的數,表示這個樣本是真圖片的概率。

'''

訓練過程

train.py

#coding:utf-8

import os
import torch as t
import torchvision as tv
import tqdm
from model import NetG,NetD
import time
import numpy as np
import scipy.io as io

# 在訓練函數前,先寫配置參數
class Config(object):
	data_path = 'data/' # 數據集存放路徑
	num_workers = 4 # 多進程加載數據所用的進程數
	image_size = 96 # 圖片尺寸
	batch_size = 256 # 批處理數
	max_epoch = 200 # 訓練的總輪數
	last_epoch = 0 # 上次訓練到的位置,默認爲0
	lr1 = 2e-4 # 生成器的學習率
	lr2 = 2e-4 # 判別器的學習率
	beta1 = 0.5 # Adam優化器的beta1參數
    beta2 = 0.999 # Adam優化器的beta2參數
	gpu = True # 是否使用GPU
	nz = 100 # 噪聲維度,用於生成器生成圖片
	ngf = 64 # 生成器feature map數
	ndf = 64 # 判別器feature map數

	save_path = 'imgs' # 生成圖片保存路徑
	d_every = 1 # 每1個batch訓練一次判別器
	g_every = 5 # 每5個batch訓練一次生成器
	save_every = 10 # 每10個epoch保存一次模型
	netd_path = None # 'netd_num.pth' 模型參數文件
	netg_path = None # 'netg_num.pth'

opt = Config()
'''
這些只是模型的默認參數,可以利用Fire等工具通過命令行傳入,覆蓋默認值。

'''

# 訓練過程
def train(**kwargs):
	for k_,v_ in kwargs.items(): # 加載參數
		setattr(opt,k_,v_)

	# 數據預處理
	transforms = tv.transforms.Compose([
		tv.transforms.Resize(opt.image_size), # 調整圖片大小
		tv.transforms.CenterCrop(opt.image_size), # 中心裁剪
		tv.transforms.ToTensor(),
		tv.transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5)) # 中心化
	])

	# 加載數據集
	dataset = tv.datasets.ImageFolder(opt.data_path,transform=transforms)
	dataloader = t.utils.data.DataLoader(dataset,batch_size=opt.batch_size,shuffle=True,num_workers=opt.num_workers,drop_last=True) # drop_last表示不用數據集最後不足一個batch的數據
	print("dataset:"+str(len(dataset))+",dataloader:"+str(len(dataloader)))
	# 網絡,使用gpu
	if opt.gpu:
		if t.cuda.is_available():
			netg,netd = NetG(opt).cuda(),NetD(opt).cuda()
			print("Train CUDA OK!")
	else:
		netg,netd = NetG(opt),NetD(opt)
	
	# 在加載預訓練模型時,最好指定map_location
	# 因爲如果程序之前在GPU上運行,那麼模型就會被存成torch.cuda.Tensor,這樣加載時會默認將數據加載至顯存。
	# 如果運行改程序的計算機中沒有GPU,加載就會報錯,故通過指定map_location將Tensor默認加載入內存中,待有需要時再移至顯存中。
	map_location = lambda storage,loc: storage
	if opt.netd_path:
		netd.load_state_dict(t.load('checkpoints/%s'%opt.netd_path,map_location=map_location))
		print("%s"%opt.netd_path,"loading...OK!")
	if opt.netg_path:
		netg.load_state_dict(t.load('checkpoints/%s'%opt.netg_path,map_location=map_location))
		print("%s"%opt.netg_path,"loading...OK!")


	# 定義優化器和損失函數
	optimizer_g = t.optim.Adam(netg.parameters(),opt.lr1,betas=(opt.beta1,opt.beta2))
	optimizer_d = t.optim.Adam(netd.parameters(),opt.lr2,betas=(opt.beta1,opt.beta2))
	criterion = t.nn.BCELoss()

	# 真圖片label爲1,假圖片label爲0
	# noises爲生成網絡的輸入
	true_labels = t.ones(opt.batch_size).cuda()
	fake_labels = t.zeros(opt.batch_size).cuda()
	noises = t.randn(opt.batch_size,opt.nz,1,1).cuda()
	fix_noises = t.randn(opt.batch_size,opt.nz,1,1).cuda()
    
	# 使用已經保存的噪聲,保存生成的fix_noises的方法,會在下面顯示出來
	# mat_noises = io.loadmat('noises_double.mat') # 讀取文件加載noises
	# fix_noises = t.from_numpy(mat_noises['np_noises']).cuda() # 重新轉換成tensor


	now = time.clock()
	epochs = range(opt.last_epoch,opt.max_epoch)
	for epoch in iter(epochs):
		g_loss = 0 # 這裏爲了平均訓練一個epoch的損失值
		d_loss = 0
		for ii,(img,_) in tqdm.tqdm(enumerate(dataloader)):
			real_img = img.cuda()

			# 每1個batch訓練一次判別器
			if ii%opt.d_every == 0:
				# 訓練判別器
				optimizer_d.zero_grad() # 梯度清零
				# 儘可能的把真圖片判別爲1
				output = netd(real_img)
				error_d_real = criterion(output,true_labels)
				error_d_real.backward() # 真圖片,反向傳播

				# 儘可能把假圖片(生成器生成的)判別爲0
				noises.data.copy_(t.randn(opt.batch_size,opt.nz,1,1)) # noises的值改變了,copy_直接覆蓋原有的noises值
				fake_img = netg(noises).detach() # 根據噪聲生成假圖 .detach()安全的獲得out的值,比.data安全,避免梯度傳遞到G上,因爲訓練D時不更新G
				output = netd(fake_img)
				error_d_fake = criterion(output,fake_labels)
				error_d_fake.backward() # 假圖片,反向傳播
				optimizer_d.step() # 更新參數

				error_d = error_d_real + error_d_fake
				d_loss += error_d.item()

			# 每5個batch訓練一次生成器
			if ii % opt.g_every == 0:
				# 訓練生成器
				optimizer_g.zero_grad()

				noises.data.copy_(t.randn(opt.batch_size,opt.nz,1,1))
				fake_img = netg(noises)
				output = netd(fake_img)
				error_g = criterion(output,true_labels)
				error_g.backward()
				optimizer_g.step()
				
				g_loss += error_g.item()


		# 輸出友好信息
		print("Epoch:{},D_Loss:{:.6f},G_Loss:{:.6f},Time:{:.4f}s".format(epoch,2*d_loss/len(dataset),5*g_loss/len(dataset),time.clock()-now))


		
		# 保存模型、圖片,這裏每次保存一次圖片
		# 噪聲可以用我們之前保存的noises.mat文件中的noises
		fix_fake_imgs = netg(fix_noises)
		tv.utils.save_image(fix_fake_imgs.data[:64],'%s/%s.png'%(opt.save_path,epoch),normalize=True,range=(-1,1)) # 這裏只保存前64張96*96圖片,它們是拼在一起的
        if epoch%opt.save_every == 0: # 這樣做就可以每次10個epoch保存一個checkpoint
			t.save(netd.state_dict(),'checkpoints/netd_%s.pth'%epoch)
			t.save(netg.state_dict(),'checkpoints/netg_%s.pth'%epoch)
			# t.cuda.empty_cache() # 週期性的清理顯存


if __name__ == '__main__':
	import fire
	fire.Fire()

​ 這裏可以每1個batch訓練一次判別器並訓練一次生成器,也可以每1個batch訓練一次生成器並3個batch才訓練一次生成器,這些模型都會收斂,只是速度的快慢。

​ 但是,我實驗了每1個batch訓練一次生成器的同時每3個batch訓練一次生成器,這樣的模型訓練不起來。雖然g_loss會比上面那樣的低,但是這並不代表結果就好。

​ 我們可以通過下面代碼保存我們的隨機生成的noise。這樣我們在訓練過程中就可以通過這保存文件中的noise來生成圖片,進而可以方便觀察模型收斂的過程。

import numpy as np
import scipy.io as io

noises = t.randn(64,100,1,1) # B×C×w×H
np_noises = np.array(noises) # 先將tensor轉換爲array
io.savemat('noises_double.mat',{'np_noises':np_noises}) # 以鍵值對的形式,保存在.mat文件中
mat_noises = io.loadmat('noises.mat') # 讀取這個文件
noises = t.from_numpy(mat_noises['np_noises']) # 重新轉換成tensor
noises = noises.cuda()

​ 運行時可以通過終端敲這樣的形式運行訓練程序,參數以–開頭,字符串的雙引號可以省略。

python train.py train --netg_path=net_800.pth ……

訓練結果

由於我是斷斷續續訓練的,打印的損失值的信息沒有保存下來。由於剛開始沒有意識到把noise存下來的好處,而且剛開始我在存圖片的時候也是每10個epoch才存一張,所以下面圖片不連貫。,但是可以清楚的發現,模型是在不停的收斂的。(我本來做了一個gif,但是CSDN傳不了那麼大的,只能找中間幾張圖貼出來了)

第9個batch
第9個batch第109個batch
第109個batch第209個batch
第209個batch第309個batch
第309個batch第409個batch
第409個batch第509個batch
第509個batch第609個batch
第609個batch第709個batch
第709個batch第800個batch
第800個batch
​ 我總共訓練800個epoch(時間花了很久,1060每batch也花了1分多鐘),在300個epoch左右生成的圖片很少有包含嘴巴的,到後邊嘴巴慢慢生成了,這說明模型還在收斂。在訓練到600-800epoch時,模型幾乎已經不能再變好了,有的圖片已經很逼真了,但是相比較訓練數據集中的真實圖片還是有區別的。而且圖片的分辨率才96*96,太小了,所以看起來不是很高清。

​ 我在想是不是模型太小了,生成網絡NetG和判別網絡NetD內主要都只是由5層卷積層組成的,圖片的有些特徵是不是還沒有被學習到?還是損失函數的選擇,會不會有更好的選擇?怎樣才能生成更高清的圖片呢?

​ 這需要去實驗,我已經換一個模型在訓練了,用的是DRGAN中的網絡,之後會整理再發到博客上。

測試過程

用tkinter寫一個簡單的GUI來顯示測試生成的圖片。

test.py

#coding:utf-8

from tkinter import *
from PIL import Image,ImageTk
from torch import nn
import torch as t
import torchvision as tv

class tk_main:
	def __init__(self):
		# 創建窗口,標題,大小
		self.window = Tk()
		self.window.title("Image")
		self.window.geometry('900x900')


	# 初始化模型
	def model_init(self):
		# 模型參數文件
		netd_path = 'netd_800.pth' # 這裏放訓練到最後生成的模型參數文件
		netg_path = 'netg_800.pth'
		if t.cuda.is_available():
				netg,netd = NetG(64).cuda().eval(),NetD(64).cuda().eval() # 默認的ngf和ndf都是64,所以這裏我直接傳給模型64
				print("Test CUDA OK!")
		else:
			netg,netd = NetG(),NetD()

		# 將模型參數加載到內存中
		map_location = lambda storage,loc: storage
		if netd_path:
			netd.load_state_dict(t.load(netd_path,map_location=map_location))
			print("%s"%netd_path,"loading...OK!")
		if netg_path:
			netg.load_state_dict(t.load(netg_path,map_location=map_location))
			print("%s"%netg_path,"loading...OK!")
		return netg,netd


	def Generate(self,netg,netd):
		"""
		隨機生成動漫圖片,並根據netd的分數選擇較好的
		"""
		# 生成圖片存放地址
		img_path = 'result.png'

		with t.no_grad():
			# 噪聲的生成,2048*100*1*1
			noises = t.randn(2048,100,1,1)
			noises = noises.cuda()


			# 生成圖片,並計算圖片在判別器的分數
			fake_img = netg(noises)
			scores = netd(fake_img).detach()
			# 挑選最好的某幾張,這裏是從2048張圖片中挑出64張
			indexs = scores.topk(64)[1] # 這裏是因爲topk()返回兩個列表,一個是具體值,一個是具體值在原輸入張量中的索引
			result = []
			for ii in indexs:
				result.append(fake_img.data[ii])
			# 保存圖片,這裏用到這個stack()函數,是因爲我們是挑選出來的圖片,需要將它們拼接在一起
			tv.utils.save_image(t.stack(result),img_path, normalize=True, range=(-1, 1))
			print('圖片存儲成功!')
			
			# 加載圖片
			load = Image.open(img_path)
			render = ImageTk.PhotoImage(load)

			img = Label(image=render)
			img.image = render
			img.place(x=57, y=57) # 圖片居中顯示


	def run(self):
		netg,netd = self.model_init()
		# 生成圖片
		Button(self.window,text='單擊生成64張動漫圖片',command=lambda :self.Generate(netg,netd)).pack()

		# 主窗口循環顯示
		self.window.mainloop()

if __name__ == '__main__':
	root = tk_main()
	root.run()

​ 這裏有個問題,在判別器判別生成器時輸出的是一個數(得分),我們借這個數來排序找到得分最高的64張圖片顯示出來。但是就是這樣一個得分的判斷有問題。生成的有的圖片,我們人爲看起來很明顯它不符合要求,但是它的得分卻很高。我認爲有可能它符合判別器的判斷標準。不過綜合來看,生成圖片中有些圖片還是符合要求的。

在這裏插入圖片描述

​ 可以很明顯看到,生成的圖片是有很多缺陷的,有的人物的雙眼是不同顏色的,有的沒有嘴巴,有的少一隻眼睛等。圖片基本上能生成,下面就該思考如何讓模型更加強悍。

參考

https://github.com/chenyuntc/pytorch-book/tree/master/chapter7-GAN%E7%94%9F%E6%88%90%E5%8A%A8%E6%BC%AB%E5%A4%B4%E5%83%8F

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章