pytorch GAN

from __future__ import print_function
#%matplotlib inline
import argparse
import os
import random
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
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

# 爲再現性設置隨機seem
manualSeed = 999
#manualSeed = random.randint(1, 10000) # 如果你想要新的結果就是要這段代碼
print("Random Seed: ", manualSeed)
random.seed(manualSeed)
torch.manual_seed(manualSeed)

dataroot = "data/celeba"
workers = 2  #加載數據的工作線程數
batch_size = 128
image_size = 64 #訓練圖像的空間大小。所有圖像將使用變壓器調整爲此大小。
nc = 3 #訓練圖像中的通道數。對於彩色圖像，這是3
nz = 100 #潛在向量 z 的大小(例如： 生成器輸入的大小)
ngf = 64 #生成器中特徵圖的大小
ndf = 64 #判別器中的特徵映射的大小
num_epochs = 5 #訓練epochs的大小
lr = 0.0002 #優化器的學習速率
beta1 = 0.5 #適用於Adam優化器的Beta1超級參數
ngpu = 1 #可用的GPU數量。使用0表示CPU模式。

#### 數據
# 在本教程中，我們將使用[Celeb-A Faces](http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html)數據集，該數據集可以在鏈接或Google Drive中下載。
# 數據集將下載爲名爲*img_align_celeba.zip*的文件。下載後，創建名爲*celeba*的目錄並將zip文件解壓縮到該目錄中。然後，將此筆記中
# 的數據對象輸入設置爲剛剛創建的celeba目錄。生成的目錄結構應該是：
# ```buildoutcfg
# /path/to/celeba
#     -> img_align_celeba
#         -> 188242.jpg
#         -> 173822.jpg
#         -> 284702.jpg
#         -> 537394.jpg

# 這是一個重要的步驟，因爲我們將使用ImageFolder數據集類，它要求在數據集的根文件夾中有子目錄。現在，
# 我們可以創建數據集，創 建數據加載器，設置要運行的設備，以及最後可視化一些訓練數據。
# 我們可以按照設置的方式使用圖像文件夾數據集。
# 創建數據集
dataset = dset.ImageFolder(root=dataroot,
                           transform=transforms.Compose([
                               transforms.Resize(image_size),
                               transforms.CenterCrop(image_size),
                               transforms.ToTensor(),
                               transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
                           ]))
# 創建加載器
dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size,
                                         shuffle=True, num_workers=workers)

# 選擇我們運行在上面的設備
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)))

# 實現
# 通過設置輸入參數和準備好的數據集，我們現在可以進入真正的實現步驟。我們將從權重初始化策略開始，
# 然後詳細討論生成器，鑑別器， 損失函數和訓練循環。

# 權重初始化
# 在DCGAN論文中，作者指出所有模型權重應從正態分佈中隨機初始化，mean = 0，stdev = 0.02。
# weights_init函數將初始化模型作爲 輸入，並重新初始化所有卷積，卷積轉置和batch標準化層以滿足此標準。
# 初始化後立即將此函數應用於模型。
# custom weights initialization called on netG and 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)
        
# 生成器
# 生成器代碼
class Generator(nn.Module):
    def __init__(self, ngpu):
        super(Generator, self).__init__()
        self.ngpu = ngpu
        self.main = nn.Sequential(
            # 輸入是Z，進入卷積
            nn.ConvTranspose2d( nz, ngf * 8, 4, 1, 0, bias=False),
            nn.BatchNorm2d(ngf * 8),
            nn.ReLU(True),
            # state size. (ngf*8) x 4 x 4
            nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 4),
            nn.ReLU(True),
            # state size. (ngf*4) x 8 x 8
            nn.ConvTranspose2d( ngf * 4, ngf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 2),
            nn.ReLU(True),
            # state size. (ngf*2) x 16 x 16
            nn.ConvTranspose2d( ngf * 2, ngf, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf),
            nn.ReLU(True),
            # state size. (ngf) x 32 x 32
            nn.ConvTranspose2d( ngf, nc, 4, 2, 1, bias=False),
            nn.Tanh()
            # state size. (nc) x 64 x 64
        )

    def forward(self, input):
        return self.main(input)

# 現在，可以實例化生成器並應用weights_init函數。查看打印的模型以查看生成器對象的結構。
# 創建生成器
netG = Generator(ngpu).to(device)
if (device.type == 'cuda') and (ngpu > 1): netG = nn.DataParallel(netG, list(range(ngpu))) # 如果需要，管理multi-gpu
netG.apply(weights_init) #應用weights_init函數隨機初始化所有權重，mean= 0，stdev = 0.2。
print(netG) #打印模型

