6-Gans-03_Face_Generator人脸生成

"""
人脸生成项目
本项目中：使用生成式对抗网络（Generative Adversarial Nets）来生成新的人脸图像。
### 获取数据
使用以下数据集：
- MNIST
- CelebA
由于 CelebA 数据集非常大。所以先在 MNIST 数据集上测试 GANs模型，然后再将其应用到CelebA数据集上。
"""

import os
import helper
from glob import glob
from matplotlib import pyplot
import matplotlib as mpl
import helper
import tensorflow as tf
import numpy as np

# 设置字符集，防止中文乱码
mpl.rcParams['font.sans-serif'] = [u'simHei']
mpl.rcParams['axes.unicode_minus'] = False

# 下载数据
data_dir = './data'
helper.download_extract('mnist', data_dir)
helper.download_extract('celeba', data_dir)

# 探索数据-1、Mnist
def explore_mnist():
    show_n_images = 1
    # glob是一个API 返回路径下所有文件或图片的index,并组成1个list.
    mnist_images = helper.get_batch(glob(os.path.join(data_dir, 'mnist/*.jpg'))[:show_n_images], 28, 28, 'L')
    pyplot.imshow(helper.images_square_grid(mnist_images, 'L'), cmap='gray')
    pyplot.show()


# 探索数据-2、CelebA
# [CelebFaces Attributes Dataset (CelebA)](http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html)
# 是一个包含 20 多万张名人图片的数据集。用此数据集生成人脸。
def explore_CelebA():
    show_n_images = 25
    mnist_images = helper.get_batch(glob(os.path.join(data_dir, 'img_align_celeba/*.jpg'))[:show_n_images], 28, 28, 'RGB')
    pyplot.imshow(helper.images_square_grid(mnist_images, 'RGB'))
    pyplot.show()

"""
## 预处理数据（Preprocess the Data）
经过数据预处理，MNIST 和 CelebA 数据集的值在 28×28 维度图像的 [-0.5, 0.5] 范围内。CelebA 数据集中的图像裁剪了
非脸部的图像部分，然后调整到 28x28 维度。
MNIST 数据集中的图像是单[通道]的黑白图像，CelebA 数据集中的图像是 [三通道的 RGB 彩色图像]。
"""

# todo 建立神经网络（Build the Neural Network）
""" 
 GANs网络的主要组成部分:
    - `model_inputs`
    - `discriminator`
    - `generator`
    - `model_loss`
    - `model_opt`
    - `train`
"""

# todo 输入
"""
部署 `model_inputs` 函数以创建用于神经网络的 [占位符 (TF Placeholders)]
- 输入图像占位符: 使用 `image_width`，`image_height` 和 `image_channels` 设置为 维度= 4。
- 输入 Z 占位符: 设置为 维度= 2，并命名为 `z_dim`。
- 学习速率占位符: 设置为 维度为 0。

返回占位符元组的形状为 (tensor of real input images, tensor of z data, learning rate)。
"""


def model_inputs(image_width, image_height, image_channels, z_dim):
    """
    Create the model inputs
    :param image_width: The input image width
    :param image_height: The input image height
    :param image_channels: The number of image channels
    :param z_dim: The dimension of Z
    :return: Tuple of (tensor of real input images, tensor of z data, learning rate)
    """
    # 需要编程
    inputs_real = tf.placeholder(tf.float32, [None, image_width, image_height, image_channels], name='inputs_real')
    inputs_z = tf.placeholder(tf.float32, [None, z_dim], name='inputs_z')
    learning_rate = tf.placeholder(tf.float32, name='learning_rate')

    return inputs_real, inputs_z, learning_rate

