import tensorflow as tf

# 1. create data
from tensorflow.examples.tutorials.mnist import input_data
mnist = input_data.read_data_sets(‘../MNIST_data‘, one_hot=True)

X = tf.placeholder(tf.float32, shape=(None, 784), name=‘X‘)
y = tf.placeholder(tf.int32, shape=(None), name=‘y‘)
is_training = tf.placeholder(tf.bool, None, name=‘is_training‘)

# 2. define network
he_init = tf.contrib.layers.variance_scaling_initializer()
with tf.name_scope(‘dnn‘):
    hidden1 = tf.layers.dense(X, 300, kernel_initializer=he_init, name=‘hidden1‘,
                              kernel_regularizer=tf.contrib.layers.l2_regularizer(0.001))
    hidden1 = tf.layers.batch_normalization(hidden1, momentum=0.9)
    hidden1 = tf.nn.relu(hidden1)
    hidden2 = tf.layers.dense(hidden1, 100, kernel_initializer=he_init, name=‘hidden2‘,
                              kernel_regularizer=tf.contrib.layers.l2_regularizer(0.001))
    hidden2 = tf.layers.batch_normalization(hidden2, training=is_training, momentum=0.9)
    hidden2 = tf.nn.relu(hidden2)
    logits = tf.layers.dense(hidden2, 10, kernel_initializer=he_init, name=‘output‘,
                             kernel_regularizer=tf.contrib.layers.l2_regularizer(0.001))

# 3. define loss
with tf.name_scope(‘loss‘):
    loss = tf.losses.softmax_cross_entropy(onehot_labels=y, logits=logits) # label is one_hot
    # =================
    reg_losses = tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)
    reg_loss = tf.reduce_sum(reg_losses)
    # loss = tf.add_n([loss] + reg_losses)
    # =================
    l2_reg_loss = tf.constant(0.0, tf.float32)
    for vv in tf.trainable_variables():
        if ‘bn‘ in vv.name or ‘batchnorm‘ in vv.name or ‘batch_norm‘ in vv.name                 or ‘batch_normalization‘ in vv.name or ‘gn‘ in vv.name:
            continue
        else:
            l2_reg_loss = tf.add(l2_reg_loss, tf.nn.l2_loss(vv))
    l2_reg_loss *= 0.001
    # loss = loss + l2_reg_loss
    # =================

# 4. define optimizer
learning_rate_init = 0.01
global_step = tf.Variable(0, trainable=False)
with tf.name_scope(‘train‘):
    learning_rate = tf.train.polynomial_decay(  # 多项式衰减
        learning_rate=learning_rate_init,  # 初始学习率
        global_step=global_step,  # 当前迭代次数
        decay_steps=22000,  # 在迭代到该次数实际，学习率衰减为 learning_rate * dacay_rate
        end_learning_rate=learning_rate_init / 10,  # 最小的学习率
        power=0.9,
        cycle=False
    )
    update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)  # for batch normalization
    with tf.control_dependencies(update_ops):
        optimizer_op = tf.train.MomentumOptimizer(
            learning_rate=learning_rate, momentum=0.9).minimize(
            loss=loss,
            var_list=tf.trainable_variables(),
            global_step=global_step # 不指定的话学习率不更新
        )

with tf.name_scope(‘eval‘):
    correct = tf.nn.in_top_k(logits, tf.argmax(y, axis=1), 1) # 目标是否在前K个预测中, label‘s dtype is int*
    accuracy = tf.reduce_mean(tf.cast(correct, tf.float32))

# 5. initialize
init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer())
saver = tf.train.Saver()

# 5. train & test
n_epochs = 1
batch_size = 55000
with tf.Session() as sess:
    saver.restore(sess, ‘./my_model_final.ckpt‘)
    for epoch in range(n_epochs):
        for iteration in range(mnist.train.num_examples // batch_size):
            X_batch, y_batch = mnist.train.next_batch(batch_size)
            loss_, l2_reg_loss_, reg_loss_ = sess.run([loss, l2_reg_loss, reg_loss], feed_dict={X: X_batch, y: y_batch, is_training:True})
        acc_train = accuracy.eval(feed_dict={X: X_batch, y: y_batch, is_training:False}) # 最后一个 batch 的 accuracy
        acc_test = accuracy.eval(feed_dict={X: mnist.test.images, y: mnist.test.labels, is_training:False})
        loss_test = loss.eval(feed_dict={X: mnist.test.images, y: mnist.test.labels, is_training:False})
        l2_reg_loss_test = l2_reg_loss.eval(feed_dict={X: mnist.test.images, y: mnist.test.labels, is_training:False})
        reg_loss_test = reg_loss.eval(feed_dict={X: mnist.test.images, y: mnist.test.labels, is_training:False})
        print("Train loss:", loss_, "Train l2_reg_loss:", l2_reg_loss_, "Train reg_loss:", reg_loss_, "Train accuracy:", acc_train)
        print("Test loss:", loss_test, "Test l2_reg_loss:", l2_reg_loss_test, "Test reg_loss:", reg_loss_test, "Test accuracy:", acc_test)

"""
# =================
Train loss: 0.000636433 Train l2_reg_loss: 0.48696715 Train reg_loss: 0.48683384 Train accuracy: 1.0
Test loss: 0.059231624 Test l2_reg_loss: 0.48696715 Test reg_loss: 0.48683384 Test accuracy: 0.983
"""