pytorch实战10：基于pytorch简单实现CGAN

基于pytorch简单实现CGAN

前言

最近在看经典的卷积网络架构，打算自己尝试复现一下，在此系列文章中，会参考很多文章，有些已经忘记了出处，所以就不贴链接了，希望大家理解。

完整的代码在最后。

本系列必须的基础

python基础知识、CNN原理知识、pytorch基础知识

本系列的目的

一是帮助自己巩固知识点；

二是自己实现一次，可以发现很多之前的不足；

三是希望可以给大家一个参考。

目录结构

文章目录

- 基于pytorch简单实现CGAN
- - 1. 前言：
  - 2. 数据介绍与数据加载器：
  - 3. 生成器：
  - 4. 判别器：
  - 5. 训练：
  - 6. 训练结果展示与思考：
  - 7. 预测：
  - 8. 总结：

1. 前言：

今天来实现一下‘干’系列中的CGAN，相比于原始的GAN来说，CGAN更加稳定，且引入了条件参数，可以在一定程度上控制输出。

本次实现CGAN，使用pytorch库和手写数字数据集MNIST（注定了很简单的）。另外，本次的目的是简单实现一下，更深入的优化，就靠大家了。

最后，补充一下我的项目结构：（特别注意：我的数据集文件夹在其它文件夹内）

├─fake_images
|  └─ 用于保存生成器生成的图片的文件夹
├─network_files
│  └─ 用于存放网络结构的文件夹
├─save_weights
|  └─ 用于存放训练好权重的文件夹
|_ 数据加载器、训练、测试文件

2. 数据介绍与数据加载器：

数据介绍

MNIST数据集大家应该都不陌生吧。我这里简单说明一下，这个数据集是灰度图（通道数为1），所有图片大小为28*28，内容为0-9的数字。

在这里插入图片描述

数据加载器

这里不需要像前面一样，实现复杂的数据加载器。因为官方已经为我们封装好了现成的加载器，这里我们只是将之封装到文件的函数中即可，如下：

from torchvision import transforms as T
from torchvision.datasets import MNIST

# 这个简单，MNIST数据集很经典，因此可以直接调用官方的方法
def My_Dataset(train=True):
    # 定义预处理方法
    transforms = T.Compose([
        T.Resize(28), # 这个其实不需要，只是后期如果要改为其它图片，肯行需要修改
        T.ToTensor(),
        T.Normalize([0.5], [0.5])
    ])
    if train:
        #  路径需要修改为自己的
        # 另外，这个数据集会自动下载
        data = MNIST('../data/mnist',train=True,download=True,transform=transforms)
    else:
        data = MNIST('../data/mnist', train=False, download=True, transform=transforms)
    return data

3. 生成器：

生成器很好定义，只需要注意最后的输出为图像尺寸即可。关于生成器的具体结构，其实论文中并没有给出，但是网上有别人实现的架构，我们可以拿来用：（下图来自GitHub项目pytorch-MNIST-CelebA-cGAN-cDCGAN）

在这里插入图片描述

其中，上图中G输出784，是因为MNIST图像大小为28*28=784。另外，上图中z是输入的噪声向量，y为其对应的标签。

按照上图架构，可以实现生成器：

import torch
from torch import nn
import numpy as np