# todo-构建辨别器网络
"""
Gans-D（辨别器）
- 典型的CNN（但没有池化层），将图片局部信息泛化到全局信息，并最终作分类判定该图片为真或假（1或0）
- 1、隐藏层激活：用Lrelu替代了Relu，减少Relu带来的梯度消失（大于0的，梯度为1，反之梯度为0）
- 2、无池化层
- 3、需要做批归一化（作用：网络训练更快，权重初始化更容易，使激活函数有更多选择，创建深网络更简单；总之，使网络效果更佳）

部署 `discriminator` 函数创建辨别器神经网络以辨别 `images`。
该函数返回：(tensor output of the discriminator, tensor logits of the discriminator) 。
- 注意事项:
  - tf.layers.batch_normalization(x2,training=True)中的参数：training
  -  在 [`tf.variable_scope`]中使用 "discriminator" 的变量空间名来重复使用该函数中的变量。 
  - Lrelu激活函数实现，用tf.maximum()
"""
def discriminator(images, reuse=False):
    """
    Create the discriminator network
    :param image: Tensor of input image(s)
    :param reuse: Boolean if the weights should be reused
    :return: Tuple of (tensor output of the discriminator, tensor logits of the discriminator)
    """
    # todo-需要编程 ：
    alpha=0.2
    with tf.variable_scope('discriminator',reuse=reuse):
        #input layer tensor 28*28*3 to 14*14*64 no batch Normalization
        x1=tf.layers.conv2d(images,64,5,strides=2,padding='same',kernel_initializer=tf.contrib.layers.xavier_initializer())
        x1=tf.maximum(x1*alpha,x1)
        x1=tf.nn.dropout(x1,0.8)
        #14*14*64 to 8*8*128
        x2=tf.layers.conv2d(x1,128,5,strides=2,padding='same',kernel_initializer=tf.contrib.layers.xavier_initializer())
        x2=tf.layers.batch_normalization(x2,training=True)
        x2=tf.maximum(x2*alpha,x2)
        x2=tf.nn.dropout(x2,0.8)
        #8*8*128 to 4*4*256
        x3=tf.layers.conv2d(x2,256,5,strides=2,padding='same',kernel_initializer=tf.contrib.layers.xavier_initializer())
        x3=tf.layers.batch_normalization(x3,training=True)
        x3=tf.maximum(x3*alpha,x3)
        x3=tf.nn.dropout(x3,0.8)
        #4*4*256  flatten to 4*4*256
        x4=tf.reshape(x3,[-1,4*4*256])
        logits=tf.layers.dense(x4,1,activation=None)
        out=tf.sigmoid(logits)

    return out,logits
# todo-构建生成器网络
"""
生成器：
- 1、随机生成一个长度为100的Z向量，全连接为4\*4\*1024长度的向量，并重塑为shape=4\*4\*1024的图片；
- 2、通过转置卷积（conv2d_transpose），进行上采样（upsampling），逐步减少depth，增加高和宽，最终生成shape和真图片shape一样的图片（包括值域）。real shape=28\*28\*3
注意事项：
- 1、隐藏层激活：Lrelu
- 2、无池化层
- 3、需要做批归一化
- 4、最后一步激活函数使用：tanh双曲正切函数

部署 `generator` 函数以使用 `z` 生成图像。 
该函数应返回： 28 x 28 x `out_channel_dim` 维度图像。
- 编程注意事项：
  - 在 [`tf.variable_scope`] 中使用 "generator" 的变量空间名来重复使用该函数中的变量。reuse参数应设置为False
"""

def generator(z, out_channel_dim, is_train=True):
    """
    Create the generator network
    :param z: Input z
    :param out_channel_dim: The number of channels in the output image
    :param is_train: Boolean if generator is being used for training
    :return: The tensor output of the generator
    """
    # 需要编程：
    alpha = 0.2
    reuse = False if is_train == True else True
    with tf.variable_scope('generator', reuse=reuse):
        # FC 4*4*512  and reshape
        x1 = tf.layers.dense(z, 4 * 4 * 512, activation=None)
        x1 = tf.reshape(x1, [-1, 4, 4, 512])
        x1 = tf.layers.batch_normalization(x1, training=is_train)
        x1 = tf.maximum(x1 * alpha, x1)
        x1 = tf.nn.dropout(x1, 0.8)
        # 4*4*512  to 7*7*256
        x1 = tf.layers.conv2d_transpose(x1, 256, 4, strides=1, padding='valid',
                                        kernel_initializer=tf.contrib.layers.xavier_initializer())
        x1 = tf.layers.batch_normalization(x1, training=is_train)
        x1 = tf.maximum(x1 * alpha, x1)
        x1 = tf.nn.dropout(x1, 0.8)
        # 7*7*256 to 14*14*128
        x2 = tf.layers.conv2d_transpose(x1, 128, 5, strides=2, padding='same',
                                        kernel_initializer=tf.contrib.layers.xavier_initializer())
        x2 = tf.layers.batch_normalization(x2, training=is_train)
        x2 = tf.maximum(x2 * alpha, x2)
        x2 = tf.nn.dropout(x2, 0.8)
        # 14*14*128 to 28*28*out_channel
        logits = tf.layers.conv2d_transpose(x2, out_channel_dim, 5, strides=2, padding='same',
                                            kernel_initializer=tf.contrib.layers.xavier_initializer())
        out = tf.tanh(logits)

    return out

