上一篇博文《TensorFlow模型参数的保存和加载》介绍了如何保存和加载TensorFlow模型训练参数,保存对象主要是Tensor/Variables。这一节我们介绍如何保存和复用op。
和Tensor一样,保存op需要在训练时为op指定名字,如下所示:
softmax = tf.nn.softmax(tf.matmul(x, W) + b,name="op_softmax")
在识别阶段,调用get_operation_by_name()函数,以op名字作为参数,如下所示:
op_softmax =sess.graph.get_operation_by_name("op_softmax").outputs[0]
这里需要注意的是,在最后面要加上“.outputs[0]”,否则会出现异常。
如果要直接运行sess.run(op_softmax),需要指定feed_dict。以官方mnist训练案例为例,调用格式为sess.run(op_softmax,feed_dict={x: mnist.test.images})。
在加载复用阶段,W和b的值已经保存在checkpoint数据中,故不需要再次声明W和b。但是,需要通过get_tensor_by_name()获取到x的声明,如下所示:
x = sess.graph.get_tensor_by_name("x:0")
上述op保存和加载操作可总结为:
1. 在训练代码中,为op指定名字;
2. 在复用阶段,通过get_operation_by_name().outputs[0]获取op;
3. 通过get_tensor_by_name()获取到feed_dict输入tensor(即x)的声明,然后执行sess.run(op, feed_dict={x:new_data})。
完整演示代码如下:
#!/usr/bin/python3.5
# -*- coding: utf-8 -*-
import os
import sys
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data
SAVER_DIR = "train-saver/"
mnist = input_data.read_data_sets("MNIST_data/", one_hot=True)
print ('本脚本须输入参数save或restore')
print ('如果当前目录下没有MNIST_data数据,可能需要花费几分钟等待mnist数据下载')
print ('如果下载缓慢,可以从百度网盘http://pan.baidu.com/s/1c2k3gkw直接下载,放到运行脚本同一目录下即可')
if __name__ =='__main__' and sys.argv[1]=='save':
x = tf.placeholder(tf.float32, [None, 784], name="x")
labels = tf.placeholder(tf.float32, [None, 10], name="labels")
W = tf.Variable(tf.zeros([784, 10]), name="var_W")
b = tf.Variable(tf.zeros([10]), name="var_b")
# 构建网络op
softmax = tf.nn.softmax(tf.matmul(x, W) + b, name="op_softmax")
cross_entropy = tf.reduce_mean(-tf.reduce_sum(labels * tf.log(softmax), reduction_indices=[1]))
train_step = tf.train.GradientDescentOptimizer(0.5).minimize(cross_entropy)
sess = tf.InteractiveSession()
tf.global_variables_initializer().run()
for i in range(1000):
if i%10 == 0:
print ("正在进行第 %d 次训练迭代......" % (i))
batch_xs, batch_ys = mnist.train.next_batch(100)
sess.run(train_step, feed_dict={x: batch_xs, labels: batch_ys})
# 保存训练结果
if not os.path.exists(SAVER_DIR):
print ('不存在训练数据保存目录,现在创建保存目录')
os.makedirs(SAVER_DIR)
# 初始化saver
saver = tf.train.Saver()
saver_path = saver.save(sess, "%smodel.ckpt"%(SAVER_DIR))
print ('mnist手写体数字训练结果已保存!')
if __name__ =='__main__' and sys.argv[1]=='restore':
sess = tf.InteractiveSession()
# 导入保存训练结果的文件
saver = tf.train.import_meta_graph("%smodel.ckpt.meta"%(SAVER_DIR))
model_file=tf.train.latest_checkpoint(SAVER_DIR)
saver.restore(sess, model_file)
# 通过指定变量名获取训练结果中的变量值
x = sess.graph.get_tensor_by_name("x:0")
# 执行识别
op_softmax = sess.graph.get_operation_by_name("op_softmax").outputs[0]
recog = tf.argmax(op_softmax,1)
print("mnist手写体数字测试集识别结果为:%s" % (sess.run(recog, feed_dict={x: mnist.test.images})))
python op-restore.py save
识别时启动命令为:
python op-restore.py restore
当第一次运行该脚本的时候,如果当前目录没有mnist数据集,则会自动下载数据集,如果网络不稳定,那么下载过程会很缓慢。为方便使用,我把mnist数据集上传到百度网盘,链接地址如下:
http://pan.baidu.com/s/1c2k3gkw
下载后将整个MNIST_data文件夹放到脚本同一目录,运行时就不会触发下载了。