tensorflow 2 keras gan手寫數字 對抗學習註解

該段實例來源於網絡，做了一些修改和標註

from __future__ import absolute_import, division, print_function, unicode_literals
import tensorflow as tf
from tensorflow import keras
import glob
import imageio
import matplotlib.pyplot as plt
import numpy as np
import os
import PIL
from tensorflow.keras import layers
import time
import sys

# 使用手寫字體或單品樣本做訓練  這裏注意的是 我們只需要訓練數據，不需要答案和測試數據集。
(train_images, _), (_, _) = keras.datasets.mnist.load_data()
# (train_images, _), (_, _) = keras.datasets.fashion_mnist.load_data()

# 因爲卷積層的需求，增加色深維度
train_images = train_images.reshape(train_images.shape[0], 28, 28, 1).astype('float32')
# 規範化爲-1 - +1
train_images = (train_images - 127.5) / 127.5

BUFFER_SIZE = 60000  # 以供60000個樣本
BATCH_SIZE = 256  # 256張爲一組
train_dataset = tf.data.Dataset.from_tensor_slices(train_images).shuffle(BUFFER_SIZE).batch(BATCH_SIZE)

# 圖片生成模型
def make_generator_model():  # 根據長度爲100的隨機數組，生成一張28，28，1的矩陣
    model = tf.keras.Sequential()
    # 全聯接層，輸入緯度爲[[100],[n]],  輸出爲7*7*256 = 12544的節點  use_bias=False不使用偏差
    model.add(layers.Dense(7*7*256, use_bias=False, input_shape=(100,)))
    # BatchNormalization層：該層在每個batch上將前一層的激活值重新規範化，即使得其輸出數據的均值接近0，其標準差接近1
    # 該層作用：（1）加速收斂（2）控制過擬合，可以少用或不用Dropout和正則（3）降低網絡對初始化權重不敏感（4）允許使用較大的學習率
    model.add(layers.BatchNormalization())
    # ReLU是將所有的負值都設爲零，相反，Leaky ReLU是給所有負值賦予一個非零斜率（負數）
    model.add(layers.LeakyReLU())
    # 將平鋪的節點轉爲7*7*256的shape
    model.add(layers.Reshape((7, 7, 256)))
    # 驗證圖形。非則斷點
    assert model.output_shape == (None, 7, 7, 256)  # Note: None is the batch size
    # 通俗的講這個解卷積，也就做反捲積，也叫做轉置卷積（最貼切），我們就叫做反捲積吧，它的目的就是卷積的反向操作
    # 個人理解，正常的卷積是提取卷積核特徵，反捲積就是用卷積核反向修改圖像，風格遷移應該也是這麼回事，那麼問題來了在這個gan中，卷積特徵從哪來？
    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())
    #64, (5, 5), strides=(2, 2), 希望得到64個特徵核，步長2，2
    #model.output_shape == (None, 14, 14, 64) 輸出的節點數64就是上面的特徵核，由於padding='same'，所以卷積後無變化，
    #14，14 是因爲步長 2，2  所以7*2
    model.add(layers.Conv2DTranspose(64, (5, 5), strides=(2, 2), padding='same', use_bias=False))
    assert model.output_shape == (None, 14, 14, 64)  # 12544
    model.add(layers.BatchNormalization())
    model.add(layers.LeakyReLU())
    model.add(layers.Conv2DTranspose(1, (5, 5), strides=(2, 2), padding='same', use_bias=False, activation='tanh'))
    #驗證上層是否輸出一個 28.28.1 總量爲784的圖像矩陣
    assert model.output_shape == (None, 28, 28, 1)
    return model
generator = make_generator_model()

def make_discriminator_model():  # 原圖、生成圖辨別網絡  這個模型與上一個相反，難道是爲了提取特徵？
    model = tf.keras.Sequential()
    # 將 28.28.1的圖像卷積 輸出64個節點
    model.add(layers.Conv2D(64, (5, 5), strides=(2, 2), padding='same', input_shape=[28, 28, 1]))
    model.add(layers.LeakyReLU())
    model.add(layers.Dropout(0.3))
    # 接着卷積出128個節點
    model.add(layers.Conv2D(128, (5, 5), strides=(2, 2), padding='same'))
    #激活函數 爲非0的斜率
    model.add(layers.LeakyReLU())
    model.add(layers.Dropout(0.3))
    # 平鋪 並輸出一個數字
    model.add(layers.Flatten())
    model.add(layers.Dense(1))
    return model
discriminator = make_discriminator_model()

