GAN

#coding:utf-8

from tensorflow.examples.tutorials.mnist import input_data #连接远程服务器，下载数据集
import tensorflow as tf
import numpy as np
from skimage.io import imsave
import os #获取路径
import shutil
import sys
import six

#############
#定义模型参数
#############
# 图片尺寸
img_height= 28
img_width = 28
img_size = img_height * img_width#784

###################
#模型保存与加载参数
###################
to_train = True
to_restore = False #为了判断是否加载模型
output_path = 'D:\Python\GANTest\output'# 保存地址


###################
#网络及训练参数
###################
#生成模型的输入是100，随机变量的维度是100，
# 最后要变成784维，所以隐层是增的，第一个隐层是150，第二个是300，这是生成模型
# 判别模型要倒着来，输入是784，第一层300，第二层150，最后的输出是1维，映射到一个值，用0或者1全宇判别
z_size = 100
h1_size = 150
h2_size = 300
# 一次训练256个，因为图像像素也不高，可以一次多训练一点
batch_size = 256
# 深度学习一般都是1000，实际上收敛的时候，一般小于1000
max_epoch = 1000

###################
#DNN网络构建（多层感知机）
###################
# 生成模型做的： 就是给一个随机变量，返回一个生成的数值x_generate，和一套参数
def 生成模型(z_prior):
    # 从输入层向第一隐层的w，应该是100行，150列

    # truncated_normal：从截断的正态分布中输出随机值。
    # shape表示生成张量的维度，mean是均值，stddev是标准差。这个函数产生正太分布，均值和标准差自己设定。
    # 这是一个截断的产生正太分布的函数，就是说产生正太分布的值如果与均值的差值大于两倍的标准差，那就重新生成
    w1 = tf.Variable(tf.truncated_normal([z_size, h1_size], \
                                         # 方差 stddev = 0.1，均值不用写，均值默认是0
                                         stddev = 0.1), name = 'g_w1', \
                                         dtype=tf.float32)
    # 注意此时是：zeros
    b1 = tf.Variable(tf.zeros([h1_size]), name = 'g_b1', dtype = tf.float32)
    # 公式上是Wx+b,但矩阵运算时，x放前面
    h1 = tf.nn.relu(tf.matmul(z_prior, w1) + b1)

    # 从第一个隐层到第二个隐层
    w2 = tf.Variable(tf.truncated_normal([h1_size, h2_size], \
                                         stddev = 0.1), name = 'g_w2', \
                                         dtype=tf.float32)
    b2 = tf.Variable(tf.zeros([h2_size]), name = 'g_b2', dtype = tf.float32)

    # 注意此时是h1，也就是 隐层1激活之后的结果值
    h2 = tf.nn.relu(tf.matmul(h1, w2) + b2)

    w3 = tf.Variable(tf.truncated_normal([h2_size, img_size], \
                                         stddev = 0.1), name = 'g_w3', \
                                         dtype=tf.float32)
    b3 = tf.Variable(tf.zeros([img_size]), name = 'g_b3', dtype = tf.float32)
    # h3也能很大，也可能很小，如果是做分类，就用softmax激活，这里不用分类，选tanh激活
    h3 = tf.matmul(h2, w3) + b3
    x_generate = tf.nn.tanh(h3)#fake，生成模型，假的
    g_params = [w1, b1, w2, b2, w3, b3]#生成模型的参数，更新的参数，需要保存
    return x_generate, g_params

