基于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')



















