import tensorflow as tf

import numpy as np

from tensorflow.contrib import rnn

from tensorflow.examples.tutorials.mnist import input_data

import matplotlib.pyplot as plt

# 第一步数据读取

mnist = input_data.read_data_sets('/data', one_hot=True)

# 第二步：超参数设置

lr = 1e-3 # 学习率的设置

input_num = 28

timestamp_size = 28

lstm_num = 2

num_hidden = 256

num_classes = 10

# 第三步：输入参数的初始化操作

_X = tf.placeholder(tf.float32, [None, 784])

y = tf.placeholder(tf.float32, [None, num_classes])

batch_size = tf.placeholder(tf.int32, [])

keep_prob = tf.placeholder(tf.float32, [])

# 改变输入的维度，变成N*28*28

X = tf.reshape(_X, [-1, timestamp_size, input_num])

# 第四步：构造函数创建多层的LSTM

def create_rnn_layer():

    cell = rnn.LSTMCell(num_hidden, reuse=tf.get_variable_scope().reuse)

    return rnn.DropoutWrapper(cell, output_keep_prob=keep_prob)

# 第五步：使用tf.contrib.rnn.MultiRNNCell

mlstm_cell = tf.contrib.rnn.MultiRNNCell([create_rnn_layer() for _ in range(lstm_num)], state_is_tuple=True)

# 第六步：使用mlstm_cell.zero_state构造初始化的state

init_state = mlstm_cell.zero_state(batch_size, dtype=tf.float32)

# 第七步：对每一个输入的数据进行迭代，获得state，并作为下一次的输入，重新输入到LSTM网络中，最后的输出即为预测结果

outputs = list()

state = init_state

with tf.variable_scope('RNN'):

    for timestamp in range(timestamp_size):

        if timestamp > 0:

            tf.get_variable_scope().reuse_variables()

        (output, state) = mlstm_cell(X[:, timestamp, :], state)

        outputs.append(output)

# 第八步：使用最后一个结果，作为输出值

h_state = outputs[-1]

# 第九步：定义W和b，用于将LSTM的输出变成预测的输出

W = tf.Variable(tf.truncated_normal([num_hidden, num_classes], stddev=0.1), dtype=tf.float32)

b = tf.Variable(tf.constant(0.1, shape=[num_classes]), dtype=tf.float32)

# 第十步：使用tf.matmul进行线性变换，使用softmax将结果转换为概率值

y_pre = tf.nn.softmax(tf.matmul(h_state, W) + b)

# 第十一步：使用tf.log(y_pre) * y 获得交叉熵损失函数

loss = -tf.reduce_mean(tf.log(y_pre) * y)

# 第十二步：使用tf.train.Adam降低损失值

train_op = tf.train.AdamOptimizer(lr).minimize(loss)

# 第十三步：使用tf.equal和tf.reduce_mean 获得准确率

correct_pred = tf.equal(tf.argmax(y_pre, 1), tf.argmax(y, 1))

accurracy = tf.reduce_mean(tf.cast(correct_pred, 'float'))

# 进行初始化操作

sess = tf.Session()

sess.run(tf.global_variables_initializer())

# 一个batch_size的大小

_batch_size = 64

for i in range(2000):

    # 循环，第十四步：获得一个batch的数据

    batch = mnist.train.next_batch(_batch_size)

    # 第十五步：训练train_op, keep_prob等于0.5， 当迭代次数为200时，打印训练的准确度

    if i % 200 == 0:

        train_accur = sess.run(accurracy, feed_dict={_X:batch[0],

                                                       y:batch[1],

                                                       keep_prob:1.0,

                                                       batch_size:_batch_size})

        print('train epoch %d train iteration %d train accur %g'%(mnist.train.epochs_completed, i, train_accur))

    sess.run(train_op, feed_dict={_X:batch[0], y:batch[1], keep_prob:0.5, batch_size:_batch_size})

# 第十六步：打印测试集的准确率

print('test accur %g'%(sess.run(accurracy, feed_dict={_X:mnist.test.images, y:mnist.test.labels, keep_prob:1.0, batch_size:mnist.test.images.shape[0]})))

# 第二部分：打印outputs的维度

# 第一步：定义输入的_batch_size

_batch_size = 5

# 第二步：获得mnist中Train的数据

X_batch, y_batch = mnist.train.next_batch(_batch_size)

# 第三步：使用sess.run获得，outputs的实际值

_outputs, state = sess.run([outputs, state], feed_dict={_X:X_batch, y:y_batch, keep_prob:1.0, batch_size:_batch_size})

# 第四步：打印_outputs的维度

print('outputs shape=', np.asarray(_outputs).shape)

# 第三部分：对每一部分的输出值进行概率的作图

# 第一步：图片的输入，即将图片的维度进行变换，对应标签的输入，对图片的标签进行维度变换

image = mnist.train.images[4]

image.shape = [-1, 784]

image_size = image.reshape([28, 28])

plt.imshow(image_size, cmap='gray')

plt.show()

y_labels = mnist.train.labels[4]

y_labels.shape = [-1, num_classes]

# 第二步：使用sess.run(outputs)获得outputs的输出值

output_array = np.array(sess.run(outputs, feed_dict={_X:image, y:y_labels, keep_prob:1.0, batch_size:1}))

# 对outputs转换为二维类型，即[28, 1, 256] 转换为[28, 256]

output_array.shape = [28, 256]

# 第三步：使用sess.run获得W和b的实际值

h_W = sess.run(W, feed_dict={_X:image, y:y_labels, keep_prob:1.0, batch_size:1})

h_b = sess.run(b, feed_dict={_X:image, y:y_labels, keep_prob:1.0, batch_size:1})

# 构建bar_index的X轴坐标

bar_index = np.arange(num_classes)

plt.figure()

for i in range(output_array.shape[0]):

    plt.subplot(4, 7, i+1)

    # 将[1, 256]重新转换为[1, 256]做一个维度的验证

    X3_h_state = output_array[i, :].reshape([-1, num_hidden])

    # 第四步：使用tf.nn.softmax将输出值转换为各个类别的概率值

    pre = sess.run(tf.nn.softmax(tf.matmul(X3_h_state, h_W) + h_b))

    # 第五步：画直方图操作

    plt.bar(bar_index, pre[0], width=0.2, align='center')

    plt.axis('off')

plt.show()