TensorFlow 使用pb文件保存和恢复模型计算图及参数详解

Nora

发布时间：2026-10-12 | 1人浏览过

原创

本文详解如何在TensorFlow中使用graph_util.convert_variables_to_constants将计算图与参数保存为pb文件，并通过GraphDef对象恢复模型。文章涵盖序列化原理、关键参数说明、代码实现及操作节点分析，帮助开发者掌握模型部署与迁移的核心技术。

一、保存模型：

利用 graph_util.convert_variables_to_constants 函数，可以将当前 Session 中的计算图序列化为一个二进制字节流。该函数主要包含三个参数：第一个参数是当前活跃的 Session，其中包含了所有的变量信息；

第二个参数是 GraphDef 对象，用于描述整个计算网络的结构；

第三个参数是一个列表，指定了图中需要输出的节点名称。

该函数的返回值是一个精简版的 GraphDef 对象，它保留了原始输入 GraphDef 以及 Session 中的网络结构和变量信息。通过调用其成员函数 SerializeToString()，可以将这些信息序列化为字节流，并写入到文件中：

constant_graph = graph_util.convert_variables_to_constants( sess, sess.graph_def , ['sum_operation'] )
with open( pbName, mode='wb') as f:
f.write(constant_graph.SerializeToString())

需要注意的是，如果原始张量（即包含在参数1和参数2中的组成部分）不参与由参数3指定的输出节点列表所涉及的张量计算，那么这些张量将不会出现在返回的 GraphDef 对象中，也不会被序列化写入到 pb 文件里。

二、恢复模型：

在恢复阶段，需要创建一个 GraphDef 对象，从上述生成的文件中加载数据，然后将其导入到当前的 Session 中：

    graph0 = tf.GraphDef()
    with open( pbName, mode='rb') as f:
      graph0.ParseFromString( f.read() )
      tf.import_graph_def( graph0 , name = '' )

三、完整代码示例：

 
import tensorflow as tf 
from tensorflow.python.framework import graph_util
 
pbName = 'graphA.pb'
def graphCreate() :
  with tf.Session() as sess :
    var1 = tf.placeholder ( tf.int32 , name='var1' ) 
    var2 = tf.Variable( 20 , name='var2' )#实参name='var2'指定了操作名，该操作返回的张量名是在
                       #'var2'后面:0 ,即var2:0 是返回的张量名，也就是说变量
                       # var2的名称是'var2:0'
    var3 = tf.Variable( 30 , name='var3' )
    var4 = tf.Variable( 40 , name='var4' )
    var4op = tf.assign( var4 , 1000 , name = 'var4op1' )
    sum = tf.Variable( 4, name='sum' )
    sum = tf.add ( var1 , var2, name = 'var1_var2' ) 
    sum = tf.add( sum , var3 , name='sum_var3' )
    sumOps = tf.add( sum , var4 , name='sum_operation' )
    oper = tf.get_default_graph().get_operations()
    with open( 'operation.csv','wt' ) as f:
      s = 'name,type,output\n'
      f.write( s ) 
      for o in oper:
        s = o.name
        s += ','+ o.type 
        inp = o.inputs
        oup = o.outputs
        for iip in inp :
          s #s += ','+ str(iip)
        for iop in oup :
          s += ',' + str(iop)
        s += '\n'
        f.write( s ) 
         
      for var in tf.global_variables():
        print('variable=> ' , var.name) #张量是tf.Variable/tf.Add之类操作的结果，
                        #张量的名字使用操作名加:0来表示
    init = tf.global_variables_initializer()
    sess.run( init )
    sess.run( var4op )
    print('sum_operation result is Tensor ' , sess.run( sumOps , feed_dict={var1:1}) )
 
    constant_graph = graph_util.convert_variables_to_constants( sess, sess.graph_def , ['sum_operation'] )
    with open( pbName, mode='wb') as f:
      f.write(constant_graph.SerializeToString())
 
def graphGet() :
  print("start get:" )
  with tf.Graph().as_default():
    graph0 = tf.GraphDef()
    with open( pbName, mode='rb') as f:
      graph0.ParseFromString( f.read() )
      tf.import_graph_def( graph0 , name = '' )
    with tf.Session() as sess :
      init = tf.global_variables_initializer()
      sess.run(init)
      v1 = sess.graph.get_tensor_by_name('var1:0' )
      v2 = sess.graph.get_tensor_by_name('var2:0' )
      v3 = sess.graph.get_tensor_by_name('var3:0' )
      v4 = sess.graph.get_tensor_by_name('var4:0' )
      
      sumTensor = sess.graph.get_tensor_by_name("sum_operation:0")
      print('sumTensor is : ' , sumTensor )
      print( sess.run( sumTensor , feed_dict={v1:1} ) ) 
  
graphCreate()
graphGet()

四、保存 pb 函数代码中的操作名称、类型及返回张量分析：

operation name	operation type	output
var1	Placeholder	Tensor("var1:0"	dtype=int32)
var2/initial_value	Const	Tensor("var2/initial_value:0"	shape=()	dtype=int32)
var2	VariableV2	Tensor("var2:0"	shape=()	dtype=int32_ref)
var2/Assign	Assign	Tensor("var2/Assign:0"	shape=()	dtype=int32_ref)
var2/read	Identity	Tensor("var2/read:0"	shape=()	dtype=int32)
var3/initial_value	Const	Tensor("var3/initial_value:0"	shape=()	dtype=int32)
var3	VariableV2	Tensor("var3:0"	shape=()	dtype=int32_ref)
var3/Assign	Assign	Tensor("var3/Assign:0"	shape=()	dtype=int32_ref)
var3/read	Identity	Tensor("var3/read:0"	shape=()	dtype=int32)
var4/initial_value	Const	Tensor("var4/initial_value:0"	shape=()	dtype=int32)
var4	VariableV2	Tensor("var4:0"	shape=()	dtype=int32_ref)
var4/Assign	Assign	Tensor("var4/Assign:0"	shape=()	dtype=int32_ref)
var4/read	Identity	Tensor("var4/read:0"	shape=()	dtype=int32)
var4op1/value	Const	Tensor("var4op1/value:0"	shape=()	dtype=int32)
var4op1	Assign	Tensor("var4op1:0"	shape=()	dtype=int32_ref)
sum/initial_value	Const	Tensor("sum/initial_value:0"	shape=()	dtype=int32)
sum	VariableV2	Tensor("sum:0"	shape=()	dtype=int32_ref)
sum/Assign	Assign	Tensor("sum/Assign:0"	shape=()	dtype=int32_ref)
sum/read	Identity	Tensor("sum/read:0"	shape=()	dtype=int32)
var1_var2	Add	Tensor("var1_var2:0"	dtype=int32)
sum_var3	Add	Tensor("sum_var3:0"	dtype=int32)
sum_operation	Add	Tensor("sum_operation:0"	dtype=int32)