我寫的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])
生成的效果圖：
