tensorflow完成mnist数据集训练和验证

coding=utf-8

“”"
author:lei
function:
“”"

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

mnist = input_data.read_data_sets("./mnist_data/", one_hot=True)
trainimgs, trainlabels, testimgs, testlabels = mnist.train.images, mnist.train.labels, mnist.test.images, mnist.test.labels

55000 10000 784 10

ntrain, ntest, dim, nclasses = trainimgs.shape[0], testimgs.shape[0], trainimgs.shape[1], trainlabels.shape[1]

print(ntrain, ntest, dim, nclasses)

diminput = 28
dimhidden = 128
dimoutput = nclasses
nsteps = 28

weights = {‘hidden’: tf.Variable(tf.random_normal([diminput, dimhidden])), ‘out’: tf.Variable(tf.random_normal([dimhidden, dimoutput]))}

biases = {‘hidden’: tf.Variable(tf.random_normal([dimhidden])), ‘out’: tf.Variable(tf.random_normal([dimoutput]))}

def _RNN(_X, _W, _b, _nsteps, _name):
# [batchsize, nsteps, diminput] --> [nsteps, batchsize, diminput]
_X = tf.transpose(_X, [1, 0, 2])
# reshape input to [nsteps * batchsize, diminput]
_X = tf.reshape(_X, [-1, diminput])

# input_layer --> hidden layer
_H = tf.matmul(_X, _W['hidden'] + _b['hidden'])

_Hsplit = tf.split(0, _nsteps, _H)

with tf.variable_scope(_name) as scope:
    lstm_cell = tf.nn.rnn_cell.BasicLSTMCell(dimhidden, forget_bias=1.0)
    _LSTM_O, _LSTM_S = tf.nn.rnn(lstm_cell, _Hsplit, dtype=tf.float32)

# 拿到最后一个lstm的值
_O = tf.matmul(_LSTM_O[-1], _W['out']) + _b['out']

return {
    'X': _X, 'H': _H, 'Hsplit': _Hsplit,
    'LSTM_O': _LSTM_O, 'LSTM_S': _LSTM_S, 'O': _O,
}

learning_rate = 0.001
x = tf.placeholder(tf.float32, [None, nsteps, diminput])
y = tf.placeholder(tf.float32, [None, dimoutput])

myrnn = _RNN(x, weights, biases, nsteps, ‘basic’)
pred = myrnn[‘O’]

cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(pred, y))
optm = tf.train.GradientDescentOptimizer(learning_rate).minimize(cost)
accr = tf.reduce_mean(tf.cast(tf.equal(tf.argmax(pred, 1), tf.argmax(y, 1)), tf.float32))

init = tf.global_variables_initializer()

training_epochs = 5
batch_size = 16
display_step = 1

with tf.compat.v1.Session() as sess:
sess.run(init)

for epoch in range(training_epochs):
    avg_cost = 0
    total_batch = 100
    for i in range(total_batch):
        batch_xs, batch_ys = mnist.train.next_batch(batch_size)
        batch_xs = batch_xs.reshape((batch_size, nsteps, diminput))
        feeds = {x: batch_xs, y: batch_ys}

        sess.run(optm, feed_dict=feeds)

        avg_cost += sess.run(cost, feed_dict=feeds) / total_batch

    if epoch % display_step == 0:
        print("Epoch:{} cost:{}".format(epoch, avg_cost))

        feeds = {x: batch_xs, y: batch_ys}

        train_acc = sess.run(accr, feed_dict=feeds)
        print("training accuracy:{}".format(train_acc))