教你基於MindSpore用DCGAN生成漫畫頭像

本文分享自華為雲社群《【昇思25天學習打卡營打卡指南-第二十天】DCGAN生成漫畫頭像》，作者：JeffDing。

DCGAN生成漫畫頭像

在下面的教程中，我們將透過示例程式碼說明DCGAN網路如何設定網路、最佳化器、如何計算損失函式以及如何初始化模型權重。在本教程中，使用的動漫頭像資料集共有70,171張動漫頭像圖片，圖片大小均為96*96。

GAN基礎原理

這部分原理介紹參考GAN影像生成。

DCGAN原理

DCGAN（深度卷積對抗生成網路，Deep Convolutional Generative Adversarial Networks）是GAN的直接擴充套件。不同之處在於，DCGAN會分別在判別器和生成器中使用卷積和轉置卷積層。

它最早由Radford等人在論文Unsupervised Representation Learning With Deep Convolutional Generative Adversarial Networks中進行描述。判別器由分層的卷積層、BatchNorm層和LeakyReLU啟用層組成。輸入是3x64x64的影像，輸出是該影像為真影像的機率。生成器則是由轉置卷積層、BatchNorm層和ReLU啟用層組成。輸入是標準正態分佈中提取出的隱向量𝑧，輸出是3x64x64的RGB影像。

本教程將使用動漫頭像資料集來訓練一個生成式對抗網路，接著使用該網路生成動漫頭像圖片。

資料準備與處理

首先我們將資料集下載到指定目錄下並解壓。示例程式碼如下：

from download import download

url = "https://download.mindspore.cn/dataset/Faces/faces.zip"

path = download(url, "./faces", kind="zip", replace=True)

資料處理

首先為執行過程定義一些輸入：

batch_size = 128          # 批次大小
image_size = 64           # 訓練影像空間大小
nc = 3                    # 影像彩色通道數
nz = 100                  # 隱向量的長度
ngf = 64                  # 特徵圖在生成器中的大小
ndf = 64                  # 特徵圖在判別器中的大小
num_epochs = 3           # 訓練週期數
lr = 0.0002               # 學習率
beta1 = 0.5               # Adam最佳化器的beta1超引數

定義create_dataset_imagenet函式對資料進行處理和增強操作。

import numpy as np
import mindspore.dataset as ds
import mindspore.dataset.vision as vision

def create_dataset_imagenet(dataset_path):
    """資料載入"""
    dataset = ds.ImageFolderDataset(dataset_path,
                                    num_parallel_workers=4,
                                    shuffle=True,
                                    decode=True)

    # 資料增強操作
    transforms = [
        vision.Resize(image_size),
        vision.CenterCrop(image_size),
        vision.HWC2CHW(),
        lambda x: ((x / 255).astype("float32"))
    ]

    # 資料對映操作
    dataset = dataset.project('image')
    dataset = dataset.map(transforms, 'image')

    # 批次操作
    dataset = dataset.batch(batch_size)
    return dataset

dataset = create_dataset_imagenet('./faces')

透過create_dict_iterator函式將資料轉換成字典迭代器，然後使用matplotlib模組視覺化部分訓練資料。

import matplotlib.pyplot as plt

def plot_data(data):
    # 視覺化部分訓練資料
    plt.figure(figsize=(10, 3), dpi=140)
    for i, image in enumerate(data[0][:30], 1):
        plt.subplot(3, 10, i)
        plt.axis("off")
        plt.imshow(image.transpose(1, 2, 0))
    plt.show()

sample_data = next(dataset.create_tuple_iterator(output_numpy=True))
plot_data(sample_data)

構造網路

當處理完資料後，就可以來進行網路的搭建了。按照DCGAN論文中的描述，所有模型權重均應從mean為0，sigma為0.02的正態分佈中隨機初始化。

生成器

生成器G的功能是將隱向量z對映到資料空間。由於資料是影像，這一過程也會建立與真實影像大小相同的 RGB 影像。在實踐場景中，該功能是透過一系列Conv2dTranspose轉置卷積層來完成的，每個層都與BatchNorm2d層和ReLu啟用層配對，輸出資料會經過tanh函式，使其返回[-1,1]的資料範圍內。

DCGAN論文生成影像如下所示：

圖片來源：Unsupervised Representation Learning With Deep Convolutional Generative Adversarial Networks.

我們透過輸入部分中設定的nz、ngf和nc來影響程式碼中的生成器結構。nz是隱向量z的長度，ngf與透過生成器傳播的特徵圖的大小有關，nc是輸出影像中的通道數。

以下是生成器的程式碼實現：

import mindspore as ms
from mindspore import nn, ops
from mindspore.common.initializer import Normal

weight_init = Normal(mean=0, sigma=0.02)
gamma_init = Normal(mean=1, sigma=0.02)

