TensorFlow模型保存与恢复：使用pb文件格式详解

凌晨两点半

发布时间：2026-10-13 | 4人浏览过

原创

本文介绍如何在TensorFlow中将训练好的模型保存为pb文件，并演示如何从该文件中恢复模型进行推理。通过graph_util.convert_variables_to_constants将变量转换为常量，实现模型的序列化存储与加载，适用于模型部署场景。

TensorFlow模型保存为pb文件

在TensorFlow中，将训练好的模型保存为pb文件是一种常见的做法，特别是在需要将模型部署到生产环境时。以下代码展示了如何创建一个简单的计算图，并将其保存为pb文件。首先，我们定义了几个变量和操作，然后使用`graph_util.convert_variables_to_constants`将变量转换为常量，最后将计算图写入到模型文件中。


import tensorflow as tf
from tensorflow.python.framework import graph_util
var1 = tf.Variable(1.0, dtype=tf.float32, name='v1')
var2 = tf.Variable(2.0, dtype=tf.float32, name='v2')
var3 = tf.Variable(2.0, dtype=tf.float32, name='v3')
x = tf.placeholder(dtype=tf.float32, shape=None, name='x')
x2 = tf.placeholder(dtype=tf.float32, shape=None, name='x2')
addop = tf.add(x, x2, name='add')
addop2 = tf.add(var1, var2, name='add2')
addop3 = tf.add(var3, var2, name='add3')
initop = tf.global_variables_initializer()
model_path = './Test/model.pb'
with tf.Session() as sess:
  sess.run(initop)
  print(sess.run(addop, feed_dict={x: 12, x2: 23}))
  output_graph_def = graph_util.convert_variables_to_constants(sess, sess.graph_def, ['add', 'add2', 'add3'])
  # 将计算图写入到模型文件中
  model_f = tf.gfile.FastGFile(model_path, mode="wb")
  model_f.write(output_graph_def.SerializeToString())

从pb文件恢复模型

要从保存的pb文件中恢复模型，可以使用`tf.import_graph_def`函数。以下代码展示了如何从pb文件中加载模型，并执行特定的操作。通过指定`return_elements`参数，我们可以选择性地加载模型中的某些部分，从而实现灵活的模型恢复。


import tensorflow as tf
with tf.Session() as sess:
  model_f = tf.gfile.FastGFile("./Test/model.pb", mode='rb')
  graph_def = tf.GraphDef()
  graph_def.ParseFromString(model_f.read())
  c = tf.import_graph_def(graph_def, return_elements=["add2:0"])
  c2 = tf.import_graph_def(graph_def, return_elements=["add3:0"])
  x, x2, c3 = tf.import_graph_def(graph_def, return_elements=["x:0", "x2:0", "add:0"])

  print(sess.run(c))
  print(sess.run(c2))
  print(sess.run(c3, feed_dict={x: 23, x2: 2}))

以上这篇TensorFlow实现保存训练模型为pd文件并恢复就是小编分享给大家的全部内容了，希望能给大家一个参考，也希望大家多多支持本站。