tensorflow模型保存为pb的各种姿势

tensorflow模型保存的各种姿势

一，直接保存pb

1, 首先我们当然可以直接在tensorflow训练中直接保存为pb为格式，保存pb的好处就是使用场景是实现创建模型与使用模型的解耦，使得创建模型与使用模型的解耦，使得前向推导inference代码统一。另外的好处就是保存为pb的时候，模型的变量会变成固定的，导致模型的大小会大大减小。

这里稍稍解释下pb：是MetaGraph的protocol buffer格式的文件，MetaGraph包括计算图，数据流，以及相关的变量和输入输出

主要使用tf.SavedModelBuilder来完成这个工作，并且可以把多个计算图保存到一个pb文件中，如果有多个MetaGraph，那么只会保留第一个MetaGraph的版本号。

保持pb的文件代码：

import tensorflow as tf
import os
from tensorflow.python.framework import graph_util
 
pb_file_path = os.getcwd()
 
with tf.Session(graph=tf.Graph()) as sess:
    x = tf.placeholder(tf.int32, name='x')
    y = tf.placeholder(tf.int32, name='y')
    b = tf.Variable(1, name='b')
    xy = tf.multiply(x, y)
    # 这里的输出需要加上name属性
    op = tf.add(xy, b, name='op_to_store')
 
    sess.run(tf.global_variables_initializer())
 
    # convert_variables_to_constants 需要指定output_node_names，list()，可以多个
    constant_graph = graph_util.convert_variables_to_constants(sess, sess.graph_def, ['op_to_store'])
 
    # 测试 OP
    feed_dict = {x: 10, y: 3}
    print(sess.run(op, feed_dict))
 
    # 写入序列化的 PB 文件
    with tf.gfile.FastGFile(pb_file_path+'model.pb', mode='wb') as f:
        f.write(constant_graph.SerializeToString())
 
    # 输出
    # INFO:tensorflow:Froze 1 variables.
    # Converted 1 variables to const ops.
    # 31

其实主要是:

 # convert_variables_to_constants 需要指定output_node_names，list()，可以多个
    constant_graph = graph_util.convert_variables_to_constants(sess, sess.graph_def, ['op_to_store'])

    # 写入序列化的 PB 文件
    with tf.gfile.FastGFile(pb_file_path+'model.pb', mode='wb') as f:
        f.write(constant_graph.SerializeToString())

1.1 加载测试代码

from tensorflow.python.platform import gfile
 
sess = tf.Session()
with gfile.FastGFile(pb_file_path+'model.pb', 'rb') as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())
    sess.graph.as_default()
    tf.import_graph_def(graph_def, name='') # 导入计算图
 
# 需要有一个初始化的过程    
sess.run(tf.global_variables_initializer())
 
# 需要先复原变量
print(sess.run('b:0'))
# 1
 
# 输入
input_x = sess.graph.get_tensor_by_name('x:0')
input_y = sess.graph.get_tensor_by_name('y:0')
 
op = sess.graph.get_tensor_by_name('op_to_store:0')
 
ret = sess.run(op,  feed_dict={input_x: 5, input_y: 5})
print(ret)
# 输出 26

2，第二种就是采用上述的那API来进行保存

import tensorflow as tf
import os
from tensorflow.python.framework import graph_util
 
pb_file_path = os.getcwd()
 
with tf.Session(graph=tf.Graph()) as sess:
    x = tf.placeholder(tf.int32, name='x')
    y = tf.placeholder(tf.int32, name='y')
    b = tf.Variable(1, name='b')
    xy = tf.multiply(x, y)
    # 这里的输出需要加上name属性
    op = tf.add(xy, b, name='op_to_store')
 
    sess.run(tf.global_variables_initializer())
 
    # convert_variables_to_constants 需要指定output_node_names，list()，可以多个
    constant_graph = graph_util.convert_variables_to_constants(sess, sess.graph_def, ['op_to_store'])
 
    # 测试 OP
    feed_dict = {x: 10, y: 3}
    print(sess.run(op,