數據集
數據集我們用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
第109個batch
第209個batch
第309個batch
第409個batch
第509個batch
第609個batch
第709個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