# -*- coding: utf-8 -*-
"""
@Time ： 2021/8/14 上午2:10
@Auth ： 陈伟峰
@File ：mnist_tensorflow.py
@phone: 15882085601
@IDE ：PyCharm
@Motto：ABC(Always Be Coding)

"""
import tensorflow as tf
import numpy as np
from tensorflow.examples.tutorials.mnist import input_data
import os
# from tensorflow.contrib.learn.python.learn.datasets.mnist import read_data_sets
os.environ["CUDA_VISIBLE_DEVICES"] = '0' #use GPU with ID=0
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.5 # maximun alloc gpu50% of MEM
config.gpu_options.allow_growth = True #allocate dynamically


mnist = input_data.read_data_sets("mnist_data/",one_hot=True)
trX,trY,teX,teY = mnist.train.images,mnist.train.labels,mnist.test.images,mnist.test.labels


trX = trX.reshape(-1,28,28,1)
teX = teX.reshape(-1,28,28,1)

# print(trX.shape)
# print(trY.shape)


X = tf.placeholder(tf.float32,[None,28,28,1])
Y = tf.placeholder(tf.float32,[None,10])

p_keep_conv = tf.placeholder(tf.float32)
p_keep_hidden = tf.placeholder(tf.float32)



def init_weights(shape):
    return tf.Variable(tf.random_normal(shape,stddev=0.01))

# kernel_size = (3,3),input_size,out_size
w = init_weights([3,3,1,32])
w2 = init_weights([3,3,32,64])
w3 = init_weights([3,3,64,128])
# 全链接
w4 = init_weights([128*4*4,625])
# 输出层
w_o = init_weights([625,10])



def model(X,w,w2,w3,w4,w_o,p_keep_conv,p_keep_hidden):

    result = tf.nn.conv2d(X,w,strides=[1,1,1,1],padding="SAME")
    result = tf.nn.relu(result)
    result = tf.nn.max_pool(result,ksize=[1,2,2,1],strides=[1,2,2,1],padding="SAME")
    result = tf.nn.dropout(result,p_keep_conv)

    result2 = tf.nn.conv2d(result, w2, strides=[1, 1, 1, 1], padding="SAME")
    result2 = tf.nn.relu(result2)
    result2 = tf.nn.max_pool(result2, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding="SAME")
    result2 = tf.nn.dropout(result2, p_keep_conv)


    result3 = tf.nn.conv2d(result2, w3, strides=[1, 1, 1, 1], padding="SAME")
    result3 = tf.nn.relu(result3)
    result3 = tf.nn.max_pool(result3, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding="SAME")
    result3 = tf.nn.dropout(result3, p_keep_conv)

    # print(result3.shape)
    # print(w4)
    # exit()
    #
    result4 = tf.reshape(result3,[-1,w4.get_shape().as_list()[0]])
    result4 = tf.nn.dropout(result4,p_keep_conv)

    result4 = tf.nn.relu(tf.matmul(result4,w4))
    result4 = tf.nn.dropout(result4,p_keep_hidden)

    out = tf.matmul(result4,w_o)
    return out


out = model(X,w,w2,w3,w4,w_o,p_keep_conv,p_keep_hidden)
# 损失函数
cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=out,labels=Y))
train_op = tf.train.GradientDescentOptimizer(learning_rate=0.001).minimize(cost)

predict_op = tf.argmax(out,axis=1)


batch_size = 128
test_size = 256


with tf.Session(config=config) as sess:

    # 全局初始化
    tf.global_variables_initializer().run()
    for i in range(100):
        train_batch = zip(range(0,len(trX),batch_size),
                          range(batch_size,len(trX)+1,batch_size)
                          )
        for start,end in train_batch:
            # print(trX[start:end].shape)
            # print(trY[start:end].shape)
            # exit()
            sess.run(train_op,feed_dict={
                X:trX[start:end],
                Y:trY[start:end],
                p_keep_conv:0.8,
                p_keep_hidden:0.5
            })
    test_indices = np.arange(len(teX))

    np.random.shuffle(test_indices)

    test_indices = test_indices[0:test_size]
    print(i,np.mean(np.argmax(teY[test_indices],axis=1)==sess.run(predict_op,feed_dict={
        X:teX[test_indices],
        p_keep_conv:1.0,
        p_keep_hidden:1.0
    })))