class Generator(nn.Module):
    def __init__(self,in_dim,img_shape,num_classes=10):
        '''
        :param in_dim: 噪声的维度
        :param num_classes: 类别个数，为数字0-9，即10个类别
        :param img_shape: 生成的图片尺寸，一般与真实图片一致，即28*28
        '''
        super(Generator, self).__init__()
        # 初始化
        self.in_dim = in_dim
        self.num_classes = num_classes
        self.img_shape = img_shape
        self.label_embedding = nn.Embedding(self.num_classes,self.num_classes)

        # 定义模型
        self.model = nn.Sequential(
            # 不要忘记输入是真实标签+噪声
            self.block(self.in_dim+self.num_classes,128,normalize=False),
            self.block(128,256),
            self.block(256,512),
            self.block(512,1024),
            #np.prod 输入a是数组，返回指定轴上的乘积，不指定轴默认是所有元素的乘积
            nn.Linear(1024,int(np.prod(self.img_shape))),
            nn.Tanh(),
        )

    # 构建基础块
    def block(self,in_channels,out_channels,normalize=True):
        layers = []
        layers.append(nn.Linear(in_channels,out_channels))
        # 是否初始化
        if normalize:
            layers.append(nn.BatchNorm1d(out_channels))
        # 不要忘记加上激活函数
        layers.append(nn.LeakyReLU(0.2,inplace=True))
        return nn.Sequential(*layers)

    def forward(self,noise,labels):
        # 将输入标签和噪声拼接在一起
        # temp2 = labels # 256
        # temp = self.label_embedding(labels) # torch.Size([256, 10])
        # 将标签转为了向量，目的是为和噪声进行拼接
        input_x = torch.cat((self.label_embedding(labels),noise),-1) # torch.Size([256, 110])
        # 送入模型
        output_img = self.model(input_x)
        # 将输出构造成真实图片的shape信息【batch，w，h】
        output_img = output_img.view(output_img.size(0),*(self.img_shape))
        return output_img

4. 判别器：

同样根据上图实现判别器，只是注意最后的输出为1个通道，这是因为我们同时输入了图像和对应标签，最后就直接输出一个概率值，这个概率值越接近1，表示判别其为真实图像，越接近0表示假图像。

同样，根据上图写出判别器架构，具体可以看注释：

import numpy as np
import torch.nn as nn
import torch

class Discriminator(nn.Module):
    def __init__(self,img_shape,num_classes=10):
        '''
        :param img_shape: 图像尺寸，这里为28*28
        :param num_classes: 类别数，即10类
        '''
        super(Discriminator, self).__init__()
        # 初始化一下
        self.num_classes = num_classes
        self.img_shape = img_shape
        self.label_embedding = nn.Embedding(self.num_classes,self.num_classes)
        # 定义模型
        self.model = nn.Sequential(
            nn.Linear(self.num_classes+int(np.prod(self.img_shape)),512),
            nn.LeakyReLU(0.2,inplace=True),
            nn.Dropout(0.5),
            nn.Linear(512,512),
            nn.Dropout(0.5),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(512,512),
            nn.Dropout(0.5),
            nn.LeakyReLU(0.2, inplace=True),
            # 最后输出为1，表示分类
            nn.Linear(512,1)
        )

    def forward(self,img,labels):
        input_x = torch.cat((img.view(img.size(0), -1), self.label_embedding(labels)),-1)
        output_class = self.model(input_x)
        return output_class

5. 训练：

相比于网络结构的定义，CGAN如何训练的，更重要。

首先，导入基本的包和我们自己定义的数据加载器、网络结构等：

import os
from My_Dataset import My_Dataset
from network_files.Discriminator import Discriminator as D
from network_files.Generator import Generator as G

import numpy as np
import torch
from torch import nn
from torch.utils.data import DataLoader
from torch.autograd import Variable
from torch import optim
from torchvision.utils import save_image # 用于生成图片、保持图片

接着，定义基本的参数：

# 定义基本参数
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') # 设备
batch_size = 256 # batch大小
epoch = 400 # 训练批次
lr = 0.002 # 初始学习率
img_shape = (1,28,28) # 图像shape
num_classes = 10 # 类别数目
in_dim = 100 # 噪声维度

然后，创建模型、定义优化器（Adam）、损失函数（MSE）并加载数据，这都是基本操作：

# 创建模型
Generator = G(in_dim,img_shape,num_classes)
Discriminator = D(img_shape,num_classes)
Generator.to(device)
Discriminator.to(device)
# 定义损失函数
loss = nn.MSELoss()
# 优化器
optim_G = optim.Adam(Generator.parameters(),lr=lr,betas=(0.5, 0.999))
optim_D = optim.Adam(Discriminator.parameters(),lr=lr,betas=(0.5, 0.999))
# 加载数据
train_dataset = My_Dataset(train=True)
train_loader = DataLoader(train_dataset,batch_size=batch_size,shuffle=True,drop_last=True)

