通过PyTorch用DCGAN生成动漫头像

数据集

数据集我们用AnimeFaces数据集,共5万多张动漫头像。

链接:https://pan.baidu.com/s/1cp-A8ZV74YBelkSuKxuM6A
提取码:face

要把所有的图片保存于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

  • 11
    点赞
  • 25
    收藏
    觉得还不错? 一键收藏
  • 3
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 3
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值