# 判別器
class Discriminator(nn.Module):
    def __init__(self, ngpu):
        super(Discriminator, self).__init__()
        self.ngpu = ngpu
        self.main = nn.Sequential(
            # input is (nc) x 64 x 64
            nn.Conv2d(nc, ndf, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf) x 32 x 32
            nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 2),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf*2) x 16 x 16
            nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 4),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf*4) x 8 x 8
            nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 8),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf*8) x 4 x 4
            nn.Conv2d(ndf * 8, 1, 4, 1, 0, bias=False),
            nn.Sigmoid()
        )

    def forward(self, input):
        return self.main(input)

# 現在，與生成器一樣，我們可以創建判別器，應用weights_init函數，並打印模型的結構。
# 創建判別器
netD = Discriminator(ngpu).to(device)
# Handle multi-gpu if desired
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)

# 損失函數和優化器
# 初始化BCELoss函數
criterion = nn.BCELoss()
# 創建一批潛在的向量，我們將用它來可視化生成器的進程
fixed_noise = torch.randn(64, nz, 1, 1, device=device)
# 在訓練期間建立真假標籤的慣例
real_label = 1
fake_label = 0
# 爲 G 和 D 設置 Adam 優化器
optimizerD = optim.Adam(netD.parameters(), lr=lr, betas=(beta1, 0.999))
optimizerG = optim.Adam(netG.parameters(), lr=lr, betas=(beta1, 0.999))

# 訓練
# Training Loop
# Lists to keep track of progress
img_list = []
G_losses = []
D_losses = []
iters = 0

print("Starting Training Loop...")
# For each epoch
for epoch in range(num_epochs):
    # 對於數據加載器中的每個batch
    for i, data in enumerate(dataloader, 0):

        ############################
        # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))
        ###########################
        ## Train with all-real batch
        netD.zero_grad()
        # Format batch
        real_cpu = data[0].to(device)
        b_size = real_cpu.size(0)
        label = torch.full((b_size,), real_label, device=device)
        # Forward pass real batch through D
        output = netD(real_cpu).view(-1)
        # Calculate loss on all-real batch
        errD_real = criterion(output, label)
        # Calculate gradients for D in backward pass
        errD_real.backward()
        D_x = output.mean().item()

        ## Train with all-fake batch
        # Generate batch of latent vectors
        noise = torch.randn(b_size, nz, 1, 1, device=device)
        # Generate fake image batch with G
        fake = netG(noise)
        label.fill_(fake_label)
        # Classify all fake batch with D
        output = netD(fake.detach()).view(-1)
        # Calculate D's loss on the all-fake batch
        errD_fake = criterion(output, label)
        # Calculate the gradients for this batch
        errD_fake.backward()
        D_G_z1 = output.mean().item()
        # Add the gradients from the all-real and all-fake batches
        errD = errD_real + errD_fake
        # Update D
        optimizerD.step()

        ############################
        # (2) Update G network: maximize log(D(G(z)))
        ###########################
        netG.zero_grad()
        label.fill_(real_label)  # fake labels are real for generator cost
        # Since we just updated D, perform another forward pass of all-fake batch through D
        output = netD(fake).view(-1)
        # Calculate G's loss based on this output
        errG = criterion(output, label)
        # Calculate gradients for G
        errG.backward()
        D_G_z2 = output.mean().item()
        # Update G
        optimizerG.step()

        # Output training stats
        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))

        # Save Losses for plotting later
        G_losses.append(errG.item())
        D_losses.append(errD.item())

        # Check how the generator is doing by saving G's output on fixed_noise
        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

# 結果
# 下面是D＆G的損失與訓練迭代的關係圖。
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()


# G的過程可視化
# 記住在每個訓練epoch之後我們如何在fixed_noise batch中保存生成器的輸出。現在，我們可以通過動畫可視化G的訓練進度。按播放按鈕 開始動畫。
#%%capture
fig = plt.figure(figsize=(8,8))
plt.axis("off")
ims = [[plt.imshow(np.transpose(i,(1,2,0)), animated=True)] for i in img_list]
ani = animation.ArtistAnimation(fig, ims, interval=1000, repeat_delay=1000, blit=True)

HTML(ani.to_jshtml())

#參考網址：http://www.pytorch123.com/SixthSection/Dcgan/