接着开始训练：

# 开始训练
for e in range(epoch):

首先，训练到一定程度时调整学习率，这个方法也是我最近看到的，感觉可以很方便的实现调整学习率：

# 开始训练
for e in range(epoch):
    #学习率调整
    if (e + 1) == 100:
        optim_D.param_groups[0]['lr'] /= 10
        optim_G.param_groups[0]['lr'] /= 10
        print("learning rate change!")

        if (e + 1) == 250:
            optim_D.param_groups[0]['lr'] /= 10
            optim_G.param_groups[0]['lr'] /= 10
            print("learning rate change!")

然后，开始训练第一个epoch，首先把数据集放入GPU中：

# 开始训练
for e in range(epoch):
    #学习率调整
    ...... # 这里不重复了，直接省略  
    for i,(batch_img,batch_label) in enumerate(train_loader):
        # 放入设备中
        batch_img,batch_label = batch_img.to(device),batch_label.to(device)

接着，需要创建噪声向量和其对应的标签，并创建两个全为0和全为1的向量，用于后期与判别器输出比较（因为判别器只输出一个概率值，因此需要创建这两个向量）：

# 开始训练
for e in range(epoch):
    #学习率调整
    ...... # 这里不重复了，直接省略  
    for i,(batch_img,batch_label) in enumerate(train_loader):
        # 放入设备中
        batch_img,batch_label = batch_img.to(device),batch_label.to(device)
		# 创建噪声向量和标签
        # z = [batch,100] , labels = [64] 为要生成的数字图片标签
        z = Variable(torch.cuda.FloatTensor(np.random.normal(0, 1, (batch_size, in_dim))))  # 噪声样本
        g_labels = Variable(torch.cuda.LongTensor(np.random.randint(0, num_classes, batch_size)))  # 对应的标签
        #  用于计算损失的标签，一个全为1，一个全为0
        valid = Variable(torch.cuda.FloatTensor(batch_size, 1).fill_(1.0), requires_grad=False)  # torch.Size([64, 1])
        fake = Variable(torch.cuda.FloatTensor(batch_size, 1).fill_(0.0), requires_grad=False)  # torch.Size([64, 1])

然后，开始真正的训练，训练的思路如下：

在这里插入图片描述

按照上述思路，实现训练过程：

# 开始训练
for e in range(epoch):
    #学习率调整
    ...... # 这里不重复了，直接省略  
    for i,(batch_img,batch_label) in enumerate(train_loader):
        ...... # 前面那些基础的定义
        # 生成器根据噪声和标签生成假的图片
        optim_G.zero_grad()
        g_img = Generator(z,g_labels)
        # 判别器判别，输出的为[batch,1]，输出的是概率值，越接近1，表示生成的图片越真实
        g_class = Discriminator(g_img,g_labels)
        # 损失
        g_loss = loss(g_class,valid)
        g_loss.backward()
        optim_G.step()
        # 训练判别器
        optim_D.zero_grad()
        # 实图片的损失，目的就是希望其与真实标签1的损失最小
        d_real_class = Discriminator(batch_img,batch_label)
        d_real_loss = loss(d_real_class,valid)
        # g_img.detach() 不要忘记detach，不然会报错的，因为前面backward已经释放了
        # 生成的虚假图片的损失，目的就是希望其与虚假标签0的损失最小
        d_fake_class = Discriminator(g_img.detach(),g_labels)
        d_fake_loss = loss(d_fake_class,fake)
        all_loss = (d_fake_loss+d_real_loss)/2
        all_loss.backward()
        optim_D.step()
        # 每个step打印损失
        print('epoch{%d},batch{%d},D_loss: %.5f,G_loss: %.5f' % (e+1,i+1,all_loss.item(),g_loss.item()))

