import tensorflow as tf

parser.add_argument('--model_path', type=str, default='', help='the path of checkpoint file')
args = parser.parse_args()

def model():
    x = tf.placeholder(tf.float32, [None, 784], name='x')
    gt = tf.placeholder(tf.float32, [None, 10], name='groundtruth')
    with tf.variable_scope('layer1'):
        w1 = tf.get_variable('weight1', [784, 1024], initializer=tf.random_normal_initializer())
        b1 = tf.get_variable('bias1', [1024], initializer=tf.constant_initializer(0.0))
        h1 = tf.nn.relu(tf.matmul(x, w1) + b1)
    with tf.variable_scope('layer2'):
        w2 = tf.get_variable('weight2', [1024, 1024], initializer=tf.random_normal_initializer())
        b2 = tf.get_variable('bias2', [1024], initializer=tf.constant_initializer(0.0))
        h2 = tf.nn.relu(tf.matmul(h1, w2) + b2)
    with tf.variable_scope('layer3'):
        w3 = tf.get_variable('weight3', [1024, 10], initializer=tf.random_normal_initializer())
        b3 = tf.get_variable('bias3', [10], initializer=tf.constant_initializer(0.0))
        y = tf.matmul(h2, w3) + b3
    # losses
    cross_entropy = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(labels=gt, logits=y))
    # optimizer
    optimizer = tf.train.GradientDescentOptimizer(args.lr)
    # define one-step train ops
    train_op = optimizer.minimize(cross_entropy)
    return x, y, gt, train_op    
    
if __name__ == "__main__":
    max_train_step = args.max_train_step
