上一篇MNIST手写数字识别——CNN中，没有实现模型的保存和回复，现在实现了~后续跟进一下
在train.py中进行了模型训练，并通过tf.train.Saver这个类进行模型的保存与恢复。test.py中直接恢复模型，就可以进行测试了，十分方便~
其中train.py中就是新增了两行代码
model_saver = tf.train.Saver()
model_saver.save(sess, './model/mnist_model.ckpt')
train.py
import tensorflow as tf
import tensorflow.examples.tutorials.mnist.input_data as input_data
mnist = input_data.read_data_sets("data/", one_hot=True, validation_size=0) # 下载并加载mnist数据
print("mnist训练集大小：", len(mnist.train.images))
print("mnist测试集大小：", len(mnist.test.images))
x = tf.placeholder(tf.float32, [None, 784], name="x") # 输入的数据占位符
y_actual = tf.placeholder(tf.float32, shape=[None, 10], name="y_actual") # 输入的标签占位符
# 定义一个函数，用于初始化所有的权值 W
def weight_variable(shape):
initial = tf.truncated_normal(shape, stddev=0.1)
return tf.Variable(initial)
# 定义一个函数，用于初始化所有的偏置项 b
def bias_variable(shape):
initial = tf.constant(0.1, shape=shape)
return tf.Variable(initial)
# 定义一个函数，用于构建卷积层
def conv2d(x, W):
return tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding='SAME')
# 定义一个函数，用于构建池化层
def max_pool(x):
return tf.nn.max_pool(x, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
# 构建网络
x_image = tf.reshape(x, [-1, 28, 28, 1]) # 转换输入数据shape,以便于用于网络中
W_conv1 = weight_variable([5, 5, 1, 32])
b_conv1 = bias_variable([32])
h_conv1 = tf.nn.relu(conv2d(x_image, W_conv1) + b_conv1) # 第一个卷积层
h_pool1 = max_pool(h_conv1) # 第一个池化层
W_conv2 = weight_variable([5, 5, 32, 64])
b_conv2 = bias_variable([64])
h_conv2 = tf.nn.relu(conv2d(h_pool1, W_conv2) + b_conv2) # 第二个卷积层
h_pool2 = max_pool(h_conv2) # 第二个池化层
W_fc1 = weight_variable([7 * 7 * 64, 1024])
b_fc1 = bias_variable([1024])
h_pool2_flat = tf.reshape(h_pool2, [-1, 7 * 7 * 64]) # reshape成向量
h_fc1 = tf.nn.relu(tf.matmul(h_pool2_flat, W_fc1) + b_fc1) # 第一个全连接层
keep_prob = tf.placeholder("float", name="keep_prob")
h_fc1_drop = tf.nn.dropout(h_fc1, keep_prob) # dropout层
W_fc2 = weight_variable([1024, 10])
b_fc2 = bias_variable([10])
y_predict = tf.nn.softmax(tf.matmul(h_fc1_drop, W_fc2) + b_fc2, name="y_predict") # softmax层
cross_entropy = -tf.reduce_sum(y_actual * tf.log(y_predict),name = "cross_entropy") # 交叉熵
train_step = tf.train.GradientDescentOptimizer(1e-3).minimize(cross_entropy) # 梯度下降法
correct_prediction = tf.equal(tf.argmax(y_predict, 1), tf.argmax(y_actual, 1))
accuracy = tf.reduce_mean(tf.cast(correct_prediction, "float")) # 精确度计算
model_saver = tf.train.Saver()
sess = tf.InteractiveSession()
sess.run(tf.global_variables_initializer())
print("开始训练:")
for i in range(1000):
batch = mnist.train.next_batch(20)
if i % 1000 == 0:
train_accuracy = accuracy.eval(feed_dict={x: batch[0], y_actual: batch[1], keep_prob:1.0})
print('step', i, 'train_accuracy', train_accuracy)
train_step.run(feed_dict={x: batch[0], y_actual: batch[1], keep_prob: 0.5})
model_saver.save(sess, './model/mnist_model.ckpt')
print("训练结束")
print("开始测试:")
test_accuracy = 0.0
for i in range(100):
batch = mnist.test.next_batch(100)
test_accuracy_i = accuracy.eval(feed_dict={x:batch[0], y_actual:batch[1], keep_prob:1.0})
test_accuracy += test_accuracy_i
if (i+1) % 20 == 0:
print('step', i+1, 'test accuracy: ', test_accuracy_i)
print("训练集准确度: ", test_accuracy/100)
print("测试结束")
sess.close()
#这行代码保存了神经网络的输出，这个在后面使用导入模型过程中起到关键作用
#tf.add_to_collection('network_output', y_predict)
运行train.py后将会在相应的路径下生成存储模型的文件：
test.py
# -*- coding: utf-8 -*-
import tensorflow as tf
import tensorflow.examples.tutorials.mnist.input_data as input_data
mnist = input_data.read_data_sets("data/", one_hot=True) # 下载并加载mnist数据
saver = tf.train.import_meta_graph("model/mnist_model.ckpt.meta")
with tf.Session() as sess:
saver.restore(sess, "./model/mnist_model.ckpt") #恢复模型
gragh = tf.get_default_graph() # 获取当前图，为了后续训练时恢复变量
x = gragh.get_tensor_by_name("x:0")# 获取输入变量（占位符，由于保存时未定义名称，tf自动赋名称“Placeholder”）
y_actual = gragh.get_tensor_by_name('y_actual:0')# 获取输出变量
keep_prob = gragh.get_tensor_by_name('keep_prob:0')# 获取dropout的保留参数
accuracy = gragh.get_tensor_by_name('accuracy:0')# 获取dropout的保留参数
print("开始测试:")
test_accuracy = 0.0
for i in range(100):
batch = mnist.test.next_batch(100)
test_accuracy_i = accuracy.eval(feed_dict={x:batch[0], y_actual:batch[1], keep_prob:1.0})
test_accuracy += test_accuracy_i
if (i+1) % 20 == 0:
print('step', i+1, 'test accuracy: ', test_accuracy_i)
print("训练集准确度: ", test_accuracy/100)
print("测试结束")
运行结果
