TensorFlow实战框架Chp10--利用TensorFlow-Slim在MNIST数据集上实现LeNet-5模型

TensorFlow-Slim在MNIST数据集上实现LeNet-5模型

# -*- coding: utf-8 -*-
"""
Created on Mon Jul  9 22:00:18 2018

@author: muli
"""

import tensorflow as tf
import tensorflow.contrib.slim as slim

from tensorflow.examples.tutorials.mnist import input_data

# 通过TensorFlow-Slim来定义LeNet-5的网络结构。
def lenet5(inputs):
    inputs = tf.reshape(inputs, [-1, 28, 28, 1])
    net = slim.conv2d(inputs, 32, [5, 5], padding='SAME', scope='layer1-conv')
    net = slim.max_pool2d(net, 2, stride=2, scope='layer2-max-pool')
    net = slim.conv2d(net, 64, [5, 5], padding='SAME', scope='layer3-conv')
    net = slim.max_pool2d(net, 2, stride=2, scope='layer4-max-pool')
    net = slim.flatten(net, scope='flatten')
    net = slim.fully_connected(net, 500, scope='layer5')
    net = slim.fully_connected(net, 10, scope='output')
    return net


# 2. 和之前的章节类似的训练模型
def train(mnist):
    # 训练数据 及 标签
    x = tf.placeholder(tf.float32, [None, 784], name='x-input')
    y_ = tf.placeholder(tf.float32, [None, 10], name='y-input')
    # 对数据进行训练
    y = lenet5(x)

    # 交叉熵
    cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(
            logits=y, labels=tf.argmax(y_, 1))
    # 计算损失
    loss = tf.reduce_mean(cross_entropy)
    # 优化
    train_op = tf.train.GradientDescentOptimizer(0.01).minimize(loss)
    with tf.Session() as sess:
        tf.global_variables_initializer().run()
        for i in range(5001):
            xs, ys = mnist.train.next_batch(100)
            _, loss_value = sess.run([train_op, loss], feed_dict={x: xs, y_: ys})

            if i % 100 == 0:
                print("After %d training step(s), loss on training batch is %g." % (i, loss_value))

# 3. 主程序
def main(argv=None):
    mnist = input_data.read_data_sets("./datasets/MNIST_data", one_hot=True)
    train(mnist)

if __name__ == '__main__':
    main()