天天看點
傳回

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/
pytorch學習 pytorch GAN
上一篇: pytorch-Detach的作用
下一篇: PyTorch, GAN筆記

繼續閱讀