class Generator(nn.Cell):
    """DCGAN網路生成器"""

    def __init__(self):
        super(Generator, self).__init__()
        self.generator = nn.SequentialCell(
            nn.Conv2dTranspose(nz, ngf * 8, 4, 1, 'valid', weight_init=weight_init),
            nn.BatchNorm2d(ngf * 8, gamma_init=gamma_init),
            nn.ReLU(),
            nn.Conv2dTranspose(ngf * 8, ngf * 4, 4, 2, 'pad', 1, weight_init=weight_init),
            nn.BatchNorm2d(ngf * 4, gamma_init=gamma_init),
            nn.ReLU(),
            nn.Conv2dTranspose(ngf * 4, ngf * 2, 4, 2, 'pad', 1, weight_init=weight_init),
            nn.BatchNorm2d(ngf * 2, gamma_init=gamma_init),
            nn.ReLU(),
            nn.Conv2dTranspose(ngf * 2, ngf, 4, 2, 'pad', 1, weight_init=weight_init),
            nn.BatchNorm2d(ngf, gamma_init=gamma_init),
            nn.ReLU(),
            nn.Conv2dTranspose(ngf, nc, 4, 2, 'pad', 1, weight_init=weight_init),
            nn.Tanh()
            )

    def construct(self, x):
        return self.generator(x)

generator = Generator()

判別器

如前所述，判別器D是一個二分類網路模型，輸出判定該影像為真實圖的機率。透過一系列的Conv2d、BatchNorm2d和LeakyReLU層對其進行處理，最後透過Sigmoid啟用函式得到最終機率。

DCGAN論文提到，使用卷積而不是透過池化來進行下采樣是一個好方法，因為它可以讓網路學習自己的池化特徵。

判別器的程式碼實現如下：

class Discriminator(nn.Cell):
    """DCGAN網路判別器"""

    def __init__(self):
        super(Discriminator, self).__init__()
        self.discriminator = nn.SequentialCell(
            nn.Conv2d(nc, ndf, 4, 2, 'pad', 1, weight_init=weight_init),
            nn.LeakyReLU(0.2),
            nn.Conv2d(ndf, ndf * 2, 4, 2, 'pad', 1, weight_init=weight_init),
            nn.BatchNorm2d(ngf * 2, gamma_init=gamma_init),
            nn.LeakyReLU(0.2),
            nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 'pad', 1, weight_init=weight_init),
            nn.BatchNorm2d(ngf * 4, gamma_init=gamma_init),
            nn.LeakyReLU(0.2),
            nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 'pad', 1, weight_init=weight_init),
            nn.BatchNorm2d(ngf * 8, gamma_init=gamma_init),
            nn.LeakyReLU(0.2),
            nn.Conv2d(ndf * 8, 1, 4, 1, 'valid', weight_init=weight_init),
            )
        self.adv_layer = nn.Sigmoid()

    def construct(self, x):
        out = self.discriminator(x)
        out = out.reshape(out.shape[0], -1)
        return self.adv_layer(out)

discriminator = Discriminator()

模型訓練

損失函式

當定義了D和G後，接下來將使用MindSpore中定義的二進位制交叉熵損失函式BCELoss。

# 定義損失函式
adversarial_loss = nn.BCELoss(reduction='mean')

最佳化器

這裡設定了兩個單獨的最佳化器，一個用於D，另一個用於G。這兩個都是lr = 0.0002和beta1 = 0.5的Adam最佳化器。

# 為生成器和判別器設定最佳化器
optimizer_D = nn.Adam(discriminator.trainable_params(), learning_rate=lr, beta1=beta1)
optimizer_G = nn.Adam(generator.trainable_params(), learning_rate=lr, beta1=beta1)
optimizer_G.update_parameters_name('optim_g.')
optimizer_D.update_parameters_name('optim_d.')

訓練模型

訓練分為兩個主要部分：訓練判別器和訓練生成器。

• 訓練判別器

