[深度學習][Python]DCGAN訓練生成手寫阿拉伯數字_訓練篇

本文參考TensorFlo官網Deep Convolutional Generative Adversarial Network的程式碼來加以實作說明。

示範如何使用深度卷積生成對抗網路(DCGAN) 生成手寫數位影像。程式碼是使用帶有訓練循環的Keras Sequential APItf.GradientTape編寫的。

使用 MNIST 資料集來訓練生成器和判別器。生成器將生成類似於 MNIST 資料的手寫數字。

動畫顯示

以下動畫顯示了生成器經過 50 個 epoch 訓練後生成的一系列圖像。這些圖像一開始是隨機噪聲，隨著時間的推移越來越像手寫數字。

本文在Colab上執行訓練，因本機電腦過於老舊XD

Python及套件版本

Python version: 3.10.12
imageio version: 2.34.2
matplotlib version: 3.7.1
numpy version: 1.25.2
PIL version: 9.4.0
tensorflow version: 2.15.0
IPython version: 7.34.0

1.載入相關套件

import glob
import imageio
import matplotlib.pyplot as plt
import numpy as np
import os
import PIL
import tensorflow as tf
from tensorflow.keras import layers
import time
from IPython import display

2. 載入訓練資料

# 取得 MNIST 訓練資料
(train_images, train_labels), (_, _) = tf.keras.datasets.mnist.load_data()
# 像素標準化
train_images = train_images.reshape(train_images.shape[0], 28, 28, 1).astype('float32')
train_images = (train_images - 127.5) / 127.5  # 使像素值介於 [-1, 1]

# 參數設定
BUFFER_SIZE = 60000 # 緩衝區大小
BATCH_SIZE = 256    # 訓練批量

# 轉為 Dataset
train_dataset = tf.data.Dataset.from_tensor_slices(train_images) \
        .shuffle(BUFFER_SIZE).batch(BATCH_SIZE).prefetch(BUFFER_SIZE).cache()

3. 定義GAN模型

3.1 定義生成模型

將從一個隨機噪聲向量中生成28x28圖像

# 生成神經網路
def make_generator_model():
    model = tf.keras.Sequential()
    # 第一層 Dense
    #輸入是一個100維的隨機向量，輸出是7x7x256個特徵圖（feature maps）。
    model.add(layers.Dense(7*7*256, use_bias=False, input_shape=(100,))) 
    #對輸出進行批量標準化（Batch Normalization），有助於穩定和加速訓練過程。
    model.add(layers.BatchNormalization())
    #應用Leaky ReLU激活函數，這是一種具有輕微負斜率的ReLU激活函數，可以解決ReLU的“死亡神經元”問題
    model.add(layers.LeakyReLU())
	  ​​#將輸出重塑為7x7x256的形狀，以便於後續的卷積操作。
    model.add(layers.Reshape((7, 7, 256)))
    #檢查
    assert model.output_shape == (None, 7, 7, 256)  # None 代表批量不檢查
	  # 第一層 Conv2DTranspose​
	  # 將低分辨率的特徵圖轉換為高分辨率的特徵圖 
    model.add(layers.Conv2DTranspose(128, (5, 5), strides=(1, 1), 
                                     padding='same', use_bias=False))
    assert model.output_shape == (None, 7, 7, 128)
    model.add(layers.BatchNormalization())
    model.add(layers.LeakyReLU())
		# 第二層 Conv2DTranspose​
    model.add(layers.Conv2DTranspose(64, (5, 5), strides=(2, 2), 
                                     padding='same', use_bias=False))
    assert model.output_shape == (None, 14, 14, 64)
    model.add(layers.BatchNormalization())
    model.add(layers.LeakyReLU())
		# 第三層 Conv2DTranspose​
	  # 將特徵圖尺寸放大到28x28，並使用'tanh'激活函數將輸出範圍限制在[-1, 1]之間。​
    model.add(layers.Conv2DTranspose(1, (5, 5), strides=(2, 2), 
                     padding='same', use_bias=False, activation='tanh'))
    assert model.output_shape == (None, 28, 28, 1)

    return model

顯示生成雜訊

# 產生生成神經網路
generator = make_generator_model()

# 測試產生的雜訊
noise = tf.random.normal([1, 100])
generated_image = generator(noise, training=False)

# 顯示雜訊生成的圖像
plt.imshow(generated_image[0, :, :, 0], cmap='gray')

3.2 定義判別器模型

判别器模型在訓練過程中會學習區分真實圖像和生成圖像，並將其用於改進生成器模型，使得生成的圖像越來越真實。

# 判别神經網路
def make_discriminator_model():
    model = tf.keras.Sequential() #定義模型：使用tf.keras.Sequential定義了一個順序模型。
    model.add(layers.Conv2D(64, (5, 5), strides=(2, 2), padding='same',
                                     input_shape=[28, 28, 1])) #卷積層， 輸入28*28
    model.add(layers.LeakyReLU()) #使用Leaky ReLU作為激活函數
    model.add(layers.Dropout(0.3)) #應用Dropout層，隨機丟棄30%的神經元，防止過擬

    model.add(layers.Conv2D(128, (5, 5), strides=(2, 2), padding='same'))
    model.add(layers.LeakyReLU())
    model.add(layers.Dropout(0.3))

    model.add(layers.Flatten()) # 全連接層（Dense layer）
    model.add(layers.Dense(1)) #一個神經元的全連接層，輸出一個單一的實值，用於判斷輸入圖像是真實還是生成的。

    return model

測試判別器功能

# 測試判别神經網路
discriminator = make_discriminator_model()

# 預測值越大代表越像
decision = discriminator(generated_image)
print (f'預測值={decision}')

4.定義損失函數及優化器

