《Estimator工程实现》系列三：模型导出示例

最新推荐文章于 2025-07-30 01:00:00 发布

原创最新推荐文章于 2025-07-30 01:00:00 发布 · 4.6k 阅读

4 ·

CC 4.0 BY-SA版权

文章标签：

#estimator #tensorflow #deep learning

深度学习同时被 2 个专栏收录

1 篇文章

订阅专栏

Tensorflow

1 篇文章

订阅专栏

本文详述了如何使用TensorFlow Estimator进行模型训练后的导出流程，包括serving_input_receiver_fn()与export_outputs的定义及应用，最终实现模型的pb文件输出，便于后续服务部署。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

本人原文简书博客地址为: https://www.jianshu.com/p/72058da4d7f7

使用场景

利用estimator进行完模型训练后，我们需要将模型导出成pb文件，以便日后服务部署。在部署serving之前，我们需要先将模型导出成 standard SavedModel format。

利用estimator 进行模型导出，只需要指明模型的输入和输出，然后导出即可。主要涉及到的方法为：serving_input_receiver_fn(), export_outputs, tf.estimator.Estimator.export_savedmodel

博客中所构建代码，之后会上传至github中，方便大家使用estimator 进行模型导出服务。

1. `serving_input_receiver_fn()`官方API描述

在训练期间，input_fn()摄取数据并准备供模型使用。在服务时，类似地， serving_input_receiver_fn()接受推理请求并为模型准备它们。此功能具有以下用途：

要向图表添加占位符，服务系统将使用推理请求进行提供。
添加将输入格式的数据转换Tensor为模型预期的功能所需的任何其他操作。

该函数返回一个tf.estimator.export.ServingInputReceiver对象，该对象将占位符和结果特征打包Tensor在一起。

2. `serving_input_receiver_fn()`实现示例

tensorflow 提供了好几种serving_input_fn()。在这里我使tf.estimator.export.ServingInputReceiver进行构建输入方法。
在代码中，features指代传入到model_fn()的 feature dict。也就是传入的特征字典。其中的images对应本人使用的模型中对应的images tensor。而receiver_tensors 指代我们要传入的数据字典。中间可以添加功能块，对传入的数据tensor进行处理，转换成model_fn中需要的feature。

def raw_serving_input_fn():
    serialized_tf_example = tf.placeholder(tf.float32, shape=[None, FLAGS.train_image_size,FLAGS.train_image_size,3], name="images")
    features = {"images": serialized_tf_example}
    receiver_tensors = {'predictor_inputs': serialized_tf_example}
    return tf.estimator.export.ServingInputReceiver(features, receiver_tensors)

3. `export_outputs`官方API描述

编写自定义时model_fn，必须填充返回值的export_outputs元素tf.estimator.EstimatorSpec。这是 {name: output}描述在服务期间要导出和使用的输出签名的单词。

在进行单个预测的通常情况下，此dict包含一个元素，并且它name是无关紧要的。在多头模型中，每个头部由该词典中的条目表示。在这种情况下name，您可以选择一个字符串，用于在服务时请求特定的头部。

每个output值必须是一个ExportOutput对象，例如 tf.estimator.export.ClassificationOutput，tf.estimator.export.RegressionOutput，或 tf.estimator.export.PredictOutput。

4. `export_outputs`实现示例

订好好输出的字典后，我们需要将该输出字典传入到tf.estimator.EstimatorSpec的export_outputs,以确定该estimator的导出模型的输出内容。

    predictions = {'pred_x_ratio': pred_x_ratio, 'pred_y_ratio': pred_y_ratio, 'pred_v': visiable_pre_argmax}

    output = {'serving_default': tf.estimator.export.PredictOutput(predictions)}
    if mode == tf.estimator.ModeKeys.PREDICT:
        return tf.estimator.EstimatorSpec(
                              mode=mode,
                              predictions=predictions,
                              loss=None, train_op=None,
        export_outputs=output )

5. 执行模型导出

通过export_savedmodel导出模型即可。

estimator.export_savedmodel(export_dir_base, serving_input_receiver_fn,
                            strip_default_attrs=True)

6. 检查输出的模型结构

进入导出的模型文件夹路径，通过官方提供的指令查看图输出结果

saved_model_cli show --dir . --all
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | \
  sudo apt-key add -

本人模型的输出图结果如下，对应着以上的输入输出定义结果。

MetaGraphDef with tag-set: 'serve' contains the following SignatureDefs:

signature_def['serving_default']:
  The given SavedModel SignatureDef contains the following input(s):
    inputs['predictor_inputs'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 384, 384, 3)
        name: images:0
  The given SavedModel SignatureDef contains the following output(s):
    outputs['pred_v'] tensor_info:
        dtype: DT_INT64
        shape: (-1, 24)
        name: Identity_2:0
    outputs['pred_x_ratio'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 24)
        name: Reshape_1:0
    outputs['pred_y_ratio'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 24)
        name: Reshape_2:0
  Method name is: tensorflow/serving/predict

参考文档：

https://tensorflow.google.cn/versions/r1.9/guide/saved_model?hl=en#using_savedmodel_with_estimators