为了能够看到训练的效果，在每经历一个epoch后，都用生成器去生成图像，并保持到指定路径（需要自己修改）：

# 开始训练
for e in range(epoch):
    #学习率调整
    ...... # 这里不重复了，直接省略  
    for i,(batch_img,batch_label) in enumerate(train_loader):
		...... # 整个训练过程省略
	# 每一个epoch，看看生成器生成的图片咋样了
    z = Variable(torch.cuda.FloatTensor(np.random.normal(0, 1, (10 ** 2, 100))))
    labels = np.array([num for _ in range(10) for num in range(10)])
    labels = Variable(torch.cuda.LongTensor(labels))
    gen_imgs = Generator(z, labels)
    # 路径自己改
    save_image(gen_imgs.data, "fake_images1/%d.png" % (e+1), nrow=10, normalize=True)

完成所有的训练后，保持一下权重，方便后期用：

# 保存权重
torch.save(Generator.state_dict(),'./save_weights/G_params.pkl')
torch.save(Discriminator.state_dict(),'./save_weights/D_params.pkl')

完整训练代码

def main():
    # 定义基本参数
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') # 设备
    batch_size = 256 # batch大小
    epoch = 400 # 训练批次
    lr = 0.002 # 初始学习率
    img_shape = (1,28,28) # 图像shape
    num_classes = 10 # 类别数目
    in_dim = 100 # 噪声维度
    # 创建模型
    Generator = G(in_dim,img_shape,num_classes)
    Discriminator = D(img_shape,num_classes)
    Generator.to(device)
    Discriminator.to(device)
    # 定义损失函数
    loss = nn.MSELoss()
    # 优化器
    optim_G = optim.Adam(Generator.parameters(),lr=lr,betas=(0.5, 0.999))
    optim_D = optim.Adam(Discriminator.parameters(),lr=lr,betas=(0.5, 0.999))
    # optim_G = optim.SGD(Generator.parameters(),lr=lr,momentum=0.9)
    # optim_D = optim.SGD(Discriminator.parameters(),lr=lr,momentum=0.9)
    # 加载数据
    train_dataset = My_Dataset(train=True)
    train_loader = DataLoader(train_dataset,batch_size=batch_size,shuffle=True,drop_last=True)

    # 开始训练
    for e in range(epoch):
        #学习率调整
        if (e + 1) == 100:
            optim_D.param_groups[0]['lr'] /= 10
            optim_G.param_groups[0]['lr'] /= 10
            print("learning rate change!")

        if (e + 1) == 250:
            optim_D.param_groups[0]['lr'] /= 10
            optim_G.param_groups[0]['lr'] /= 10
            print("learning rate change!")
        for i,(batch_img,batch_label) in enumerate(train_loader):
            # 放入设备中
            batch_img,batch_label = batch_img.to(device),batch_label.to(device)
            # 创建噪声向量和标签
            # z = [batch,100] , labels = [64] 为要生成的数字图片标签
            z = Variable(torch.cuda.FloatTensor(np.random.normal(0, 1, (batch_size, in_dim))))  # 噪声样本
            g_labels = Variable(torch.cuda.LongTensor(np.random.randint(0, num_classes, batch_size)))  # 对应的标签
            #  用于计算损失的标签，一个全为1，一个全为0
            valid = Variable(torch.cuda.FloatTensor(batch_size, 1).fill_(1.0), requires_grad=False)  # torch.Size([64, 1])
            fake = Variable(torch.cuda.FloatTensor(batch_size, 1).fill_(0.0), requires_grad=False)  # torch.Size([64, 1])
            # 生成器根据噪声和标签生成假的图片
            optim_G.zero_grad()
            g_img = Generator(z,g_labels)
            # 判别器判别，输出的为[batch,1]，输出的是概率值，越接近1，表示生成的图片越真实
            g_class = Discriminator(g_img,g_labels)
            # 损失
            g_loss = loss(g_class,valid)
            g_loss.backward()
            optim_G.step()
            # 训练判别器
            optim_D.zero_grad()
            # 实图片的损失，目的就是希望其与真实标签1的损失最小
            d_real_class = Discriminator(batch_img,batch_label)
            d_real_loss = loss(d_real_class,valid)
            # g_img.detach() 不要忘记detach，不然会报错的，因为前面backward已经释放了
            # 生成的虚假图片的损失，目的就是希望其与虚假标签0的损失最小
            d_fake_class = Discriminator(g_img.detach(),g_labels)
            d_fake_loss = loss(d_fake_class,fake)
            all_loss = (d_fake_loss+d_real_loss)/2
            all_loss.backward()
            optim_D.step()
            # 每个step打印损失
            print('epoch{%d},batch{%d},D_loss: %.5f,G_loss: %.5f' % (e+1,i+1,all_loss.item(),g_loss.item()))
        # 每一个epoch，看看生成器生成的图片咋样了
        z = Variable(torch.cuda.FloatTensor(np.random.normal(0, 1, (10 ** 2, 100))))
        labels = np.array([num for _ in range(10) for num in range(10)])
        labels = Variable(torch.cuda.LongTensor(labels))
        gen_imgs = Generator(z, labels)
        # 路径自己改
        save_image(gen_imgs.data, "fake_images1/%d.png" % (e+1), nrow=10, normalize=True)
    # 保存权重
    torch.save(Generator.state_dict(),'./save_weights/G_params.pkl')
    torch.save(Discriminator.state_dict(),'./save_weights/D_params.pkl')

