【Tensorflow】模型存储，生成“pb”文件与读取

#通过convert_variables_to_constants函数将计算图中的变量及其取值通过常量的方式保存于一个文件中
 
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
 
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    # 导出当前计算图的GraphDef部分，即从输入层到输出层的计算过程部分
    graph_def = tf.get_default_graph().as_graph_def()
    output_graph_def = graph_util.convert_variables_to_constants(sess,
                                                        graph_def, ['add'])
 
    with tf.gfile.GFile("Model/combined_model.pb", 'wb') as f:
        f.write(output_graph_def.SerializeToString())
 
 
#  载入包含变量及其取值的模型
 
import tensorflow as tf
from tensorflow.python.platform import gfile
 
with tf.Session() as sess:
    model_filename = "Model/combined_model.pb"
    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)) # [array([ 3.], dtype=float32)]

最后

以上就是忧心高跟鞋为你收集整理的【Tensorflow】模型存储，生成“pb”文件与读取的全部内容，希望文章能够帮你解决【Tensorflow】模型存储，生成“pb”文件与读取所遇到的程序开发问题。

如果觉得靠谱客网站的内容还不错，欢迎将靠谱客网站推荐给程序员好友。