用pytorch實現對抗生成網絡

最近在學習深度學習編程，採用的深度學習框架是pytorch，看的書主要是陳雲編著的《深度學習框架PyTorch入門與實踐》、廖星宇編著的《深度學習入門之PyTorch》、肖志清的《神經網絡與PyTorch實踐》，都是入門的學習材料，適合初學者。

通過近1個多月的學習，基本算是入門了，後面將深度學習與實踐。這裏分享一個《神經網絡與PyTorch實踐》中對抗生成網絡的例子。它是用對抗生成網絡的方法，訓練CIFAR-10的數據集，訓練模型。

生成網絡gnet將大小爲(64,11)的潛在張量轉化爲大小爲(3,32,32)的假數據;鑑別網絡dnet將大小爲(3,32,32)的數據轉化爲大小爲
(1,1,1)的對數賠率張量。下面是整個模型的python代碼，包括（1）數據加載，（2）模型搭建，（3）模型訓練與模型測試。

import torch
import torch.nn as nn
import torch.nn.init as init
import torch.optim
from torch.utils.data import DataLoader
from torchvision.datasets import CIFAR10,CIFAR100
import torchvision.transforms as transforms
from torchvision.utils import save_image
from torchviz import make_dot

dataset = CIFAR100(root='./data',
                  download=True,
                  transform= transforms.ToTensor())
dataloader = DataLoader(dataset, batch_size=64, shuffle=True)

#check the data
#for batch_idx, data in enumerate(dataloader):
#    real_images, _ = data
#    print('real_images size = {}'.format(real_images.size()))
#    batch_size = real_images.size(0)
#    print('#{} has {} images.'.format(batch_idx, batch_size))
#    if batch_idx %100 ==0:
#        path = './data/CIFAR10_shuffled_batch{:03d}.png'.format(batch_idx)
#        save_image(real_images, path, normalize=True)


#construct the generator and discrimiter network
latent_size=64  #潛在大小
n_channel=3   #輸出通道數
n_g_feature=64   #生成網絡隱藏層大小
#construct the generator 
gnet= nn.Sequential(
    #輸入大小 == (64, 1, 1)
    nn.ConvTranspose2d(latent_size, 4 * n_g_feature, kernel_size=4, bias=False),
    nn.BatchNorm2d(4*n_g_feature),
    nn.ReLU(),
    #大小 = （256，4，4）
    nn.ConvTranspose2d(4*n_g_feature, 2 * n_g_feature, kernel_size=4, stride=2, padding=1, bias=False),
    nn.BatchNorm2d(2*n_g_feature),
    nn.ReLU(),
    #大小 = （128， 8，8）
    nn.ConvTranspose2d(2*n_g_feature, n_g_feature, kernel_size=4, stride=2, padding=1, bias= False),
    nn.BatchNorm2d(n_g_feature),
    nn.ReLU(),
    #大小 = （64，16，16）
    nn.ConvTranspose2d(n_g_feature, n_channel, kernel_size=4, stride=2, padding=1),
    nn.Sigmoid(),
    #圖片大小 = （3， 32， 32）
)

#define the instance of GeneratorNet
print(gnet)
if torch.cuda.is_available():
    gnet.to(torch.device('cuda:0'))

#construct the discrimator
n_d_feature = 64  #鑑別網絡隱藏層大小
dnet = nn.Sequential(
    #圖片大小 = （3，32，32）
    nn.Conv2d(n_channel, n_d_feature, kernel_size=4, stride=2, padding=1),
    nn.LeakyReLU(0.2),
    #大小 = （63，16，16）
    nn.Conv2d(n_d_feature, 2*n_d_feature, kernel_size=4, stride=2, padding=1, bias= False),
    nn.BatchNorm2d(2*n_d_feature),
    nn.LeakyReLU(0.2),
    #大小 = （128， 8，8）
    nn.Conv2d(2*n_d_feature, 4*n_d_feature, kernel_size=4, stride=2, padding=1, bias= False),
    nn.BatchNorm2d(4*n_d_feature),
    nn.LeakyReLU(0.2),
    #大小 = （256，4，4）
    nn.Conv2d(4*n_d_feature, 1, kernel_size=4),
    #對數賠率張量大小=（1，1，1）    
    #nn.Sigmoid()
)
print(dnet)
if torch.cuda.is_available():
    dnet.to(torch.device('cuda:0'))