# 使用 TensorFlow 定義的二分類交叉熵損失函數
cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True)

# 定義判别神經網路損失函數為 真實影像 + 生成影像 的損失函數
def discriminator_loss(real_output, fake_output): 
	'''
	真實影像損失 (real_loss)：
	這部分計算真實影像的損失，即將真實影像的輸出與全為1的標籤比較。
	生成影像損失 (fake_loss)：
	這部分計算生成影像的損失，即將生成影像的輸出與全為0的標籤比較。
	'''​
    real_loss = cross_entropy(tf.ones_like(real_output), real_output)
    fake_loss = cross_entropy(tf.zeros_like(fake_output), fake_output)
    total_loss = real_loss + fake_loss
    return total_loss

# 定義生成神經網路損失函數為 生成影像 的損失函數
def generator_loss(fake_output):
#將生成影像的輸出與全為1的標籤比較。生成器的目標是騙過判别器，使其認為生成的影像是真實的，因此使用全1標籤。
    return cross_entropy(tf.ones_like(fake_output), fake_output)

# 優化器均為 Adam
generator_optimizer = tf.keras.optimizers.Adam(1e-4)
discriminator_optimizer = tf.keras.optimizers.Adam(1e-4)

5.儲存檢查點

儲存和復原模型，這在長時間運行的訓練任務中斷時非常有用。

checkpoint_dir = './dcgan_training_checkpoints'
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")
checkpoint = tf.train.Checkpoint(generator_optimizer=generator_optimizer,
                  discriminator_optimizer=discriminator_optimizer,
                  generator=generator,
                  discriminator=discriminator)

6.參數設定

這樣的設計確保了生成器和判别器在每個訓練步驟中都能相互對抗，從而提升生成影像的質量。

# 參數設定
EPOCHS = 50                    # 訓練執行週期，即整個訓練資料集將被完整地訓練50次。
noise_dim = 100                # 雜訊向量大小，這個向量將作為生成神經網路的輸入。
num_examples_to_generate = 16  # 每次生成的樣本數量。

# 產生亂數(雜訊)'
seed = tf.random.normal([num_examples_to_generate, noise_dim]) 

# 定義梯度下降，分別對判别神經網路、生成神經網路進行訓練
@tf.function  # 產生運算圖 裝飾器將函數編譯成 TensorFlow 運算圖以提高性能。
def train_step(images):
    noise = tf.random.normal([BATCH_SIZE, noise_dim])

    with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
        # 生成神經網路進行訓練
        generated_images = generator(noise, training=True)

        # 判别神經網路進行訓練
        real_output = discriminator(images, training=True)           # 真實影像
        fake_output = discriminator(generated_images, training=True) # 生成影像

        # 計算損失
        gen_loss = generator_loss(fake_output)
        disc_loss = discriminator_loss(real_output, fake_output)

    # 梯度下降
    gradients_of_generator = gen_tape.gradient(gen_loss, generator.trainable_variables)
    gradients_of_discriminator = disc_tape.gradient(disc_loss, 
                                                    discriminator.trainable_variables)

    # 更新權重
    generator_optimizer.apply_gradients(zip(gradients_of_generator, 
                                            generator.trainable_variables))
    discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, 
                                                discriminator.trainable_variables))

7.定義訓練函數及產生圖像並存檔

def train(dataset, epochs):
    for epoch in range(epochs):
        start = time.time()

        for image_batch in dataset:
            train_step(image_batch)

        # 產生圖像
        display.clear_output(wait=True)
        generate_and_save_images(generator, epoch + 1, seed)

        # 每 10 個執行週期存檔一次
        if (epoch + 1) % 10 == 0:
            checkpoint.save(file_prefix = checkpoint_prefix)

        print ('epoch {} 花費 {} 秒'.format(epoch + 1, time.time()-start))

    # 顯示最後結果
    display.clear_output(wait=True)
    generate_and_save_images(generator, epochs, seed)

# 產生圖像並存檔
def generate_and_save_images(model, epoch, test_input):
    # 預測
    predictions = model(test_input, training=False)

    # 顯示 4x4 的格子
    fig = plt.figure(figsize=(4, 4))
    for i in range(predictions.shape[0]):
            plt.subplot(4, 4, i+1)
            plt.imshow(predictions[i, :, :, 0] * 127.5 + 127.5, cmap='gray')
            plt.axis('off')

    # 存檔
    plt.savefig('./GAN_result/image_at_epoch_{:04d}.png'.format(epoch))
    plt.show()

8.訓練模型

#產生儲存結果的資料夾
import os
if not os.path.exists('./GAN_result'):
  os.makedirs('./GAN_result'

train(train_dataset, EPOCHS) #需要一段時間

9.顯示結果

# 顯示最後結果
def display_image(epoch_no):
    return PIL.Image.open('./GAN_result/image_at_epoch_{:04d}.png'.format(epoch_no))

display_image(EPOCHS)

10.將訓練過程的存檔圖像轉為GIF 檔，並顯示GIF檔

要先安裝tensorflow_docs才可以使用它的embed在colab上顯示gif檔

!pip install git+https://github.com/tensorflow/docs

# 產生 GIF 檔
anim_file = './GAN_result/dcgan.gif'
with imageio.get_writer(anim_file, mode='I') as writer: #創建一個 GIF 檔案的寫入器​
    filenames = glob.glob('./GAN_result/image*.png') #使用 glob.glob 搜尋符合 image*.png 的所有檔案名稱
    filenames = sorted(filenames)
    for filename in filenames:
        # print(filename)
        image = imageio.imread(filename)
        writer.append_data(image)
    
# 顯示 GIF 檔
import tensorflow_docs.vis.embed as embed

embed.embed_file(anim_file)

11.儲存模型

generator.save('DCGAN.h5')

參考文獻