TensorFlow学习笔记2:pb文件的保存和加载方法

pb文件的保存方法

使用 tf.train.Saver 会保存运行 TensorFlow 程序所需要的全部信息，然而有时并不需要某些信息。比如在测试或者离线预测时，只需要知道如何从神经网络的输入层经过前向传播计算得到输出层即可，而不需要类似于变量初始化、模型保存等辅助节点的信息。而且，将变量取值和计算图结构分成不同的文件存储有时候也不方便，于是 TensorFlow 提供了 convert_variables_to_constants 函数，通过
这个函数可以将计算图中的变量及其取值通过常量的方式保存，这样整个 TensorFlow 计算图可以统一存放在一个文件中。以下程序提供了一个样例。

import tensorflow as tf
from tensorflow.python.framework import graph_util

v1 = tf.Variable(tf.constant(1.0, shape=[1]), name = "v1")
v2 = tf.Variable(tf.constant(2.0, shape=[1]), name = "v2")
result = v1 + v2

init_op = tf.global_variables_initializer()
with tf.Session() as sess:
    sess.run(init_op)
    #导出当前计算图的 GraphDef 部分，只需要这一部分就可以完成从输入层到输出层的计算
    #过程     				
    graph_def = tf.get_default_graph().as_graph_def()
    #张量的名称后面有：0，表示是某个计算节点的第一个输出。而计算节点本身的名称后
    #是没有0 的。
    output_graph_def = graph_util.convert_variables_to_constants(sess, graph_def, ['add'])
    #将导出的模型存入文件。
    with tf.gfile.GFile("Saved_model/combined_model.pb", "wb") as f:
           f.write(output_graph_def.SerializeToString())

加载pb文件。

import tensorflow as tf
from tensorflow.python.platform import gfile
with tf.Session() as sess:
    model_filename = "Saved_model/combined_model.pb"
   #读取保存的模型文件，并将文件解析成对应的 GraphDef Protocol Buffer 。
    with gfile.FastGFile(model_filename, 'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())

    result = tf.import_graph_def(graph_def, return_elements=["add:0"])
    print(sess.run(result))

将graph_def 中保存的图加载到当前的图中。return_elements=[“add:0”]给出了返回的张量的名称。在保存的时候给出的是计算节点的名称 ，所以为“ add”。在加载的时候给出的是张量的名称，所以是add:0。