mnist_forward

import tensorflow as tf
import numpy as np

IMGSIZE = 28
IMG_CH = 1
CONV_1_W_SIZE = 5
CONV_1_KER_NUM = 32
CONV_2_W_SIZE = 5
CONV_2_KER_NUM = 64
FC_1_NUM = 512
OUT_NUM = 10


def get_weight(shape,reg,istrain):
    w = tf.Variable(tf.truncated_normal(shape,stddev= 0.1))
    if istrain :
        tf.add_to_collection('losses',tf.contrib.layers.l2_regularizer(reg)(w))
    return w

def get_bias(shape):
    return tf.Variable(tf.zeros(shape))

def conv2d_relu_maxpool(x,w,b):
    y = tf.nn.conv2d(x,w,strides=[1,1,1,1],padding='SAME')
    relu1 = tf.nn.relu(tf.nn.bias_add(y,b))
    pool1 = tf.nn.max_pool(relu1,ksize=[1,2,2,1],strides=[1,2,2,1],padding='SAME')
    return pool1



def forward(x,reg,istrain):
    conv_w1 = get_weight([CONV_1_W_SIZE,CONV_1_W_SIZE,IMG_CH,CONV_1_KER_NUM],reg,istrain)
    conv_b1 = get_bias([CONV_1_KER_NUM])
    pool1 = conv2d_relu_maxpool(x,conv_w1,conv_b1)

    conv_w2 = get_weight([CONV_2_W_SIZE,CONV_2_W_SIZE,CONV_1_KER_NUM,CONV_2_KER_NUM],reg,istrain)
    conv_b2 = get_bias([CONV_2_KER_NUM])
    pool2 = conv2d_relu_maxpool(pool1, conv_w2, conv_b2)

    shpae = pool2.get_shape().as_list()
    NODES =shpae[1] * shpae[2]*shpae[3]
    shpd = tf.reshape(pool2,[shpae[0],NODES])

    fc_w1 = get_weight([NODES,FC_1_NUM],reg,istrain)
    fc_b1 = get_bias([FC_1_NUM])
    fc1 = tf.nn.relu(tf.matmul(shpd,fc_w1)+fc_b1)
    if istrain:
        fc1 = tf.nn.dropout(fc1,keep_prob=0.5)

    fc_w2 = get_weight([FC_1_NUM,OUT_NUM],reg,istrain)
    fc_b2 = get_bias([OUT_NUM])
    fc2 = tf.matmul(fc1,fc_w2)+fc_b2
    return fc2