# x_data: 就是通过input_data下载的数据集，在后面的会话中给他喂，先写一个形参x_data就是y值
# x_generate：传两个
#  keep_prob：Dropout率，意思是每个元素被保留的概率，那么 keep_prob:1就是所有元素全部保留的意思。
# 一般在大量数据训练时，为了防止过拟合，添加Dropout层，设置一个0~1之间的小数
def 判别模型(x_data, x_generate, keep_prob):
    # 这两个值是并行输入的
    # tf.concat：tensorflow中用来拼接张量的函数。 tf.concat()拼接的张量只会改变一个维度，其他维度是保存不变的。
    # 比如两个shape为[2,3]的矩阵拼接，要么通过axis=0变成[4,3]，要么通过axis=1变成[2,6]。改变的维度索引对应axis的值。
    # x_data（y值）和x_generate（y^值）都是一维的，所以拼接成两行的向量，但是他们不是相加，一起进网络，一起出网络
    # x_in代表判别模型的输入
    x_in = tf.concat([x_data, x_generate], 0)
    w1 = tf.Variable(tf.truncated_normal([img_size, h2_size], \
                                         stddev = 0.1), name = 'd_w1', \
                                         dtype=tf.float32)
    b1 = tf.Variable(tf.zeros([h2_size]), name = 'd_b1', dtype = tf.float32)
    # 这里比之前可以加一个dropout率
    h1 = tf.nn.dropout(tf.nn.relu(tf.matmul(x_in, w1) + b1),keep_prob)

    w2 = tf.Variable(tf.truncated_normal([h2_size, h1_size], \
                                         stddev = 0.1), name = 'd_w2', \
                                         dtype=tf.float32)
    b2 = tf.Variable(tf.zeros([h1_size]), name = 'd_b2', dtype = tf.float32)
    h2 = tf.nn.dropout(tf.nn.relu(tf.matmul(h1, w2) + b2),keep_prob)
    # w3只有一个值，注意b和后面的网络深度相同，因为前一层的b连的是后一层的cell
    w3 = tf.Variable(tf.truncated_normal([h1_size, 1], stddev = 0.1), \
                                        name = 'd_w3', dtype = tf.float32)
    b3 = tf.Variable(tf.zeros([1]), name = 'd_b3', dtype = tf.float32)
    h3 = tf.matmul(h2, w3) + b3
    print(h3)
    # 两维数组，上面一行是1*256(因为一次训练256个)
    # tf.slice(input_, begin, size)：
    # “input_”是你输入的tensor，就是被切的那个。
    # “begin”是每一个维度的起始位置，这个下面详细说。
    # “size”相当于问每个维度拿几个元素出来。
    # 需要把y和y^分开，因为之前是黏在一起的，然后算损失
    #  为什么[batch_size, -1]是[256, 784]？累加？全连接层

    y_data = tf.nn.sigmoid(tf.slice(h3, [0, 0], [batch_size, -1]), name = None) #y
    # 二分类直接用sigmoid去激活
    # [batch_size, 0]中batch_size：从第256开始切，[-1, -1]：第一个-1，切剩下所有的
    y_generated = tf.nn.sigmoid(tf.slice(h3, [batch_size, 0], [-1, -1]), name = None) #y^
    d_params = [w1, b1, w2, b2, w3, b3]#判别模型的参数
    return y_data, y_generated, d_params

