TensorFlow 保存和导入计算图中的部分节点

最新推荐文章于 2024-08-13 22:48:20 发布

hiptonese

最新推荐文章于 2024-08-13 22:48:20 发布

阅读量4.8k

点赞数 1

文章标签： Tensorflow

本文链接：https://blog.csdn.net/u012560212/article/details/73202015

版权

这篇博客详细介绍了如何在TensorFlow中保存和导入计算图中的特定节点。通过`convert_variables_to_constants`函数，将变量转换为常量并保存了'add_res'和'add_res1'两个节点。然后使用`tf.import_graph_def`成功导入保存的节点，运行并打印了这两个节点的输出结果。

摘要由CSDN通过智能技术生成

#保存计算图中的节点
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')
res = tf.add(v1,v2,name='add_res')
res2 = tf.add(res,v2,name='add_res1')

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    
    #获取图
    graph_def = tf.get_default_graph().as_graph_def() 
    
    #保存了两个节点。 
    #convert_variables_to_constants：通过这个函数可以将计算图中的变量及其取值通过常量保存。
    #add_res 没有“:0”，表示这是计算节点，而“add_res:0” 表示节点计算后的输出张量。
    output_graph_def = graph_util.convert_variables_to_constants(sess,graph_def,['add_res','add_res1'])     
    
    #文件句柄，保存文件。
    with tf.gfile.GFile('./save/combined_model.pb','wb') as f:
        f.write(output_graph_def.SerializeToString())
        
#导入已保存的节点