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)