#initialization for gnet and dnet
def weights_init(m):
    if type(m) in [nn.ConvTranspose2d, nn.Conv2d]:
        init.xavier_normal_(m.weight)
    elif type(m) == nn.BatchNorm2d:
        init.normal_(m.weight, 1.0, 0.02)
        init.constant_(m.bias, 0)

gnet.apply(weights_init)
dnet.apply(weights_init)

#網絡的訓練和使用
#要構造一個損失函數並對它進行優化
#定義損失
criterion = nn.BCEWithLogitsLoss()
#定義優化器
goptimizer = torch.optim.Adam(gnet.parameters(), lr=0.0002, betas=(0.5, 0.999))
doptimizer = torch.optim.Adam(dnet.parameters(), lr=0.0002, betas=(0.5, 0.999))

#用於測試的噪聲，用來查看相同的潛在張量在訓練過程中生成圖片的變換
batch_size=64
fixed_noises = torch.randn(batch_size, latent_size, 1,1)

#save the net to file for check
y=gnet(fixed_noises)
vise_graph = make_dot(y, params=dict(gnet.named_parameters()))
vise_graph.view(filename='gnet')


y=dnet(y)
vise_graph = make_dot(y)
vise_graph.view(filename='dnet')


#訓練過程
epoch_num=10
for epoch in range(epoch_num):
    for batch_idx, data in enumerate(dataloader):
        #載入本批次數據
        real_images,_ = data
        batch_size = real_images.size(0)
        
        #訓練鑑別網絡
        labels = torch.ones(batch_size) #設置真實數據對應標籤爲1
        preds = dnet(real_images)  #對真實數據進行判別
        outputs = preds.reshape(-1)
        dloss_real = criterion(outputs, labels)  #真實數據的鑑別損失
        dmean_real = outputs.sigmoid().mean()  #計算鑑別器將多少比例的真實數據判定爲真，僅用於輸出顯示
        
        noises = torch.randn(batch_size, latent_size, 1,1)  #潛在噪聲
        fake_images = gnet(noises)  #生成假數據
        labels = torch.zeros(batch_size)  #假數據對應標籤爲0
        fake = fake_images.detach()  #是的梯度的計算不回溯到生成網絡，可用於加快訓練速度。刪去此步，結果不變
        preds = dnet(fake)
        outputs = preds.view(-1)
        dloss_fake = criterion(outputs, labels)  #假數據的鑑別損失
        dmean_fake = outputs.sigmoid().mean()  #計算鑑別器將多少比例的假數據判定爲真，僅用於輸出顯示
        
        dloss = dloss_real+dloss_fake
        dnet.zero_grad()
        dloss.backward()
        doptimizer.step()
        
        #訓練生成網絡
        labels = torch.ones(batch_size) #生成網絡希望所有生成的數據都是被認爲時真的
        preds = dnet(fake_images) #讓假數據通過假別網絡
        outputs = preds.view(-1)
        gloss = criterion(outputs, labels)  #從真數據看到的損失
        gmean_fake = outputs.sigmoid().mean() #計算鑑別器將多少比例的假數據判斷爲真，僅用於輸出顯示
        gnet.zero_grad()
        gloss.backward()
        goptimizer.step()
        
        #輸出本步訓練結果
        print('[{}/{}]'.format(epoch, epoch_num)+
              '[{}/{}]'.format(batch_idx, len(dataloader))+
              '鑑別網絡損失：{:g} 生成網絡損失：{:g}'.format(dloss, gloss)+
              '真實數據判真比例：{:g} 假數據判真比例：{:g}/{:g}'.format(dmean_real, dmean_fake, gmean_fake))
        if batch_idx %100 == 0:
            fake = gnet(fixed_noises)  #由固定潛在徵糧生成假數據
            save_image(fake, './data/images_epoch{:02d}_batch{:03d}.png'.format(epoch, batch_idx))  #保存假數據
        
        
#保存訓練的網絡
torch.save(gnet, 'gnet.pkl')
torch.save(dnet, 'dnet.pkl')

結果如下