基于pytorch简单实现DCGAN
前言
最近会把一些简单的CV领域的架构进行复现,完整的代码在最后。
本系列必须的基础
python基础知识、CNN原理知识、pytorch基础知识
本系列的目的
一是帮助自己巩固知识点;
二是自己实现一次,可以发现很多之前的不足;
三是希望可以给大家一个参考。
参考资料
GitHub项目:
https://github.com/znxlwm/pytorch-MNIST-CelebA-GAN-DCGAN
文章:
https://zhuanlan.zhihu.com/p/48501100 --- 反卷积计算公式推导
https://blog.csdn.net/qq_41605740/article/details/127816320 --- BCELoss介绍
https://blog.csdn.net/m0_62128864/article/details/123935874 --- DCGAN实现
目录结构
文章目录
- 基于pytorch简单实现DCGAN
- 1. 前言
- 2. 生成器和判别器实现
- 3. 训练前的准备
- 4. 实现训练过程
- 5. 结果分析
- 6. 总结
1. 前言
前一篇基于pytorch实现了CGAN,但是效果不是很好。于是打算试试DCGAN,因为相比于CGAN,DCGAN更加强大,而且两者整体代码差不多,也就是架构采用的不同。
必须一提的是,DCGAN不是D + CGAN,DCGAN是Deep Convolution GAN,而CGAN中的C是conditional。
另外,前一篇,将代码分开写在不同文件夹中,感觉有点没必要,因为代码量其实不多。这里就直接放一个文件中了。
我的目录结构:
---- DCGAN
---- fake_images # 用于保存生成器生成图片的文件夹
---- pytorch_dcgan.py # 主要代码文件
---- data
---- mnist # MNIST数据集文件夹
2. 生成器和判别器实现
首先,贴一张图,来自参考资料中GitHub项目的图片:
这张图就是生成器和判别器的主要架构图。其中需要注意的几点如下:
- 生成器全由反卷积实现,最终输出为图像大小
- 由于MNIST图像为灰度图,因此维度只有1
- 判别器最终输出一个值,并且使用sigmoid将值限定为0到1之间,表示这个输入图像为真实or假的概率值
- 上图中绿色部分就是卷积核
- 值得注意的是,作者似乎把原来28*28*1的MNIST图像转为了64*64*1,目的应该是为了方便生成器中反卷积的尺寸设计等
一般来说,有了上面的图,实现起来就非常简单了。这里先来实现判别器,因为判别器没有用到反卷积,用的都是平常常见的东西:(看注释)
# 判别器
class Discriminator(nn.Module):
def __init__(self):
super(Discriminator, self).__init__()
self.model = nn.Sequential(
# 按照图片上的实现即可
nn.Conv2d(1,128,kernel_size=4,stride=2,padding=1),
# 这里0.2是图片中采取的值
nn.LeakyReLU(0.2),
nn.Conv2d(128,256,4,2,1),
nn.BatchNorm2d(256),
nn.LeakyReLU(0.2),
nn.Conv2d(256, 512, 4, 2, 1),
nn.BatchNorm2d(512),
nn.LeakyReLU(0.2),
nn.Conv2d(512, 1024, 4, 2, 1),
nn.BatchNorm2d(1024),
nn.LeakyReLU(0.2),
# 最终输出1不要忘记即可
nn.Conv2d(1024, 1, 4, 2, 0),
# 加上一个sigmoid,控制输出到0-1
nn.Sigmoid(),
)
def forward(self,x):
result = self.model(x)
return result
接着,来实现生成器,由于生成器采用了反卷积,这里放一下反卷积常用的计算输出尺寸公式:
以上图为例,第一个反卷积:(上图已经给出了输出尺寸、通道数、卷积核大小、步长,但是没有给出padding值,所以需要自己计算一下)
输入尺寸为: 1*1,维度为100
卷积核:4*4,s=1
输出尺寸目标为: 4*4*1024
所以,输出通道数为1024
根据公式:4 = 1*(1-1)-2*p+4,得到p=0
根据上述计算过程,第一个反卷积应该这么写:
nn.ConvTranspose2d(100,1024,4,1,0),
# 对应参数:输入通道数、输出通道数、卷积核大小、步长、padding值
那么,依次计算,可以完成生成器的构建:
# 生成器
class Generator(nn.Module):
def __init__(self):
super(Generator, self).__init__()
self.model = nn.Sequential(
# 输入通道、输出通道、卷积核、步长、padding
nn.ConvTranspose2d(100,1024,4,1,0),
nn.BatchNorm2d(1024),
nn.ReLU(),
nn.ConvTranspose2d(1024, 512, 4, 2, 1),
nn.BatchNorm2d(512),
nn.ReLU(),
nn.ConvTranspose2d(512, 256, 4, 2, 1),
nn.BatchNorm2d(256),
nn.ReLU(),
nn.ConvTranspose2d(256, 128, 4, 2, 1),
nn.BatchNorm2d(128),
nn.ReLU(),
# 最后,输出通道数为1,是因为为灰度图
nn.ConvTranspose2d(128, 1, 4, 2, 1),
nn.Tanh(),
)
def forward(self,x):
result = self.model(x)
return result
3. 训练前的准备
现在来完成训练前的准备工作,即定义基本参数、优化器、损失函数等内容。
首先,定义采用的设备:(单GPU)
# 设备
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
然后,定义训练多少次,batch大小、初始学习率和损失函数:
# batch大小、epoch次数、初始学习率、损失函数
batch_size = 128
epoch = 200
lr = 0.002
loss = nn.BCELoss()
这里,需要简单介绍一下BCELoss这个损失函数,该函数叫做二值交叉熵损失函数,一看就知道是用于二值分类的。感兴趣的可以看看参考资料中的介绍。
接着,创建模型、定义优化器:
# 模型创建
D = Discriminator().to(device)
G = Generator().to(device)
# 优化器
optim_G = optim.Adam(G.parameters(),lr=lr,betas=(0.5, 0.999))
optim_D = optim.Adam(D.parameters(),lr=lr,betas=(0.5, 0.999))
最后,定义一下数据加载器,由于我们用的是MNIST数据集,而pytorch官方已经实现了这个数据集的加载,所以可以很简单的实现。但是,注意的是,模型要求输入大小为64*64,所以还需要定义一下预处理方法:
# 预处理方法
transforms_func = transforms.Compose(
[transforms.Resize(64), # 缩放到64*64
transforms.ToTensor(), # 转为tensor
transforms.Normalize(mean=0.5,std=0.5)] # 归一化处理
)
# 加载数据
dataset = MNIST('../data/mnist', # 确定自己加载的路径或者要加载保存的路径
train=True,
transform=transforms_func,
download=False) # 这个参数,根据自己有没有数据集确定
train_loader = DataLoader(dataset,batch_size=batch_size,shuffle=True,drop_last=True)
4. 实现训练过程
训练过程,其实和CGAN训练过程一样,同时也是按照原始GAN定义的训练过程进行训练的,即先训练一步D、在训练G,重复上述训练过程即可。
首先,定义一个循环:
# 开始训练
for e in range(epoch):
pass
那么,接着,需要定义调整学习率的方法,这里采取的思路是训练到指定epoch次数,调整一次学习率:
# 开始训练
for e in range(epoch):
# 调整学习率
if e+1 == 50:
optim_G.param_groups[0]['lr'] /= 10
optim_D.param_groups[0]['lr'] /= 10
print('学习率改变')
if e+1 == 100:
optim_G.param_groups[0]['lr'] /= 10
optim_D.param_groups[0]['lr'] /= 10
print('学习率改变')
然后,再定义一个循环,开始定义真正的训练过程:
# 开始训练
for e in range(epoch):
# 调整学习率
......
# 开始真正的训练
for i,(batch_img,batch_label) in enumerate(train_loader):
pass
首先,需要定义两个变量,一个全为1,一个全为0,用于与判别器的输出计算损失值:
# 开始训练
for e in range(epoch):
# 调整学习率
......
# 开始真正的训练
for i,(batch_img,batch_label) in enumerate(train_loader):
# 创建标签张量,一个全为1,一个全为0
batch = batch_img.size()[0]
label_real = torch.ones(batch, requires_grad=True)
label_fake = torch.zeros(batch, requires_grad=True)
然后,把这些值放入GPU中:
# 开始训练
for e in range(epoch):
# 调整学习率
......
# 开始真正的训练
for i,(batch_img,batch_label) in enumerate(train_loader):
# 创建标签张量,一个全为1,一个全为0
batch = batch_img.size()[0]
label_real = torch.ones(batch, requires_grad=True)
label_fake = torch.zeros(batch, requires_grad=True)
# 放入GPU中
batch_img, batch_label, label_fake, label_real = batch_img.to(device), batch_label.to(device), label_fake.to(
device), label_real.to(device)
下面,开始训练判别器D:
# 开始训练
for e in range(epoch):
# 调整学习率
......
# 开始真正的训练
for i,(batch_img,batch_label) in enumerate(train_loader):
# 定义参数
......
# 优化D
optim_D.zero_grad()
# 先喂真实图像
d_result_real = D(batch_img).squeeze() # 把结果拉平
d_real_loss = loss(d_result_real,label_real)
# 再喂假的图像
# 创建噪声向量
z = torch.randn((batch,100),requires_grad=True).view(-1,100,1,1) # 展开为[batch,100,1,1]
z = z.to(device)
g_result = G(z)
d_result_fake = D(g_result).squeeze()
d_fake_loss = loss(d_result_fake,label_fake)
all_loss = d_fake_loss + d_real_loss # 两个损失相加,同时反向传播
all_loss.backward()
optim_D.step()
接着,训练G:
# 开始训练
for e in range(epoch):
# 调整学习率
......
# 开始真正的训练
for i,(batch_img,batch_label) in enumerate(train_loader):
# 定义参数
......
# 优化D
......
# 开始训练G
optim_G.zero_grad()
# 噪声向量
z = torch.randn((batch, 100), requires_grad=True).view(-1, 100, 1, 1) # 展开为[batch,100,1,1]
z = z.to(device)
g_result = G(z)
d_result = D(g_result).squeeze()
g_loss = loss(d_result,label_real)
g_loss.backward()
optim_G.step()
在完成了训练D、G后,打印本次训练的损失函数值:
# 开始训练
for e in range(epoch):
# 调整学习率
......
# 开始真正的训练
for i,(batch_img,batch_label) in enumerate(train_loader):
# 定义参数
......
# 优化D
......
# 开始训练G
......
# 每个batch训练完毕后,打印损失
print('epoch {%d},batch {%d},g_loss:%.5f,d_loss:%.5f' % (e+1,i+1,g_loss.item(),all_loss.item()))
而在完成每epoch的训练后,用生成器生成一张图片并保存下来,用于后期分析结果:
# 开始训练
for e in range(epoch):
......
# 开始真正的训练
for i,(batch_img,batch_label) in enumerate(train_loader):
......
# 训练一个epoch,用生成器生成图片,并保存
z = torch.randn((5*5, 100), requires_grad=True).view(-1, 100, 1, 1) # 展开为[batch,100,1,1]
z = z.to(device)
gen_imgs = G(z)
# 路径自己改
save_image(gen_imgs.data, "fake_images1/%d.png" % (e + 1), nrow=10, normalize=True)
完整的训练代码
# 开始训练
for e in range(epoch):
# 调整学习率
if e+1 == 50:
optim_G.param_groups[0]['lr'] /= 10
optim_D.param_groups[0]['lr'] /= 10
print('学习率改变')
if e+1 == 100:
optim_G.param_groups[0]['lr'] /= 10
optim_D.param_groups[0]['lr'] /= 10
print('学习率改变')
# 开始真正的训练
for i,(batch_img,batch_label) in enumerate(train_loader):
# 创建标签张量,一个全为1,一个全为0
batch = batch_img.size()[0]
label_real = torch.ones(batch, requires_grad=True)
label_fake = torch.zeros(batch, requires_grad=True)
# 放入GPU中
batch_img, batch_label, label_fake, label_real = batch_img.to(device), batch_label.to(device), label_fake.to(
device), label_real.to(device)
# 优化D
optim_D.zero_grad()
# 先喂真实图像
d_result_real = D(batch_img).squeeze() # 把结果拉平
d_real_loss = loss(d_result_real,label_real)
# 再喂假的图像
# 创建噪声向量
z = torch.randn((batch,100),requires_grad=True).view(-1,100,1,1) # 展开为[batch,100,1,1]
z = z.to(device)
g_result = G(z)
d_result_fake = D(g_result).squeeze()
d_fake_loss = loss(d_result_fake,label_fake)
all_loss = d_fake_loss + d_real_loss # 两个损失相加,同时反向传播
all_loss.backward()
optim_D.step()
# 开始训练G
optim_G.zero_grad()
# 噪声向量
z = torch.randn((batch, 100), requires_grad=True).view(-1, 100, 1, 1) # 展开为[batch,100,1,1]
z = z.to(device)
g_result = G(z)
d_result = D(g_result).squeeze()
g_loss = loss(d_result,label_real)
g_loss.backward()
optim_G.step()
# 每个batch训练完毕后,打印损失
print('epoch {%d},batch {%d},g_loss:%.5f,d_loss:%.5f' % (e+1,i+1,g_loss.item(),all_loss.item()))
# 训练一个epoch,用生成器生成图片,并保存
z = torch.randn((5*5, 100), requires_grad=True).view(-1, 100, 1, 1) # 展开为[batch,100,1,1]
z = z.to(device)
gen_imgs = G(z)
# 路径自己改
save_image(gen_imgs.data, "fake_images1/%d.png" % (e + 1), nrow=10, normalize=True)
5. 结果分析
由于网络架构相比CGAN大了很多,因此我的个人电脑跑起来很慢,我计算了下,一个epoch大概需要5-8分钟才跑完,那么跑完目标的200个epoch需要16-26个小时左右。这是我个人没有办法接受的,而且电脑用了几年了,散热很垃圾。
所以,我减少了epoch次数,改为了50个epoch,同时我将网络架构按比例缩小了,最终的运行结果如下图所示:
6. 总结
虽然相比于CGAN的结果来说,DCGAN的结果却是更加好一点,这里的好我个人认为是结构上的胜利,因为DCGAN采取的结构更大,所以取得的结果比较好。即使训练的批次很小,结果也不错。
但是,同样发下一个问题,就是训练仍然不稳定,有时候训练的好好的,但是突然生成器损失值会增大很多,如下图所示:
综上,感觉想取得好结果,一是从结构上入手,二是尽量控制稳定性。
完整代码
# author: baiCai
# DCGAN --- MNIST --- pytorch
# 导包
import torch
from torch import nn
from torch.utils.data import DataLoader
from torch.nn import functional as F
from torch import optim
from torchvision.utils import save_image
from torchvision import transforms
from torchvision.datasets import MNIST
# 说明: 这里的d是后面添加的,用于控制模型的规模
# 生成器
class Generator(nn.Module):
def __init__(self,d=128):
super(Generator, self).__init__()
self.model = nn.Sequential(
# 输入通道、输出通道、卷积核、步长、padding
nn.ConvTranspose2d(100,d*8,4,1,0),
nn.BatchNorm2d(d*8),
nn.ReLU(),
nn.ConvTranspose2d(d*8, d*4, 4, 2, 1),
nn.BatchNorm2d(d*4),
nn.ReLU(),
nn.ConvTranspose2d(d*4, d*2, 4, 2, 1),
nn.BatchNorm2d(d*2),
nn.ReLU(),
nn.ConvTranspose2d(d*2, d, 4, 2, 1),
nn.BatchNorm2d(d),
nn.ReLU(),
# 最后,输出通道数为1,是因为为灰度图
nn.ConvTranspose2d(d, 1, 4, 2, 1),
nn.Tanh(),
)
def forward(self,x):
result = self.model(x)
return result
# 判别器
class Discriminator(nn.Module):
def __init__(self,d=128):
super(Discriminator, self).__init__()
self.model = nn.Sequential(
# 按照图片上的实现即可
nn.Conv2d(1,d,kernel_size=4,stride=2,padding=1),
# 这里0.2是图片中采取的值
nn.LeakyReLU(0.2),
nn.Conv2d(d,d*2,4,2,1),
nn.BatchNorm2d(d*2),
nn.LeakyReLU(0.2),
nn.Conv2d(d*2, d*4, 4, 2, 1),
nn.BatchNorm2d(d*4),
nn.LeakyReLU(0.2),
nn.Conv2d(d*4, d*8, 4, 2, 1),
nn.BatchNorm2d(d*8),
nn.LeakyReLU(0.2),
# 最终输出1不要忘记即可
nn.Conv2d(d*8, 1, 4, 2, 0),
# 加上一个sigmoid,控制输出到0-1
nn.Sigmoid(),
)
def forward(self,x):
result = self.model(x)
return result
# 定义基本参数
# 设备
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
# batch大小、epoch次数、初始学习率、损失函数
batch_size = 128
epoch = 200
lr = 0.002
loss = nn.BCELoss()
# 模型创建
D = Discriminator().to(device)
G = Generator().to(device)
# 优化器
optim_G = optim.Adam(G.parameters(),lr=lr,betas=(0.5, 0.999))
optim_D = optim.Adam(D.parameters(),lr=lr,betas=(0.5, 0.999))
# 预处理方法
transforms_func = transforms.Compose(
[transforms.Resize(64), # 缩放到64*64
transforms.ToTensor(), # 转为tensor
transforms.Normalize(mean=0.5,std=0.5)] # 归一化处理
)
# 加载数据
dataset = MNIST('../data/mnist', # 确定自己加载的路径或者要加载保存的路径
train=True,
transform=transforms_func,
download=False) # 这个参数,根据自己有没有数据集确定
train_loader = DataLoader(dataset,batch_size=batch_size,shuffle=True,drop_last=True)
# 开始训练
for e in range(epoch):
# 调整学习率
if e+1 == 50:
optim_G.param_groups[0]['lr'] /= 10
optim_D.param_groups[0]['lr'] /= 10
print('学习率改变')
if e+1 == 100:
optim_G.param_groups[0]['lr'] /= 10
optim_D.param_groups[0]['lr'] /= 10
print('学习率改变')
# 开始真正的训练
for i,(batch_img,batch_label) in enumerate(train_loader):
# 创建标签张量,一个全为1,一个全为0
batch = batch_img.size()[0]
label_real = torch.ones(batch, requires_grad=True)
label_fake = torch.zeros(batch, requires_grad=True)
# 放入GPU中
batch_img, batch_label, label_fake, label_real = batch_img.to(device), batch_label.to(device), label_fake.to(
device), label_real.to(device)
# 优化D
optim_D.zero_grad()
# 先喂真实图像
d_result_real = D(batch_img).squeeze() # 把结果拉平
d_real_loss = loss(d_result_real,label_real)
# 再喂假的图像
# 创建噪声向量
z = torch.randn((batch,100),requires_grad=True).view(-1,100,1,1) # 展开为[batch,100,1,1]
z = z.to(device)
g_result = G(z)
d_result_fake = D(g_result).squeeze()
d_fake_loss = loss(d_result_fake,label_fake)
all_loss = d_fake_loss + d_real_loss # 两个损失相加,同时反向传播
all_loss.backward()
optim_D.step()
# 开始训练G
optim_G.zero_grad()
# 噪声向量
z = torch.randn((batch, 100), requires_grad=True).view(-1, 100, 1, 1) # 展开为[batch,100,1,1]
z = z.to(device)
g_result = G(z)
d_result = D(g_result).squeeze()
g_loss = loss(d_result,label_real)
g_loss.backward()
optim_G.step()
# 每个batch训练完毕后,打印损失
print('epoch {%d},batch {%d},g_loss:%.5f,d_loss:%.5f' % (e+1,i+1,g_loss.item(),all_loss.item()))
# 训练一个epoch,用生成器生成图片,并保存
# 可以改变5*5,表示生成多数个图片
z = torch.randn((5*5, 100), requires_grad=True).view(-1, 100, 1, 1) # 展开为[batch,100,1,1]
z = z.to(device)
gen_imgs = G(z)
# 路径自己改
save_image(gen_imgs.data, "fake_images1/%d.png" % (e + 1), nrow=10, normalize=True)
# # 保存一下权值:是否保存取决于自己,另外需要注意路径哦
# torch.save(D.state_dict(),'./save_weights/D.pkl')
# torch.save(G.state_dict(),'./save_weights/G.pkl')