訓練判別器的目的是最大程度地提高判別影像真偽的機率。按照Goodfellow的方法，是希望透過提高其隨機梯度來更新判別器，所以我們要最大化logD(x)+log(1−D(G(z))的值。

• 訓練生成器

如DCGAN論文所述，我們希望透過最小化log(1−D(G(z)))來訓練生成器，以產生更好的虛假影像。

在這兩個部分中，分別獲取訓練過程中的損失，並在每個週期結束時進行統計，將fixed_noise批次推送到生成器中，以直觀地跟蹤G的訓練進度。

下面實現模型訓練正向邏輯：

def generator_forward(real_imgs, valid):
    # 將噪聲取樣為發生器的輸入
    z = ops.standard_normal((real_imgs.shape[0], nz, 1, 1))

    # 生成一批影像
    gen_imgs = generator(z)

    # 損失衡量發生器繞過判別器的能力
    g_loss = adversarial_loss(discriminator(gen_imgs), valid)

    return g_loss, gen_imgs

def discriminator_forward(real_imgs, gen_imgs, valid, fake):
    # 衡量鑑別器從生成的樣本中對真實樣本進行分類的能力
    real_loss = adversarial_loss(discriminator(real_imgs), valid)
    fake_loss = adversarial_loss(discriminator(gen_imgs), fake)
    d_loss = (real_loss + fake_loss) / 2
    return d_loss

grad_generator_fn = ms.value_and_grad(generator_forward, None,
                                      optimizer_G.parameters,
                                      has_aux=True)
grad_discriminator_fn = ms.value_and_grad(discriminator_forward, None,
                                          optimizer_D.parameters)

@ms.jit
def train_step(imgs):
    valid = ops.ones((imgs.shape[0], 1), mindspore.float32)
    fake = ops.zeros((imgs.shape[0], 1), mindspore.float32)

    (g_loss, gen_imgs), g_grads = grad_generator_fn(imgs, valid)
    optimizer_G(g_grads)
    d_loss, d_grads = grad_discriminator_fn(imgs, gen_imgs, valid, fake)
    optimizer_D(d_grads)

    return g_loss, d_loss, gen_imgs

迴圈訓練網路，每經過50次迭代，就收集生成器和判別器的損失，以便於後面繪製訓練過程中損失函式的影像。

import mindspore

G_losses = []
D_losses = []
image_list = []

total = dataset.get_dataset_size()
for epoch in range(num_epochs):
    generator.set_train()
    discriminator.set_train()
    # 為每輪訓練讀入資料
    for i, (imgs, ) in enumerate(dataset.create_tuple_iterator()):
        g_loss, d_loss, gen_imgs = train_step(imgs)
        if i % 100 == 0 or i == total - 1:
            # 輸出訓練記錄
            print('[%2d/%d][%3d/%d]   Loss_D:%7.4f  Loss_G:%7.4f' % (
                epoch + 1, num_epochs, i + 1, total, d_loss.asnumpy(), g_loss.asnumpy()))
        D_losses.append(d_loss.asnumpy())
        G_losses.append(g_loss.asnumpy())

    # 每個epoch結束後，使用生成器生成一組圖片
    generator.set_train(False)
    fixed_noise = ops.standard_normal((batch_size, nz, 1, 1))
    img = generator(fixed_noise)
    image_list.append(img.transpose(0, 2, 3, 1).asnumpy())

    # 儲存網路模型引數為ckpt檔案
    mindspore.save_checkpoint(generator, "./generator.ckpt")
    mindspore.save_checkpoint(discriminator, "./discriminator.ckpt")

結果展示

執行下面程式碼，描繪D和G損失與訓練迭代的關係圖：

plt.figure(figsize=(10, 5))
plt.title("Generator and Discriminator Loss During Training")
plt.plot(G_losses, label="G", color='blue')
plt.plot(D_losses, label="D", color='orange')
plt.xlabel("iterations")
plt.ylabel("Loss")
plt.legend()
plt.show()

視覺化訓練過程中透過隱向量fixed_noise生成的影像。

import matplotlib.pyplot as plt
import matplotlib.animation as animation

def showGif(image_list):
    show_list = []
    fig = plt.figure(figsize=(8, 3), dpi=120)
    for epoch in range(len(image_list)):
        images = []
        for i in range(3):
            row = np.concatenate((image_list[epoch][i * 8:(i + 1) * 8]), axis=1)
            images.append(row)
        img = np.clip(np.concatenate((images[:]), axis=0), 0, 1)
        plt.axis("off")
        show_list.append([plt.imshow(img)])

    ani = animation.ArtistAnimation(fig, show_list, interval=1000, repeat_delay=1000, blit=True)
    ani.save('./dcgan.gif', writer='pillow', fps=1)

showGif(image_list)

從上面的影像可以看出，隨著訓練次數的增多，影像質量也越來越好。如果增大訓練週期數，當num_epochs達到50以上時，生成的動漫頭像圖片與資料集中的較為相似，下面我們透過載入生成器網路模型引數檔案來生成影像，程式碼如下：

# 從檔案中獲取模型引數並載入到網路中
mindspore.load_checkpoint("./generator.ckpt", generator)

fixed_noise = ops.standard_normal((batch_size, nz, 1, 1))
img64 = generator(fixed_noise).transpose(0, 2, 3, 1).asnumpy()

fig = plt.figure(figsize=(8, 3), dpi=120)
images = []
for i in range(3):
    images.append(np.concatenate((img64[i * 8:(i + 1) * 8]), axis=1))
img = np.clip(np.concatenate((images[:]), axis=0), 0, 1)
plt.axis("off")
plt.imshow(img)
plt.show()

點選關注，第一時間瞭解華為雲新鮮技術~