Pytorch CGAN代码
我写的CGAN（Conditional GAN）代码是在DCGAN（Deep Convolution Generative Adversarial Networks）的基础上形成的，DCGAN的优点在于使用了卷积，比全连接层训练时更少的使用数据。因为看了几份CGAN的代码，有的train的时候再Generator时用的卷积，而不是转置卷积，感觉很奇怪，就自己手打了一份使用转置卷积的代码；如有错误，还望指出。
下面的图是训练过程：
from __future__ import print_function import torchvision import argparse import os import random import torch import torch.nn as nn import torch.nn.parallel import torch.backends.cudnn as cudnn from torch.autograd import Variable import torch.optim as optim import torch.utils.data import torchvision.datasets as dset import torchvision.transforms as transforms import torchvision.utils as vutils import numpy as np import matplotlib.pyplot as plt import matplotlib.animation as animation from IPython.display import HTML %matplotlib inline
定义一些参数：
workers = 2 # 数据加载器能够使用的进程数量 batch_size = 256 # 训练时的批大小 # 训练图片的大小，所有的图片都将改变到该大小 image_size = 32 # 转换器使用的大小. nc = 1 # 训练图片的通道数，彩色图片是3 nz = 100 # 本征向量z的大小(生成器的输入大小) ngf = 32 # 生成器中特征图个数大小 ndf = 32 # 判别器中特征图个数大小 num_epochs = 5 # 训练次数 lr = 0.0002 # 优化器学习率 beta1 = 0.5 # Adam优化器的Beta1超参 ngpu = 1 # 可利用的GPU数量，使用0将运行在CPU模式。
创建数据集：
# 创建数据集 dataset = torchvision.datasets.MNIST(root='./dataset_method_1', train=True, transform=transforms.Compose([transforms.Resize(image_size), transforms.CenterCrop(image_size), transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,)), ]), download=True) # 创建数据加载器 dataloader = torch.utils.data.DataLoader(dataset=dataset, batch_size=batch_size, shuffle=True) #600*100*([[28*28],x])
# 决定我们在哪个设备上运行 device = torch.device("cuda:0" if (torch.cuda.is_available() and ngpu > 0) else "cpu")
展示一些训练图片：
real_batch = next(iter(dataloader)) plt.figure(figsize=(8,8)) plt.axis("off") plt.title("Training Images") plt.imshow(np.transpose(vutils.make_grid(real_batch[0].to(device)[:64], padding=2, normalize=True).cpu(),(1,2,0)))
在netG和netD上调用的自定义权重初始化函数：
def weights_init(m): classname = m.__class__.__name__ if classname.find('Conv') != -1: nn.init.normal_(m.weight.data, 0.0, 0.02) elif classname.find('BatchNorm') != -1: nn.init.normal_(m.weight.data, 1.0, 0.02) nn.init.constant_(m.bias.data, 0)
生成器model（Generator）：
class Generator(nn.Module): def __init__(self, ngpu): self.ngpu = ngpu super(Generator,self).__init__() self.gen=nn.Sequential( # 输入是 Z, 对Z进行卷积 nn.ConvTranspose2d(110, ngf*8, 4, 1, 0, bias=False), nn.BatchNorm2d(ngf*8), nn.ReLU(inplace=True), # 输入特征图大小. (ngf*8) x 4 x 4 nn.ConvTranspose2d(ngf*8, ngf*4, 4, 2, 1, bias=False), nn.BatchNorm2d(ngf*4), nn.ReLU(inplace=True), # 输入特征图大小. (ngf*4) x 8 x 8 nn.ConvTranspose2d(ngf*4, ngf*2, 4, 2, 1, bias=False), nn.BatchNorm2d(ngf * 2), nn.ReLU(inplace=True), # 输入特征图大小. (ngf*2) x 16 x 16 nn.ConvTranspose2d(ngf*2, 1, 4, 2, 1, bias=False), nn.Tanh() # 输入特征图大小. (nc) x 32 x 32 ) def forward(self, x): x=self.gen(x) return x
# 创建生成器 netG = Generator(ngpu).to(device) # 如果期望使用多个GPU，设置一下。 if (device.type == 'cuda') and (ngpu > 1): netG = nn.DataParallel(netG, list(range(ngpu))) # 使用权重初始化函数 weights_init 去随机初始化所有权重 # mean=0, stdev=0.2. netG.apply(weights_init) # 输出该模型 print(netG)
判别器model（Discriminator）：
class Discriminator(nn.Module): def __init__(self, ngpu): super(Discriminator, self).__init__() self.ngpu = ngpu self.main = nn.Sequential( # state size. (ndf) x 32 x 32 nn.Conv2d(1, ngf, 4, 2, 1, bias=False), nn.LeakyReLU(0.2, inplace=True), # 输入大小. (ndf*2) x 16 x 16 nn.Conv2d(ngf, ngf * 2, 4, 2, 1, bias=False), nn.BatchNorm2d(ngf * 2), nn.LeakyReLU(0.2, inplace=True), # 输入大小. (ndf*2) x 8 x 8 nn.Conv2d(ngf * 2, ngf * 4, 4, 2, 1, bias=False), nn.BatchNorm2d(ngf * 4), nn.LeakyReLU(0.2, inplace=True), # 输入大小. (ndf*4) x 4 x 4 nn.Conv2d(ngf * 4, 11, 4, 1, 0, bias=False), nn.Sigmoid() ) def forward(self, x): return self.main(x)
# 创建判别器 netD = Discriminator(ngpu).to(device) # 如果期望使用多GPU，设置一下 if (device.type == 'cuda') and (ngpu > 1): netD = nn.DataParallel(netD, list(range(ngpu))) # 使用权重初始化函数 weights_init 去随机初始化所有权重 # mean=0, stdev=0.2. netD.apply(weights_init) # 输出该模型 print(netD)
定义一些Loss，optimizer
# 初始化 BCE损失函数 criterion = nn.BCELoss() criterion_pic = nn.CrossEntropyLoss() # 创建一个批次的本征向量用于可视化生成器训练的过程。 fixed_noise = torch.randn(64, nz + 10, 1, 1, device=device) # 建立一个在训练中使用的真实和假的标记 real_label = 1 fake_label = 0 # 为G和D都设置Adam优化器 # 刚看到说G的lr是D的lr的2倍会好一些 optimizerD = optim.Adam(netD.parameters(), lr=lr, betas=(beta1, 0.999)) optimizerG = optim.Adam(netG.parameters(), lr=lr, betas=(beta1, 0.999))
主要改的就是这里了：
# 训练循环 # 保存跟踪进度的列表 img_list = [] G_losses = [] D_losses = [] iters = 0 print("Starting Training Loop...") # 每个epoh for epoch in range(num_epochs): # 数据加载器中的每个批次 for i, data in enumerate(dataloader, 0): ############################ # (1) 更新 D 网络: 最大化 log(D(x)) + log(1 - D(G(z))) ########################### ## 使用所有真实样本批次训练 netD.zero_grad() # 格式化批 real_cpu = data[0].to(device) b_size = real_cpu.size(0) label = torch.full((b_size,), real_label, device=device) # 通过D向前传递真实批次 output = netD(real_cpu) output = output.view([b_size, 11]
real_label_label = output[:, 0] # 第一维就是预测是否是真实图片 real_label_pic = output[:, 1:] # 后十维是onehot，判断是哪个数字 # print(real_label_pic.shape) # 对所有真实样本批次计算损失 errD_real = criterion(real_label_label, label) errD_real.backward(retain_graph=True) # 连续多个backward，除了最后一个不叫retain_graph=True,其余都要加 errD_fake_pic = criterion_pic(real_label_pic, data[1].cuda(device)) errD_fake_pic.backward() # 计算后向传递中D的梯度 D_x = real_label_label.mean().item() ## 使用所有假样本批次训练 # 生成本征向量批次 # noise = torch.randn(b_size, nz+10, 1, 1, device=device) # 注释这行表示输入110维的noise，也就是DCGAN的输入 noise = torch.randn(b_size, nz, device=device) labels_onehot = np.zeros((b_size,10)) labels_onehot[np.arange(b_size),data[1].numpy()]=1 noise=np.concatenate((noise.cpu().numpy(), labels_onehot),axis=1) # 每张图片是0到9之间的数字 labels_onehot = Variable(torch.from_numpy(labels_onehot).float()).to(device) noise=noise.reshape([-1, 110, 1, 1]) noise=Variable(torch.from_numpy(noise).float()).to(device) # 使用生成器G生成假图片 fake = netG(noise) label.fill_(fake_label) # 使用判别器分类所有的假批次样本 output = netD(fake.detach()) # 计算判别器D的损失对所有的假样本批次 output = output.view([b_size, 11]) real_label_label = output[:, 0] real_label_pic = output[:, 1:] errD_fake = criterion(real_label_label, label) # 对这个批次计算梯度 errD_fake.backward(retain_graph=True) errD_fake_pic = criterion_pic(real_label_pic, data[1].cuda(device)) errD_fake_pic.backward() D_G_z1 = real_label_label.mean().item() # 把所有真样本和假样本批次的梯度加起来 errD = errD_real + errD_fake + errD_fake_pic # 更新判别器D optimizerD.step() ############################ # (2) 更新 G 网络: 最大化 log(D(G(z))) ########################### netG.zero_grad() label.fill_(real_label) # 假样本的标签对于生成器成本是真的 # 因为我们之更新了D，通过D执行所有假样本批次的正向传递 output = netD(fake).view(-1) # 基于这个输出计算G的损失 output = output.view([b_size, 11]) real_label_label = output[:, 0] real_label_pic = output[:, 1:] errG = criterion(real_label_label, label) # 为生成器计算梯度 errG.backward(retain_graph=True) errD_fake_pic = criterion_pic(real_label_pic, data[1].cuda(device)) errD_fake_pic.backward() D_G_z2 = output.mean().item() # 更新生成器G optimizerG.step() # 输出训练状态 if i % 50 == 0: print('[%d/%d][%d/%d]\tLoss_D: %.4f\tLoss_G: %.4f\tD(x): %.4f\tD(G(z)): %.4f / %.4f' % (epoch, num_epochs, i, len(dataloader), errD.item(), errG.item(), D_x, D_G_z1, D_G_z2)) # 为以后画损失图，保存损失 G_losses.append(errG.item()) D_losses.append(errD.item()) # 检查生成器generator做了什么，通过保存的fixed_noise通过G的输出 if (iters % 500 == 0) or ((epoch == num_epochs-1) and (i == len(dataloader)-1)): with torch.no_grad(): fake = netG(fixed_noise).detach().cpu() img_list.append(vutils.make_grid(fake, padding=2, normalize=True)) iters += 1
Loss图：
plt.figure(figsize=(10,5)) plt.title("Generator and Discriminator Loss During Training") plt.plot(G_losses,label="G") plt.plot(D_losses,label="D") plt.xlabel("iterations") plt.ylabel("Loss") plt.legend() plt.show()
查看效果：
# 从数据加载器中获取一批真实图像 real_batch = next(iter(dataloader)) # 画出真实图像 plt.figure(figsize=(15,15)) plt.subplot(1,2,1) plt.axis("off") plt.title("Real Images") plt.imshow(np.transpose(vutils.make_grid(real_batch[0].to(device)[:64], padding=5, normalize=True).cpu(),(1,2,0))) # 画出来自最后一次训练的假图像 plt.subplot(1,2,2) plt.axis("off") plt.title("Fake Images") plt.imshow(np.transpose(img_list[-1],(1,2,0))) plt.show()
输入某一个数字时：
noise = torch.randn(b_size, nz, device=device) labels_onehot = np.zeros((b_size,10)) labels_onehot[np.arange(b_size),9]=1 # 这个9换成想生成的数字 noise=np.concatenate((noise.cpu().numpy(), labels_onehot),axis=1) # 每张图片是0到9之间的数字 labels_onehot = Variable(torch.from_numpy(labels_onehot).float()).to(device) noise=noise.reshape([-1, 110, 1, 1]) noise=Variable(torch.from_numpy(noise).float()).to(device) # 使用生成器G生成假图片 fake = netG(noise)
plt.imshow(fake.cpu().detach().numpy()[5][0])
生成的效果图：
