GAN二次元头像生成Pytorch实现（附完整代码）
介绍
本文是李宏毅GAN课程课后作业HW3_1(二次元头像生成，Keras实现)的Pytorch版本。写这篇的原因是一方面刚开始接触GAN，二是个人比较习惯用Pytorch，所以将keras改成Pytorch版本。
实现所需要的资源：
链接：https://pan.baidu.com/s/1cLmFNQpJe1DOI96IVuvVyQ
提取码：nha2
本文一个改动就是将kernel=4变成了3，因为kernel一般都是奇数。其他和原网络基本相同。
下面是主要部分的代码，包括网络模块和训练/验证/测试两个模块。
完整的代码见 https://github.com/AsajuHuishi/Generate_a_quadratic_image_with_GAN
import torch
import torch.optim as optim
import torch.optim.lr_scheduler as lr_scheduler
import torch.nn as nn
from torch.autograd import Variable
import matplotlib.pyplot as plt
import numpy as np
import os
import argparse
import time
import visdom
1.网络模块
生成器
##定义卷积核
def default_conv(in_channels,out_channels,kernel_size,bias=True):
return nn.Conv2d(in_channels,out_channels,
kernel_size,padding=kernel_size//2, #保持尺寸
bias=bias)
##定义ReLU
def default_relu():
return nn.ReLU(inplace=True)
## reshape
def get_feature(x):
return x.reshape(x.size()[0],128,16,16)
class Generator(nn.Module):
def __init__(self,input_dim=100,conv=default_conv,relu=default_relu,reshape=get_feature):
super(Generator,self).__init__()
head = [nn.Linear(input_dim,128*16*16),
relu()]
self.reshape = reshape #16x16
body = [nn.Upsample(scale_factor=2,mode='nearest'), #32x32
conv(128,128,3),
relu(),
nn.Upsample(scale_factor=2,mode='nearest'), #64x64
conv(128,64,3),
relu(),
conv(64,3,3),
nn.Tanh()]
self.head = nn.Sequential(*head)
self.body = nn.Sequential(*body)
def forward(self,x):#x:(batchsize,input_dim)
x = self.head(x)
x = self.reshape(x)
x = self.body(x)
return x #(batchsize,3,64,64)
def name(self):
return 'Generator'
判别器
class Discriminator(nn.Module):
def __init__(self,conv=default_conv,relu=default_relu):
super(Discriminator,self).__init__()
main = [conv(3,32,3),
relu(),
conv(32,64,3),
relu(),
conv(64,128,3),
relu(),
conv(128,256,3),
relu()]
self.main = nn.Sequential(*main)
self.fc = nn.Linear(256*64*64,1)
self.sigmoid = nn.Sigmoid()
def forward(self,x):#x:(batchsize,3,64,64)
x = self.main(x)#(b,256,64,64)
x = x.view(x.size()[0],-1)#(b,256*64*64)
x = self.fc(x) #(b,1)
x = self.sigmoid(x)
return x
def name(self):
return 'Discriminator'
2.训练/验证/测试模块
相关参数、模型初始化
class GAN(nn.Module):
def __init__(self,args):
super(GAN,self).__init__()
self.img_size = 64
self.channels = 3
self.latent_dim = args.latent_dim
self.num_epoch = args.num_epoch
self.batch_size = args.batch_size
self.cuda = args.cuda
self.interval = 20 #每相邻20个epoch验证一次
self.continue_training = args.continue_training #是否是继续训练
## 生成器初始化
self.generator = Generator(self.latent_dim)
## 判别器初始化
self.discriminator = Discriminator()
self.testmodelpath = args.testmodelpath
self.datapath = args.datapath
if self.cuda:
self.generator.cuda()
self.discriminator.cuda()
self.continue_training_isrequired()
训练+dataloader数据集
def trainer(self):
## 读入图片数据,分batch
print('===> Data preparing...')
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
transform = transforms.ToTensor() ##dataloader输出是tensor，不加这个会报错
dataset = ImageFolder(self.datapath,transform=transform)
dataloader = DataLoader(dataset, batch_size=self.batch_size, shuffle=True, num_workers=0, drop_last=True)
##drop_last: dataset中的数据个数可能不是batch_size的整数倍，drop_last为True会将多出来不足一个batch的数据丢弃
num_batch = len(dataloader) #batch的数量为len(dataloader)=总图片数/batchsize
print('num_batch:',num_batch)
#dataloader: (batchsize,3,64,64) 分布0-1
## 判别值
target_real = Variable(torch.ones(self.batch_size,1))
target_false = Variable(torch.zeros(self.batch_size,1))
one_const = Variable(torch.ones(self.batch_size,1))
if self.cuda:
target_real = target_real.cuda()
target_false = target_false.cuda()
one_const = one_const.cuda()
## 优化器
optim_generator = optim.Adam(self.generator.parameters(),lr=0.0002,betas=(0.5,0.999))
optim_discriminator = optim.Adam(self.discriminator.parameters(),lr=0.0002,betas=(0.5,0.999))
## 误差函数
# content_criterion = nn.MSELoss()
adversarial_criterion = nn.BCELoss()
## 训 练 开 始
for epoch in range(self.start_epoch,self.num_epoch): ##epoch
##用于观察一个epoch不同batch的平均loss
mean_dis_loss = 0.0
mean_gen_con_loss = 0.0
mean_gen_adv_loss = 0.0
mean_gen_total_loss = 0.0
for i,data in enumerate(dataloader): ##循环次数：batch的数量为len(dataloader)=总图片数//batchsize
if epoch<3 and i%10==0:
print('epoch%d: %d/%d'%(epoch,i,len(dataloader)))
##1.1生成noise
gen_input = np.random.normal(0,1,(self.batch_size,self.latent_dim)).astype(np.float32)
gen_input = torch.from_numpy(gen_input)
gen_input = torch.autograd.Variable(gen_input,requires_grad=True)
if self.cuda:
gen_input = gen_input.cuda()
fake = self.generator(gen_input) ##生成器生成的(batchsize,3,64,64)
real, _ = data #data:list[tensor,tensor]取第零个 real:(batchsize,3,64,64)
if self.cuda:
real = real.cuda()
fake = fake.cuda()
## 1.固定G,训练判别器D
self.discriminator.zero_grad()
dis_loss1 = adversarial_criterion(self.discriminator(real),target_real)
dis_loss2 = adversarial_criterion(self.discriminator(fake.detach()),target_false)##注意经过G的网络再进入D网络之前要detach()之后再进入
dis_loss = 0.5*(dis_loss1+dis_loss2)
# print('epoch:%d--%d,判别器loss:%.6f'%(epoch,i,dis_loss))
dis_loss.backward()
optim_discriminator.step()
mean_dis_loss+=dis_loss
## 2.固定D,训练生成器G
self.generator.zero_grad()
##生成noise
gen_input = np.random.normal(0,1,(self.batch_size,self.latent_dim)).astype(np.float32)
gen_input = torch.from_numpy(gen_input)
gen_input = torch.autograd.Variable(gen_input,requires_grad=True)
if self.cuda:
gen_input = gen_input.cuda()
fake = self.generator(gen_input) ##生成器生成的(batchsize,3,64,64)
gen_con_loss = 0
gen_adv_loss = adversarial_criterion(self.discriminator(fake),one_const)##固定D更新G
gen_total_loss = gen_con_loss + gen_adv_loss
# print('epoch:%d--%d,生成器loss:%.6f'%(epoch,i,gen_total_loss))
gen_total_loss.backward()
optim_generator.step()
mean_gen_con_loss+=gen_con_loss
mean_gen_adv_loss+=gen_adv_loss
mean_gen_total_loss+=gen_total_loss
## 一个epoch输出一次
print('epoch:%d/%d'%(epoch, self.num_epoch))
print('Discriminator_Loss: %.4f'%(mean_dis_loss/num_batch))
print('Generator_total_Loss:%.4f'%(mean_gen_total_loss/num_batch))
## 保存模型
state_dis = {'dis_model': self.discriminator.state_dict(), 'epoch': epoch}
state_gen = {'gen_model': self.generator.state_dict(), 'epoch': epoch}
if not os.path.isdir('checkpoint'):
os.mkdir('checkpoint')
torch.save(state_dis, 'checkpoint/'+self.discriminator.name()+'__'+str(epoch+1)) #each epoch
torch.save(state_gen, 'checkpoint/'+self.generator.name()+'__'+str(epoch+1)) #each epoch
torch.save(state_dis, 'checkpoint/'+self.discriminator.name()) #final
torch.save(state_gen, 'checkpoint/'+self.generator.name()) #final
## 验证模型
if epoch<45 or epoch%self.interval==0:
self.validater(epoch)
print('--'.center(12,'-'))
验证
def validater(self,epoch):
vis = visdom.Visdom(env='generate_girl_epoch%d'%(epoch))
r,c = 3,3
gen_input_val = np.random.normal(0,1,(r*c,self.latent_dim)).astype(np.float32)
gen_input_val = torch.from_numpy(gen_input_val)
gen_input_val = torch.autograd.Variable(gen_input_val)
if self.cuda:
gen_input_val = gen_input_val.cuda()
output_val = self.generator(gen_input_val) #(r*c,3,64,64)
output_val = output_val.cpu()
output_val = output_val.data.numpy() #(r*c,3,64,64)
img = np.transpose(output_val,(0,2,3,1)) #(r*c,64,64,3)
fig, axs = plt.subplots(r,c)
cnt = 0
for i in range(r):
for j in range(c):
vis.image(output_val[cnt],opts={'title':'epoch%d_cnt%d'%(epoch,cnt)})
axs[i, j].imshow(img[cnt, :, :, :])
axs[i, j].axis('off')
cnt += 1
if not os.path.isdir('images'):
os.mkdir('images')
fig.savefig('images/val_%d.png'%(epoch+1)) ##保存验证结果
plt.close()
测试
def tester(self,gen_input_test): #输入(N,latent_dim)
assert gen_input_test.shape[1]==self.latent_dim, \
'dimension 1''s size expect %d,but input %d'%(self.latent_dim,gen_input_test.shape[1])
gen_input_test = gen_input_test.astype(np.float32)
gen_input_test = torch.from_numpy(gen_input_test)
gen_input_test = torch.autograd.Variable(gen_input_test)
if self.cuda:
gen_input_test = gen_input_test.cuda()
## 下载验证结果
if os.path.isdir('checkpoint'):
try:
checkpoint_gen = torch.load(self.testmodelpath)
self.generator.load_state_dict(checkpoint_gen['gen_model'])
except FileNotFoundError:
print('Can\'t found dict')
output_test = self.generator(gen_input_test)
output_test = output_test.cpu()
output_test = output_test.data.numpy() #(N,3,64,64)
img = np.transpose(output_test,(0,2,3,1)) #(N,64,64,3)
if not os.path.isdir('images'):
os.mkdir('images')
N = img.shape[0] #图像个数
for i in range(N):
plt.imshow(img[i, :, :, :])
plt.axis('off')
plt.savefig('images/test_%d.png'%(i+1)) ##保存结果
plt.close()
结果和原keras相比没什么区别，毕竟网络都差不多，也不需要过高期望，而且网络本身比较小，生成一个好看的人脸，对是五官是否协调有很大的要求，是很有挑战的事情。
输入：
np.random.normal(0,1,(1,self.latent_dim)).astype(np.float32)
部分结果：
