import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data
mnist = input_data.read_data_sets('MNIST_data',one_hot=True)

class Net:
    def __init__(self):
        self.x = tf.placeholder(dtype=tf.float32,shape=[None,28,28,1])
        self.y = tf.placeholder(dtype=tf.float32,shape=[None,10])
        self.conv1_w = tf.Variable(tf.truncated_normal([3,3,1,8],dtype=tf.float32,stddev=0.1))
        self.conv1_b = tf.Variable(tf.zeros([8],dtype=tf.float32))
        self.conv2_w = tf.Variable(tf.truncated_normal([3,3,8,16],dtype=tf.float32,stddev=0.1))
        self.conv2_b = tf.Variable(tf.zeros([16],dtype=tf.float32))
        self.w = tf.Variable(tf.truncated_normal(dtype=tf.float32,shape=[7*7*16,10],stddev=0.1))
        self.b = tf.Variable(tf.zeros([10]))
    def forward(self):
        self.conv1 = tf.nn.relu(tf.nn.conv2d(self.x,self.conv1_w,strides=[1,1,1,1],padding='SAME')+self.conv1_b)
        self.pool1 = tf.nn.max_pool(self.conv1,ksize=[1,2,2,1],strides=[1,2,2,1],padding='SAME')
        self.conv2 = tf.nn.relu((tf.nn.conv2d(self.pool1,self.conv2_w,strides=[1,1,1,1],padding='SAME')+self.conv2_b))
        self.pool2 = tf.nn.max_pool(self.conv2,ksize=[1,2,2,1],strides=[1,2,2,1],padding='SAME')
        self.flat = tf.reshape(self.pool2,[-1,7*7*16])
        self.yo = tf.nn.softmax(tf.matmul(self.flat,self.w)+self.b)
    def backward(self):
        self.loss = tf.reduce_mean((self.y-self.yo)**2)
        self.opt = tf.train.AdamOptimizer().minimize(self.loss)
def zql(y,y_hat):
    return tf.reduce_mean(tf.cast(tf.equal(tf.argmax(y,1),tf.argmax(y_hat,1)),'float'))
if __name__ == '__main__':
    net = Net()
    net.forward()
    net.backward()
    init = tf.global_variables_initializer()
    with tf.Session() as sess:
        sess.run(init)
        for i in range(1000):
            x,y = mnist.train.next_batch(100)
            x = x.reshape([100,28,28,1])
            loss,_ = sess.run([net.loss,net.opt],feed_dict={net.x:x,net.y:y})
            if i%10 == 0:
                x_test,y_test = mnist.test.next_batch(100)
                x_test = x_test.reshape([100,28,28,1])
                out = sess.run(net.yo,feed_dict={net.x:x_test})
                print(sess.run(zql(y_test,out)))