# todo-构建损失函数
"""
部署 `model_loss` 函数训练并计算 GANs 的损失。该函数应返回形如 (discriminator loss, generator loss) 的元组。
"""

def model_loss(input_real, input_z, out_channel_dim):
    """
    Get the loss for the discriminator and generator
    :param input_real: Images from the real dataset
    :param input_z: Z input
    :param out_channel_dim: The number of channels in the output image
    :return: A tuple of (discriminator loss, generator loss)
    """
    # todo-需要编程
    smooth = 0.1
    g_model = generator(input_z, out_channel_dim, is_train=True)
    d_model_real, d_logits_real = discriminator(input_real)
    d_model_fake, d_logits_fake = discriminator(g_model, reuse=True)

    d_loss_real = tf.reduce_mean(
        tf.nn.sigmoid_cross_entropy_with_logits(
            logits=d_logits_real, labels=tf.ones_like(d_model_real) * (1 - smooth)))
    d_loss_fake = tf.reduce_mean(
        tf.nn.sigmoid_cross_entropy_with_logits(
            logits=d_logits_fake, labels=tf.zeros_like(d_model_fake)))
    g_loss = tf.reduce_mean(
        tf.nn.sigmoid_cross_entropy_with_logits(
            logits=d_logits_fake, labels=tf.ones_like(d_model_fake)))
    d_loss = d_loss_real + d_loss_fake

    return d_loss, g_loss

# todo-构建优化器
"""
部署 `model_opt` 函数实现对 GANs 的优化。使用 [`tf.trainable_variables`] 获取可训练的所有变量。
通过变量空间名 `discriminator` 和 `generator` 来筛选变量。
该函数应返回形如 (discriminator training operation, generator training operation)。
"""

def model_opt(d_loss, g_loss, learning_rate, beta1):
    """
    Get optimization operations
    :param d_loss: Discriminator loss Tensor
    :param g_loss: Generator loss Tensor
    :param learning_rate: Learning Rate Placeholder
    :param beta1: The exponential decay rate for the 1st moment in the optimizer
    :return: A tuple of (discriminator training operation, generator training operation)
    """
    # todo 需要编程：
    t_vars = tf.trainable_variables()
    d_vars = [var for var in t_vars if var.name.startswith('discriminator')]
    g_vars = [var for var in t_vars if var.name.startswith('generator')]
    with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS)):
        d_train_opt = tf.train.AdamOptimizer(learning_rate,
                                             beta1=beta1).minimize(d_loss, var_list=d_vars)
        g_train_opt = tf.train.AdamOptimizer(learning_rate,
                                             beta1=beta1).minimize(g_loss, var_list=g_vars)

    return d_train_opt, g_train_opt

# 输出显示
# 使用该函数可以显示生成器 (Generator) 在训练过程中的当前输出，帮助评估 GANs 模型。
def show_generator_output(sess, n_images, input_z, out_channel_dim, image_mode):
    """
    Show example output for the generator
    :param sess: TensorFlow session
    :param n_images: Number of Images to display
    :param input_z: Input Z Tensor
    :param out_channel_dim: The number of channels in the output image
    :param image_mode: The mode to use for images ("RGB" or "L")
    """
    pyplot.ion()
    cmap = None if image_mode == 'RGB' else 'gray'
    z_dim = input_z.get_shape().as_list()[-1]
    example_z = np.random.uniform(-1, 1, size=[n_images, z_dim])
    samples = sess.run(
        generator(input_z, out_channel_dim, False),
        feed_dict={input_z: example_z})

    images_grid = helper.images_square_grid(samples, image_mode)
    pyplot.imshow(images_grid, cmap=cmap)
    pyplot.show()
    pyplot.pause(2)
    pyplot.close()

# todo-训练
"""
使用 `show_generator_output` 函数显示 `generator` 在训练过程中的输出。
**注意**：在每个批次 (batch) 中运行 `show_generator_output` 函数会显著增加训练时间。推荐每 100 批次输出一次 `generator` 的输出。 
"""