6. 训练结果展示与思考：

我尝试了训练400次，运行中间的结果展示如下：

在这里插入图片描述

其实运行的效果，我自己不是很满意。另外，在运行过程中，发现一个非常重要的事情：训练经常容易崩溃，即我运行到50epoch的时候，效果已经不错了，但是会突然损失值增大到一个难以接受的值，导致训练崩溃。

关于这一点，应该就是CGAN、GAN的缺点，就是训练不稳定，相比于原始GAN，CGAN的优点是可以控制输出长什么样，但是稳定性方面仍然有所欠缺。我尝试过换一种训练方式，比如先训练好判别器，再训练生成器，试图让优秀学生（判别器）带动差生（生成器），但是效果不佳。

因此，后期我会尝试换一种更稳定的GAN网络，看看是否仍然会遇到同样的问题。

7. 预测：

这个非常简单，由于训练完毕后，我们保存了模型的参数，因此可以直接把生成器参数拿来用。这里的预测，指的是生成一个张图片。

代码如下：

import torch
import numpy as np
from torch.autograd import Variable
from torchvision.utils import save_image

from network_files.Generator import Generator as G

img_shape = (1,32,32)
num_classes = 10
in_dim = 100

# 创建生成器模型
G_model = G(in_dim,img_shape,num_classes)
# 加载参数
G_model.load_state_dict(torch.load('./save_weights/G_params.pkl'))
# 随机生成100个图片并显示
z = Variable(torch.FloatTensor(np.random.normal(0, 1, (10 ** 2, 100))))
labels = np.array([num for _ in range(10) for num in range(10)])
labels = Variable(torch.LongTensor(labels))
gen_imgs = G_model(z, labels)
save_image(gen_imgs.data, "fake_images/%d.png" % (100), nrow=10, normalize=True)

8. 总结：

完成了CGAN，感觉生成对抗的神奇和有趣，进一步了解损失函数控制生成目标的含义。

不过缺点是CGAN生成不稳定且效果并不是很好，后期可以考虑试试其它的GAN网络，看看效果是不是更好点。

参考资料

博客：
https://blog.csdn.net/qq_42721935/article/details/126600071
https://blog.csdn.net/weixin_38052918/article/details/107911739
GitHub项目：
https://github.com/znxlwm/pytorch-MNIST-CelebA-cGAN-cDCGAN

完整代码

链接：https://pan.baidu.com/s/14ZXgMaEuXgM2ETrbHBvOHg 
提取码：izol