#########################################################
#定义图片展示的函数(即将生成模型的输出相片绘制成图片保存)
#########################################################
# 绘制一个8*8的格子，每个格子保存16张图片，每张图片就是28*28=784的图片
def 展示结果保存(batch_res, fname, grid_size=(8, 8), grid_pad=5):#show_result
    '''图片保存'''
    batch_res = 0.5 * batch_res.reshape((batch_res.shape[0], img_height, img_width)) + 0.5
    img_h, img_w = batch_res.shape[1], batch_res.shape[2]
    grid_h = img_h * grid_size[0] + grid_pad * (grid_size[0] - 1)
    grid_w = img_w * grid_size[1] + grid_pad * (grid_size[1] - 1)
    img_grid = np.zeros((grid_h, grid_w), dtype=np.uint8)
    for i, res in enumerate(batch_res):
        if i >= grid_size[0] * grid_size[1]:
            break
        img = (res) * 255
        img = img.astype(np.uint8)
        row = (i // grid_size[0]) * (img_h + grid_pad)
        col = (i % grid_size[1]) * (img_w + grid_pad)
        img_grid[row:row + img_h, col:col + img_w] = img
    imsave(fname, img_grid)

######################
#开始训练
######################
def 开始训练():
    # input_data是网上下载的，one_hot打开，one_hot指的是标签
    mnist = input_data.read_data_sets('MNIST_data', one_hot = True)
    # [batch_size, z_size]=[256,100], z_prior:输入给生成模型的,随机变量
    z_prior = tf.placeholder(tf.float32, [batch_size, z_size], name = 'z_prior')
    # [batch_size, img_size]=[256,784]   x_data：真实值
    x_data = tf.placeholder(tf.float32, [batch_size, img_size], name = 'x_data')
    keep_prob = tf.placeholder(tf.float32, name = 'keep_prob')
    # 步数，假设255步，trainable：不是训练的步长，是测试的步长
    global_step = tf.Variable(255, name = 'global_step', trainable = False)

    x_generate, g_params = 生成模型(z_prior)
    y_data, y_generated, d_params = 判别模型(x_data, x_generate, keep_prob)

    ##########
    #构建损失，两个损失，先更新d_loss，注意损失函数的区别
    ##########
    d_loss = -(tf.log(y_data)+tf.log(1-y_generated))
    g_loss = -tf.log(y_generated)
    # 优化器，要优化两次
    optimizer = tf.train.AdamOptimizer(0.0001)
    # d_params两个模型中的参数列表，run这两个模型就好了：d_train，g_train
    d_train = optimizer.minimize(d_loss, var_list = d_params)
    g_train = optimizer.minimize(g_loss, var_list = g_params)
    # 之前的写法也可以
    init = tf.initialize_all_variables()
    #####################
    #保存与加载模型
    #####################
    saver = tf.train.Saver()
    sess = tf.Session()
    # 先run(init)，初始化变量
    sess.run(init)
    if to_restore:#加载模型
        chkpt_fname = tf.train.latest_checkpoint(output_path)
        # 把模型及上面的参数加载到会话中来
        saver.restore(sess, chkpt_fname)
    elif os.path.exists(output_path):
        # 存在，则删掉，重新创建
        shutil.rmtree(output_path)
        os.mkdir(output_path)
    # os.path.exists(output_path)
    # # 存在，则删掉，重新创建
    # shutil.rmtree(output_path)
    # os.mkdir(output_path)

    # 创建随机变量 0, 1，均值为0，方差为1，(batch_size, z_size) 256行*100维
    # z_sample_val = np.random.normal(0, 1, size = (batch_size, z_size)).astype(tf.float32)

    # 60000整个数据集有60000张图片
    steps = 60000//batch_size

    #for i in range(sess.run(global_step), max_epoch):
    for i in range(max_epoch):

        for j in range(steps):

            print("epoch序次:%d, 迭代次数:%d"%(i, j))

            ############
            #传入实参
            ############
            # 切batch_size个
            x_value, _ = mnist.train.next_batch(batch_size)
            # 归一化处理，图像一般都这样处理，*2-1，让像素分布在0的两边
            x_value = 2*x_value.astype(np.float32) - 1
            # x_value：真实值。z_value
            z_value = np.random.normal(0, 1, size = (batch_size, z_size)).astype(np.float32)

            # 先run判别模型
            sess.run(d_train, feed_dict = {x_data:x_value, z_prior:z_value, keep_prob:np.sum(0.7).astype(np.float32)})

            sess.run(g_train, feed_dict = {x_data:x_value, z_prior:z_value, keep_prob:np.sum(0.7).astype(np.float32)})
            # 每运行一遍，再运行x_generate
            x_gen_val = sess.run(x_generate, feed_dict={z_prior:z_value})

            展示结果保存(x_gen_val, 'D:\Python\GANTest\output_sample\sample{0}.jpg'.format(i))
            ################
            #保存模型
            ################
            sess.run(tf.assign(global_step, i+1))
            saver.save(sess, os.path.join(output_path, 'model'), global_step = global_step)

if __name__ =='__main__':

    开始训练()