def train(epoch_count, batch_size, z_dim, learning_rate, beta1,
          get_batches, data_shape, data_image_mode):
    """
    Train the GAN
    :param epoch_count: Number of epochs
    :param batch_size: Batch Size
    :param z_dim: Z dimension
    :param learning_rate: Learning Rate
    :param beta1: The exponential decay rate for the 1st moment in the optimizer
    :param get_batches: Function to get batches
    :param data_shape: Shape of the data
    :param data_image_mode: The image mode to use for images ("RGB" or "L")
    """
    # todo-需要编程
    image_width = data_shape[-3]
    image_height = data_shape[-2]
    image_channels = data_shape[-1]
    inputs_real, inputs_z, lr = model_inputs(image_width, image_height, image_channels, z_dim)
    d_loss, g_loss = model_loss(inputs_real, inputs_z, image_channels)
    d_train_opt, g_train_opt = model_opt(d_loss, g_loss, learning_rate, beta1)

    # saver = tf.train.Saver()
    # sample_z = np.random.uniform(-1, 1, size=(72, z_dim))
    # samples, losses = [], []
    steps = 0
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        for epoch_i in range(epoch_count):
            for batch_images in get_batches(batch_size):
                # 需要编程：
                batch_images = batch_images * 2
                steps += 1
                # 生成随机噪音样本，传入G
                batch_z = np.random.uniform(-1, 1, size=(batch_size, z_dim))
                # 执行opt
                _ = sess.run(d_train_opt, feed_dict={inputs_real: batch_images, inputs_z: batch_z,
                                                     lr: learning_rate})
                _ = sess.run(g_train_opt, feed_dict={inputs_z: batch_z, inputs_real: batch_images,
                                                     lr: learning_rate})

                if steps % 10 == 0:
                    # 获取损失并打印。
                    train_loss_d = d_loss.eval({inputs_z: batch_z, inputs_real: batch_images})
                    train_loss_g = g_loss.eval({inputs_z: batch_z})

                    print("Epoch {}/{}...".format(epoch_i + 1, epoch_count),
                          "Discriminator Loss: {:.4f}...".format(train_loss_d),
                          "Generator Loss: {:.4f}".format(train_loss_g))
                    # Save losses to view after training
                    # losses.append((train_loss_d, train_loss_g))

                if steps % 100 == 0:
                    _ = show_generator_output(sess, 25, inputs_z, image_channels, data_image_mode)
        # saver.save(sess, './checkpoints/generator.ckpt')

# todo-在 MNIST 上测试你的 GANs 模型。
# 经过 2 次迭代，GANs 应该能够生成类似手写数字的图像。确保生成器 (generator) 低于辨别器 (discriminator) 的损失，
def train_minst():
    batch_size = 128
    z_dim = 100
    learning_rate = 0.001
    beta1 = 0.5
    epochs = 5

    mnist_dataset = helper.Dataset('mnist', glob(os.path.join(data_dir, 'mnist/*.jpg')))
    with tf.Graph().as_default():
        train(epochs, batch_size, z_dim, learning_rate, beta1, mnist_dataset.get_batches,
              mnist_dataset.shape, mnist_dataset.image_mode)

# todo-在CelebA上运行你的 GANs 模型
# 在一般的GPU上运行每次迭代大约需要 20 分钟。你可以运行整个迭代，或者当 GANs 开始产生真实人脸图像时停止它。
def train_CelebA():
    # 超参数调优目标：D_loss 为0.3左右时（真假难辨），生成图片效果最佳
    batch_size = 8
    z_dim = 100
    learning_rate = 0.0001
    beta1 = 0.5
    alpha = 0.01
    epochs = 1

    celeba_dataset = helper.Dataset('celeba', glob(os.path.join(data_dir, 'img_align_celeba/*.jpg')))
    with tf.Graph().as_default():
        train(epochs, batch_size, z_dim, learning_rate, beta1, celeba_dataset.get_batches,
              celeba_dataset.shape, celeba_dataset.image_mode)
"""
# 超参数建议：D_loss 为左右时（真假难辨），生成图片效果最佳
batch_size = 32,64,128,256,512
z_dim = 100,200
learning_rate = 0.01,0.001,0.0001
beta1 = 0.5
alpha=0.01--0.2
epochs=1
"""

if __name__=='__main__':
    # explore_mnist()
    # explore_CelebA()
    # train_minst()
    train_CelebA()