TensorFlow模型op的保存和加载(含演示代码)

上一篇博文《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})))

假设脚本文件名为op-restore.py,则训练时启动命令为:

python op-restore.py save

识别时启动命令为:

python op-restore.py restore

当第一次运行该脚本的时候,如果当前目录没有mnist数据集,则会自动下载数据集,如果网络不稳定,那么下载过程会很缓慢。为方便使用,我把mnist数据集上传到百度网盘,链接地址如下:

http://pan.baidu.com/s/1c2k3gkw

下载后将整个MNIST_data文件夹放到脚本同一目录,运行时就不会触发下载了。

你可能感兴趣的:(TensorFlow使用技巧)