DCGAN生成二次元頭像(Pytorch)

任務:基於Pytorch搭建DCGAN網絡進行訓練自動生成二次元頭像的模型

簡介

好久沒發Blog了,發個Pytorch的入門項目一起happy下~
數據集:網盤地址
有大佬在知乎專欄提供了公開的二次元頭像的數據集。本篇博客即在這個數據集上進行訓練。

網絡介紹

DCGAN的全稱是深度卷積生成對抗網絡,使用深度卷積網絡作爲鑑別器,用反捲積神經網絡作爲生成器。
在DCGAN的訓練過程中,生成模型的訓練目標是使得生成的圖片可以很好地欺騙判別模型,使得判別模型認爲生成模型生成的圖片是"真實"的;而判別模型的訓練目標則是儘量地正確區分生成模型生成的圖片和真實存在的圖片。於是,這種訓練方式就很自然地產生了生成模型和判別模型之間的"博弈"。
在理想情況下,我們希望DCGAN訓練好之後,生成模型生成的圖片是可以以假亂真的。
鑑別器
原數據集的圖片大小爲96*96
第一層:CNN
卷積大小4*4,步長2,輸出通道數32
輸入5*3*96*96
輸出5*32*48*48
第二層:CNN
卷積大小4*4,步長2,輸出通道數64
輸入5*32*48*48
輸出5*64*24*24
第三層:CNN
卷積大小4*4,步長2,輸出通道數128
輸入5*64*24*24
輸出5*128*12*12
第四層:CNN
卷積大小4*4,步長2,輸出通道數256
輸入5*128*12*12
輸出5*256*6*6
展開爲256*6*6一維標量
第五層:線性層
使用sigmoid函數進行二元分類

生成器
隨機生成100維向量
第一層:線性層
輸入5*100
輸出5*1024*6*6
第二層:transpose CNN
卷積大小4*4,步長2,輸出通道數512
輸入5*1024*6*6
輸出5*512*12*12
第三層:transpose CNN
卷積大小4*4,步長2,輸出通道數256
輸入5*512*12*12
輸出5*256*24*24
第四層:transpose CNN
卷積大小4*4,步長2,輸出通道數128
輸入5*256*24*24
輸出5*128*48*48
第四層:transpose CNN
卷積大小4*4,步長2,輸出通道數3
輸入5*128*48*48
輸出5*3*96*96

訓練過程
(1)固定生成器,訓練鑑別器
(2)固定鑑別器,訓練生成器

代碼詳解

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
from torchvision import datasets,transforms,models
import matplotlib.pyplot as plt
import os

#對圖片進行變換
#RandomHorizontalFlip:隨機水平翻轉給定的圖片,概率爲0.5。即:一半的概率翻轉,一半的概率不翻轉
#ToTensor:將圖片轉爲Tensor格式
#Normalize::用給定的均值和標準差分別對每個通道的數據進行正則化。
#   給定均值(M1,…,Mn),給定標準差(S1,…,Sn),其中n是通道數(一般是3),對每個通道進行
#       output[channel] = (input[channel] - mean[channel]) / std[channel]
data_transform = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])

trainset = datasets.ImageFolder('faces', data_transform) #數據集
#加載數據集的Dataloader,每次隨機打亂,加載5張圖片
trainloader = torch.utils.data.DataLoader(trainset, batch_size=5, shuffle=True)

#展示並保存圖片
def imshow(inputs, picname):
    plt.ion()
    inputs = inputs / 2 + 0.5 #加載數據集的時候使用了正則化,需要恢復一下
    inputs = inputs.numpy().transpose((1, 2, 0)) #pytorch加載圖片通道數在前,我們展示圖片圖片的時候通道數在後
    plt.imshow(inputs)
    # plt.show()
    plt.pause(1)
    plt.savefig(picname+'.jpg') #保存圖片
    plt.close()