#這段大概就是個測試
# 隨機生成一個向量，用於生成圖片
noise = tf.random.normal([1, 100])
# 生成一張，此時模型未經訓練，圖片爲噪點 generator = keras.model
generated_image = generator(noise, training=False)
# plt.imshow(generated_image[0, :, :, 0], cmap='gray')
# 判斷結果
decision = discriminator(generated_image)
# 此時的結果應當應當趨近於0，表示爲僞造圖片
print(decision)


# 交叉熵損失函數
cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True)

# 辨別模型損失函數
def discriminator_loss(real_output, fake_output):
    # 樣本圖希望結果趨近1
    real_loss = cross_entropy(tf.ones_like(real_output), real_output)
    # 自己生成的圖希望結果趨近0
    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，也就是真實樣本
    return cross_entropy(tf.ones_like(fake_output), fake_output)

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

# 訓練結果保存
checkpoint_dir = r'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)

EPOCHS = 100
noise_dim = 100
num_examples_to_generate = 16

# 初始化16個種子向量，用於生成4x4的圖片  seed shape: 16, 100
seed = tf.random.normal([num_examples_to_generate, noise_dim])


# @tf.function表示TensorFlow編譯、緩存此函數，用於在訓練中快速調用
@tf.function
def train_step(images):  #更新 模型權重數據的核心方法
    # 隨機生成一個批次的種子向量 BATCH_SIZE = 256   noise_dim = 100  ，256個長度爲100的噪音響亮
    noise = tf.random.normal([BATCH_SIZE, noise_dim]) #noise shape:[256],[100]

    #查看每一次epoch參數更新  這個GradientTape 是每次梯度更新都會調用的，這個取代了model.fit的訓練計算
    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)

    # 計算出的參量應用到模型   梯度修剪，用於改變值， 梯度修剪主要避免訓練梯度爆炸和消失問題
    #zIP是個格式轉換函數 例如：a = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]； zip(*a) = [(1, 4, 7), (2, 5, 8), (3, 6, 9)]
    generator_optimizer.apply_gradients(zip(gradients_of_generator, generator.trainable_variables))
    discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, discriminator.trainable_variables))

#訓練
def train(dataset, epochs):  
    for epoch in range(epochs+1):
        start = time.time()

        # 訓練
        for image_batch in dataset:
            train_step(image_batch)

        #保存圖片
        # 每個訓練批次生成一張圖片作爲階段成功
        print("=======================================")
        generate_and_save_images(generator, epoch + 1, seed)

        # 保存模型
        # 每20次迭代保存一次訓練數據
        # if (epoch + 1) % 20 == 0: # 註銷該行每次都保存
        checkpoint.save(file_prefix=checkpoint_prefix)

        print('Time for epoch {} is {} sec'.format(epoch + 1, time.time()-start))

# 生成圖片
def generate_and_save_images(model, epoch, test_input):
    # 設置爲非訓練狀態，生成一組圖片
    predictions = model(test_input, training=False)

    fig = plt.figure(figsize=(4, 4))

    # 4格x4格拼接
    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')

    # 保存爲png
    plt.savefig('image_at_epoch_{:04d}.png'.format(epoch))
    # plt.show()
    plt.close()

# 遍歷所有png圖片，彙總爲gif動圖
def write_gif():
    anim_file = 'dcgan.gif'
    with imageio.get_writer(anim_file, mode='I') as writer:
        filenames = glob.glob('image*.png')
        filenames = sorted(filenames)
        last = -1
        for i, filename in enumerate(filenames):
            frame = 2*(i**0.5)
            if round(frame) > round(last):
                last = frame
            else:
                continue
            image = imageio.imread(filename)
            writer.append_data(image)
        image = imageio.imread(filename)
        writer.append_data(image)

# 生成一張初始狀態的4格圖片，應當是噪點
generate_and_save_images(generator, 0000, seed)


# 如果使用train參數運行則進入訓練模式
TRAIN = True
if len(sys.argv) == 2 and sys.argv[1] == 'train':
    TRAIN = True

if TRAIN:
    # 以訓練模式運行，進入訓練狀態
    checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir)) #啓用該段，應該可以繼續訓練
    train(train_dataset, EPOCHS)
    write_gif()
else:
    # 非訓練模式，恢復訓練數據
    checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))
    print("After training:")
    # 顯示訓練完成後，生成圖片的辨別結果
    generated_image = generator(noise, training=False)
    decision = discriminator(generated_image)
    # 結果應當趨近1
    print(decision)
    # 重新生成隨機值，生成一組圖片保存
    seed = tf.random.normal([num_examples_to_generate, noise_dim])

    generate_and_save_images(generator, 9999, seed)