Tensorflow保存和读入pb文件

tensorflow保存pb文件

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

#消除警告：Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

pb_file_path = 'D:/test/'

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)

    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))

    #输出结果
    #Converted 1 variables to const ops.
    #31

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

tensorflow读入pb文件

from tensorflow.python.platform import gfile
import tensorflow as tf
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'


pb_file_path = 'D:/test/'

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})

#输出为26
print(ret)