#鑑別器
class D(nn.Module):
    def __init__(self, nc, ndf):
        super(D, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(nc, ndf, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(ndf),
            nn.LeakyReLU(0.2, inplace=True)
        )
        self.layer2 = nn.Sequential(
            nn.Conv2d(ndf, ndf*2, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(ndf*2),
            nn.LeakyReLU(0.2, inplace=True)
        )
        self.layer3 = nn.Sequential(
            nn.Conv2d(ndf*2, ndf*4, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(ndf*4),
            nn.LeakyReLU(0.2, inplace=True)
        )
        self.layer4 = nn.Sequential(
            nn.Conv2d(ndf*4, ndf*8, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(ndf*8),
            nn.LeakyReLU(0.2, inplace=True)
        )
        self.fc = nn.Sequential(
            nn.Linear(256*6*6, 1),
            nn.Sigmoid()
        )
    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.layer4(out)
        out = out.view(-1, 256*6*6)
        out = self.fc(out)
        return out

#生成器
class G(nn.Module):
    def __init__(self, nc, ngf, nz, feature_size):
        super(G, self).__init__()
        self.prj = nn.Linear(feature_size, nz*6*6)
        self.layer1 = nn.Sequential(
            nn.ConvTranspose2d(nz, ngf*4, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(ngf*4),
            nn.ReLU()
        )
        self.layer2 = nn.Sequential(
            nn.ConvTranspose2d(ngf*4, ngf*2, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(ngf*2),
            nn.ReLU()
        )
        self.layer3 = nn.Sequential(
            nn.ConvTranspose2d(ngf*2, ngf, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(ngf),
            nn.ReLU()
        )
        self.layer4 = nn.Sequential(
            nn.ConvTranspose2d(ngf, nc, kernel_size=4, stride=2, padding=1),
            nn.Tanh()
        )

    def forward(self, x):
        out = self.prj(x)
        out = out.view(-1, 1024, 6, 6)
        out = self.layer1(out)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.layer4(out)
        return out

# inputs, _ = next(iter(trainloader))
# imshow(torchvision.utils.make_grid(inputs), "RealDataSample")

d = D(3, 32) #定義鑑別器,輸入通道爲3,輸出通道爲32
g = G(3, 128, 1024, 100) #定義生成器,最後一層輸出通道爲3,輸入通道爲128,線性層輸出通道爲128,輸入通道爲100

criterion = nn.BCELoss() #使用交叉熵作爲損失函數
#生成器和鑑別器均採用Adam優化器
d_optimizer = torch.optim.Adam(d.parameters(), lr=0.0003)
g_optimizer = torch.optim.Adam(g.parameters(), lr=0.0003)

def train(d, g, criterion, d_optimizer, g_optimizer, epochs=1, show_every=1000, print_every=10):
    iter_count = 0
    for epoch in range(epochs):
        for inputs,_ in trainloader:
            real_inputs = inputs #5張真實圖片
            fake_inputs = g(torch.randn(5, 100)) #5張假圖片

            #對應標籤
            real_labels = torch.ones(real_inputs.size(0))
            fake_labels = torch.zeros(5)

            #鑑別真實圖片
            real_outputs = d(real_inputs)
            d_loss_real = criterion(real_outputs, real_labels)
            real_scores = real_outputs

            #鑑別假圖片
            fake_outputs = d(fake_inputs)
            d_loss_fake = criterion(fake_outputs, fake_labels)
            fake_scores = fake_outputs

            d_loss = d_loss_fake + d_loss_real #計算總損失
            #反向傳播
            d_optimizer.zero_grad()
            d_loss.backward()
            d_optimizer.step() #更新鑑別器參數

            #生成假圖片,並用更新後的鑑別器進行鑑別
            fake_inputs = g(torch.randn(5, 100))
            outputs = d(fake_inputs) #鑑別器的標籤
            real_labels = torch.ones(outputs.size(0)) #真實的標籤
            g_loss = criterion(outputs, real_labels)
            #反向傳播
            g_optimizer.zero_grad()
            g_loss.backward()
            g_optimizer.step() #更新鑑別器
            
            #保存圖片
            if iter_count % show_every == 0:
                print('Epoch:{}, Iter:{}, D:{:.4}, G:{:.4}'.format(epoch, iter_count, d_loss.item(), g_loss.item()))

                picname = "Epoch_"+str(epoch)+"Iter_"+str(iter_count)
                imshow(torchvision.utils.make_grid(fake_inputs.data), picname)
            #打印損失函數
            if iter_count % print_every == 0:
                print('Epoch:{}, Iter:{}, D:{:.4}, G:{:.4}'.format(epoch, iter_count, d_loss.item(), g_loss.item()))

            iter_count += 1
    print("Finsh")

train(d, g, criterion, d_optimizer, g_optimizer, epochs=100)

實驗結果

epoch=50
在這裏插入圖片描述
epoch=100
在這裏插入圖片描述
我只訓練到100次,這個時候的差多動漫人物的雛形已經有了,還有些細節要改善。大家也可以試一下調大epoch,結果應該會更好~。

參考

DCGAN論文及代碼學習

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章