Skip to main content
 首页 » 编程设计

python之如何使用 TensorFlow 中的 Estimator 将模型存储在 `.pb` 文件中

2025年05月04日24cmt

我使用 TensorFlow 的估算器训练了我的模型。看来应该用export_savedmodel来制作.pb文件,但是我真的不知道如何构造serving_input_receiver_fn。有人有什么想法吗? 欢迎使用示例代码。

额外问题:

  1. .pb 是我想要重新加载模型时唯一需要的文件吗? 变量不需要?

  2. 与使用 adam 优化器的 .ckpt 相比,.pb 会减少多少模型文件大小?

请您参考如下方法:

您可以使用 freeze_graph.py 从 .ckpt + .pbtxt 生成一个 .pb 如果您使用的是 tf.estimator.Estimator ,那么您会在 model_dir 中找到这两个文件

python freeze_graph.py \ 
    --input_graph=graph.pbtxt \ 
    --input_checkpoint=model.ckpt-308 \ 
    --output_graph=output_graph.pb 
    --output_node_names=<output_node> 
  1. Is .pb the only file I need when I want to reload the model? Variable unnecessary?

是,您必须知道您也是模型的输入节点和输出节点名称。然后使用import_graph_def加载.pb文件,使用get_operation_by_name获取输入输出操作

  1. How much will .pb reduced the model file size compared with .ckpt with adam optimizer?

.pb 文件不是压缩的.ckpt 文件,因此没有“压缩率”。

但是,有一种方法 to optimize your .pb file 可用于推理,并且此优化可能会减小文件大小,因为它删除了图形中仅用于训练操作的部分(请参阅完整说明 here)。

[comment] how can I get the input and output node names?

您可以使用 op name 参数设置输入和输出节点名称。

要在您的 .pbtxt 文件中列出节点名称,请使用以下脚本。

import tensorflow as tf 
from google.protobuf import text_format 
 
with open('graph.pbtxt') as f: 
    graph_def = text_format.Parse(f.read(), tf.GraphDef()) 
 
print [n.name for n in graph_def.node] 

[comment] I found that there is a tf.estimator.Estimator.export_savedmodel(), is that the function to store model in .pb directly? And I'm struggling in it's parameter serving_input_receiver_fn. Any ideas?

export_savedmodel() 生成一个 SavedModel,它是 TensorFlow 模型的通用序列化格式。 它应该包含适合 TensorFlow Serving APIs

serving_input_receiver_fn() 是生成 SavedModel 所需内容的一部分,它通过添加占位符来确定模型的输入签名图表。

来自文档

This function has the following purposes:

  • To add placeholders to the graph that the serving system will feed with inference requests.
  • To add any additional ops needed to convert data from the input format into the feature Tensors expected by the model.

如果您收到序列化 tf.Examples 形式的推理请求(这是一种典型模式),那么您可以使用 doc 中提供的示例。

feature_spec = {'foo': tf.FixedLenFeature(...), 
                'bar': tf.VarLenFeature(...)} 
 
def serving_input_receiver_fn(): 
  """An input receiver that expects a serialized tf.Example.""" 
  serialized_tf_example = tf.placeholder(dtype=tf.string, 
                                         shape=[default_batch_size], 
                                         name='input_example_tensor') 
  receiver_tensors = {'examples': serialized_tf_example} 
  features = tf.parse_example(serialized_tf_example, feature_spec) 
  return tf.estimator.export.ServingInputReceiver(features, receiver_tensors) 

[comment] Any idea to list the node names in '.pb'?

这取决于它是如何生成的。

如果它是 SavedModel 使用:

import tensorflow as tf 
 
with tf.Session() as sess: 
    meta_graph_def = tf.saved_model.loader.load( 
        sess, 
        [tf.saved_model.tag_constants.SERVING], 
        './saved_models/1519232535') 
    print [n.name for n in meta_graph_def.graph_def.node] 

如果它是 MetaGraph 则使用:

import tensorflow as tf 
from tensorflow.python.platform import gfile 
 
with tf.Session() as sess: 
    with gfile.FastGFile('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='') 
        print [n.name for n in graph_def.node]