基于卷积神经网络的手写数字识别分类(Tensorflow)

import numpy as np
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data
import os
%matplotlib inline
import matplotlib.pyplot as plt

mnist = input_data.read_data_sets('MNIST_data', one_hot=True)

class ConvModel(object):
    def __init__(self, lr, batch_size, iter_num):
        self.lr = lr
        self.batch_size = batch_size
        self.iter_num = iter_num

        self.X_flat = tf.placeholder(tf.float32, [None, 784])
        self.X = tf.reshape(self.X_flat, [-1, 28, 28, 1]) # 本次要用卷积进行运算，所以使用2维矩阵。从这个角度讲，利用了更多的位置信息。
        self.y = tf.placeholder(tf.float32, [None, 10])
        self.dropRate = tf.placeholder(tf.float32)

        conv1 = tf.layers.conv2d(self.X, 32, 5, padding='same', activation=tf.nn.relu,
                                     kernel_initializer=tf.truncated_normal_initializer(stddev=0.1, seed=0),
                                     bias_initializer=tf.constant_initializer(0.1))
        conv1 = tf.layers.max_pooling2d(conv1 , 2,2)
        conv2 = tf.layers.conv2d(conv1, 64, 5, padding='same', activation=tf.nn.relu,
                                     kernel_initializer=tf.truncated_normal_initializer(stddev=0.1, seed=0),
                                     bias_initializer=tf.constant_initializer(0.1))
        pool1 = tf.layers.max_pooling2d(conv2, 2,2)
        flatten = tf.reshape(pool1 , [-1, 7*7*64])
        dense1 = tf.layers.dense(flatten, 1024,  activation=tf.nn.relu, use_bias=True,
                                 kernel_initializer=tf.truncated_normal_initializer(stddev=0.1, seed=0),
                                 bias_initializer=tf.constant_initializer(0.1))
        dense1_ = tf.nn.dropout(dense1, self.dropRate)
        dense2 = tf.layers.dense(dense1_, 10, activation=tf.nn.relu, use_bias=True,
                                 kernel_initializer=tf.truncated_normal_initializer(stddev=0.1, seed=0),
                                 bias_initializer=tf.constant_initializer(0.1))

        self.loss = tf.losses.softmax_cross_entropy(onehot_labels=self.y, logits=dense2)
        self.train_step = tf.train.AdamOptimizer(1e-4).minimize(self.loss )

        # 用于模型训练
        self.correct_prediction = tf.equal(tf.argmax(self.y, axis=1), tf.argmax(dense2, axis=1))
        self.accuracy = tf.reduce_mean(tf.cast(self.correct_prediction, tf.float32))
        # 用于保存训练好的模型
        self.saver = tf.train.Saver()
    def train(self):
        with tf.Session() as sess:
            sess.run(tf.global_variables_initializer())  # 先初始化所有变量。
            for i in range(self.iter_num):
                batch_x, batch_y = mnist.train.next_batch(self.batch_size)   # 读取一批数据
                loss, _= sess.run([self.loss, self.train_step],
                                  feed_dict={self.X_flat: batch_x, self.y: batch_y, self.dropRate: 0.5})   # 每调用一次sess.run，就像拧开水管一样，所有self.loss和self.train_step涉及到的运算都会被调用一次。
                if i%1000 == 0:
                    train_accuracy = sess.run(self.accuracy, feed_dict={self.X_flat: batch_x, self.y: batch_y, self.dropRate: 1.})  # 把训练集数据装填进去
                    test_x, test_y = mnist.test.next_batch(self.batch_size)
                    test_accuracy = sess.run(self.accuracy, feed_dict={self.X_flat: test_x, self.y: test_y, self.dropRate: 1.})   # 把测试集数据装填进去
                    print ('iter\t%i\tloss\t%f\ttrain_accuracy\t%f\ttest_accuracy\t%f' % (i,loss,train_accuracy,test_accuracy))
            self.saver.save(sess, 'model/mnistModel') # 保存模型

    def test(self):
        with tf.Session() as sess:
            self.saver.restore(sess, 'model/mnistModel')
            Accuracy = []
            for i in range(int(10000/self.batch_size)):
                test_x, test_y = mnist.test.next_batch(self.batch_size)
                test_accuracy = sess.run(self.accuracy, feed_dict={self.X_flat: test_x, self.y: test_y, self.dropRate: 1.})
                Accuracy.append(test_accuracy)
            print('==' * 15)
            print( 'Test Accuracy: ', np.mean(np.array(Accuracy))   ) 

model = ConvModel(0.001, 64, 30000)   # 学习率为0.001，每批传入64张图，训练30000次
model.train()      # 训练模型
model.test()       # 预测