迁移 SavedModel 工作流程

在 TensorFlow.org 上查看

在 Google Colab 中运行

在 GitHub 上查看源代码

下载笔记本

将模型从 TensorFlow 1 的图形和会话迁移到 TensorFlow 2 API（例如 tf.function、tf.Module 和 tf.keras.Model）后，您可以迁移模型保存和加载代码。此笔记本提供了有关如何在 TensorFlow 1 和 TensorFlow 2 中使用 SavedModel 格式保存和加载的示例。以下是迁移从 TensorFlow 1 到 TensorFlow 2 的相关 API 更改的快速概述

	TensorFlow 1	迁移到 TensorFlow 2
保存	tf.compat.v1.saved_model.Builder tf.compat.v1.saved_model.simple_save	tf.saved_model.save Keras：tf.keras.models.save_model
加载	tf.compat.v1.saved_model.load	tf.saved_model.load Keras：tf.keras.models.load_model
签名：一组输入 和输出张量， 可用于运行	使用 *.signature_def 实用程序生成 （例如 tf.compat.v1.saved_model.predict_signature_def）	编写一个 tf.function，并使用 signatures 参数导出它 在 tf.saved_model.save 中。
分类 和回归: 特殊类型的签名	使用 tf.compat.v1.saved_model.classification_signature_def, tf.compat.v1.saved_model.regression_signature_def, 以及某些估计器导出生成。	这两种签名类型已从 TensorFlow 2 中删除。 如果服务库需要这些方法名称， tf.compat.v1.saved_model.signature_def_utils.MethodNameUpdater.

有关映射的更深入解释，请参阅下面的 从 TensorFlow 1 到 TensorFlow 2 的更改 部分。

设置

以下示例展示了如何使用 TensorFlow 1 和 TensorFlow 2 API 将相同的虚拟 TensorFlow 模型（在下面定义为 add_two）导出到 SavedModel 格式。首先设置导入和实用程序函数

import tensorflow as tf
import tensorflow.compat.v1 as tf1
import shutil

def remove_dir(path):
  try:
    shutil.rmtree(path)
  except:
    pass

def add_two(input):
  return input + 2

TensorFlow 1：保存和导出 SavedModel

在 TensorFlow 1 中，您可以使用 tf.compat.v1.saved_model.Builder、tf.compat.v1.saved_model.simple_save 和 tf.estimator.Estimator.export_saved_model API 来构建、保存和导出 TensorFlow 图形和会话

1. 将图形保存为带有 SavedModelBuilder 的 SavedModel

remove_dir("saved-model-builder")

with tf.Graph().as_default() as g:
  with tf1.Session() as sess:
    input = tf1.placeholder(tf.float32, shape=[])
    output = add_two(input)
    print("add two output: ", sess.run(output, {input: 3.}))

    # Save with SavedModelBuilder
    builder = tf1.saved_model.Builder('saved-model-builder')
    sig_def = tf1.saved_model.predict_signature_def(
        inputs={'input': input},
        outputs={'output': output})
    builder.add_meta_graph_and_variables(
        sess, tags=["serve"], signature_def_map={
            tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY: sig_def
    })
    builder.save()

!saved_model_cli run --dir saved-model-builder --tag_set serve \
 --signature_def serving_default --input_exprs input=10

2. 为服务构建 SavedModel

remove_dir("simple-save")

with tf.Graph().as_default() as g:
  with tf1.Session() as sess:
    input = tf1.placeholder(tf.float32, shape=[])
    output = add_two(input)
    print("add_two output: ", sess.run(output, {input: 3.}))

    tf1.saved_model.simple_save(
        sess, 'simple-save',
        inputs={'input': input},
        outputs={'output': output})

!saved_model_cli run --dir simple-save --tag_set serve \
 --signature_def serving_default --input_exprs input=10

3. 将估计器推理图形导出为 SavedModel

在估计器 model_fn 的定义中（在下面定义），您可以通过在 tf.estimator.EstimatorSpec 中返回 export_outputs 来在模型中定义签名。输出类型不同

• tf.estimator.export.ClassificationOutput
• tf.estimator.export.RegressionOutput
• tf.estimator.export.PredictOutput

这些将分别生成分类、回归和预测签名类型。

当估计器使用 tf.estimator.Estimator.export_saved_model 导出时，这些签名将与模型一起保存。

def model_fn(features, labels, mode):
  output = add_two(features['input'])
  step = tf1.train.get_global_step()
  return tf.estimator.EstimatorSpec(
      mode,
      predictions=output,
      train_op=step.assign_add(1),
      loss=tf.constant(0.),
      export_outputs={
          tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY: \
          tf.estimator.export.PredictOutput({'output': output})})
est = tf.estimator.Estimator(model_fn, 'estimator-checkpoints')

# Train for one step to create a checkpoint.
def train_fn():
  return tf.data.Dataset.from_tensors({'input': 3.})
est.train(train_fn, steps=1)

# This utility function `build_raw_serving_input_receiver_fn` takes in raw
# tensor features and builds an "input serving receiver function", which
# creates placeholder inputs to the model.
serving_input_fn = tf.estimator.export.build_raw_serving_input_receiver_fn(
    {'input': tf.constant(3.)})  # Pass in a dummy input batch.
estimator_path = est.export_saved_model('exported-estimator', serving_input_fn)

# Estimator's export_saved_model creates a time stamped directory. Move this
# to a set path so it can be inspected with `saved_model_cli` in the cell below.
!rm -rf estimator-model
import shutil
shutil.move(estimator_path, 'estimator-model')

!saved_model_cli run --dir estimator-model --tag_set serve \
 --signature_def serving_default --input_exprs input=[10]

TensorFlow 2：保存和导出 SavedModel

保存和导出使用 tf.Module 定义的 SavedModel

要在 TensorFlow 2 中导出模型，您必须定义一个 tf.Module 或一个 tf.keras.Model 来保存模型的所有变量和函数。然后，您可以调用 tf.saved_model.save 来创建 SavedModel。请参阅 使用 SavedModel 格式 指南中的“保存自定义模型”部分，以了解更多信息。

class MyModel(tf.Module):
  @tf.function
  def __call__(self, input):
    return add_two(input)

model = MyModel()

@tf.function
def serving_default(input):
  return {'output': model(input)}

signature_function = serving_default.get_concrete_function(
    tf.TensorSpec(shape=[], dtype=tf.float32))
tf.saved_model.save(
    model, 'tf2-save', signatures={
        tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature_function})

!saved_model_cli run --dir tf2-save --tag_set serve \
 --signature_def serving_default --input_exprs input=10

保存和导出使用 Keras 定义的 SavedModel

用于保存和导出 Keras API 的—Model.save 或 tf.keras.models.save_model—可以从 tf.keras.Model 导出 SavedModel。查看 保存和加载 Keras 模型 以获取更多详细信息。

inp = tf.keras.Input(3)
out = add_two(inp)
model = tf.keras.Model(inputs=inp, outputs=out)

@tf.function(input_signature=[tf.TensorSpec(shape=[], dtype=tf.float32)])
def serving_default(input):
  return {'output': model(input)}

model.save('keras-model', save_format='tf', signatures={
        tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY: serving_default})

!saved_model_cli run --dir keras-model --tag_set serve \
 --signature_def serving_default --input_exprs input=10

加载 SavedModel

使用上述任何 API 保存的 SavedModel 可以使用 TensorFlow 1 或 TensorFlow 2 API 加载。

通常，当加载到 TensorFlow 2 中时，TensorFlow 1 SavedModel 可用于推理，但只有当 SavedModel 包含资源变量时才可能进行训练（生成梯度）。您可以检查变量的 dtype—如果变量 dtype 包含“_ref”，则它是一个引用变量。

只要 SavedModel 使用签名保存，就可以从 TensorFlow 1 加载和执行 TensorFlow 2 SavedModel。

以下部分包含代码示例，展示如何加载在上一节中保存的 SavedModel，并调用导出的签名。

TensorFlow 1：使用 tf.saved_model.load 加载 SavedModel

在 TensorFlow 1 中，您可以使用 tf.saved_model.load 将 SavedModel 直接导入到当前图和会话中。您可以对张量输入和输出名称调用 Session.run

def load_tf1(path, input):
  print('Loading from', path)
  with tf.Graph().as_default() as g:
    with tf1.Session() as sess:
      meta_graph = tf1.saved_model.load(sess, ["serve"], path)
      sig_def = meta_graph.signature_def[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
      input_name = sig_def.inputs['input'].name
      output_name = sig_def.outputs['output'].name
      print('  Output with input', input, ': ', 
            sess.run(output_name, feed_dict={input_name: input}))

load_tf1('saved-model-builder', 5.)
load_tf1('simple-save', 5.)
load_tf1('estimator-model', [5.])  # Estimator's input must be batched.
load_tf1('tf2-save', 5.)
load_tf1('keras-model', 5.)

TensorFlow 2：加载使用 tf.saved_model 保存的模型

在 TensorFlow 2 中，对象被加载到一个 Python 对象中，该对象存储变量和函数。这与从 TensorFlow 1 保存的模型兼容。

查看 tf.saved_model.load API 文档和 加载和使用自定义模型 部分，了解有关 使用 SavedModel 格式 指南的详细信息。

def load_tf2(path, input):
  print('Loading from', path)
  loaded = tf.saved_model.load(path)
  out = loaded.signatures[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY](
      tf.constant(input))['output']
  print('  Output with input', input, ': ', out)

load_tf2('saved-model-builder', 5.)
load_tf2('simple-save', 5.)
load_tf2('estimator-model', [5.])  # Estimator's input must be batched.
load_tf2('tf2-save', 5.)
load_tf2('keras-model', 5.)

使用 TensorFlow 2 API 保存的模型还可以访问附加到模型的 tf.function 和变量（而不是那些作为签名导出的变量）。例如

loaded = tf.saved_model.load('tf2-save')
print('restored __call__:', loaded.__call__)
print('output with input 5.', loaded(5))

TensorFlow 2：加载使用 Keras 保存的模型

Keras 加载 API—tf.keras.models.load_model—允许您将保存的模型重新加载到 Keras 模型对象中。请注意，这仅允许您加载使用 Keras 保存的 SavedModel (Model.save 或 tf.keras.models.save_model)。

使用 tf.saved_model.save 保存的模型应使用 tf.saved_model.load 加载。您可以使用 tf.saved_model.load 加载使用 Model.save 保存的 Keras 模型，但您只会获得 TensorFlow 图。有关详细信息，请参阅 tf.keras.models.load_model API 文档和 保存和加载 Keras 模型 指南。

loaded_model = tf.keras.models.load_model('keras-model')
loaded_model.predict_on_batch(tf.constant([1, 3, 4]))

GraphDef 和 MetaGraphDef

没有直接的方法可以将原始 GraphDef 或 MetaGraphDef 加载到 TF2 中。但是，您可以使用 v1.wrap_function 将导入图的 TF1 代码转换为 TF2 concrete_function。

首先，保存 MetaGraphDef

# Save a simple multiplication computation:
with tf.Graph().as_default() as g:
  x = tf1.placeholder(tf.float32, shape=[], name='x')
  v = tf.Variable(3.0, name='v')
  y = tf.multiply(x, v, name='y')
  with tf1.Session() as sess:
    sess.run(v.initializer)
    print(sess.run(y, feed_dict={x: 5}))
    s = tf1.train.Saver()
    s.export_meta_graph('multiply.pb', as_text=True)
    s.save(sess, 'multiply_values.ckpt')

使用 TF1 API，您可以使用 tf1.train.import_meta_graph 导入图并恢复值

with tf.Graph().as_default() as g:
  meta = tf1.train.import_meta_graph('multiply.pb')
  x = g.get_tensor_by_name('x:0')
  y = g.get_tensor_by_name('y:0')
  with tf1.Session() as sess:
    meta.restore(sess, 'multiply_values.ckpt')
    print(sess.run(y, feed_dict={x: 5}))

没有用于加载图的 TF2 API，但您仍然可以将其导入到可以在急切模式下执行的具体函数中

def import_multiply():
  # Any graph-building code is allowed here.
  tf1.train.import_meta_graph('multiply.pb')

# Creates a tf.function with all the imported elements in the function graph.
wrapped_import = tf1.wrap_function(import_multiply, [])
import_graph = wrapped_import.graph
x = import_graph.get_tensor_by_name('x:0')
y = import_graph.get_tensor_by_name('y:0')

# Restore the variable values.
tf1.train.Saver(wrapped_import.variables).restore(
    sess=None, save_path='multiply_values.ckpt')

# Create a concrete function by pruning the wrap_function (similar to sess.run).
multiply_fn = wrapped_import.prune(feeds=x, fetches=y)

# Run this function
multiply_fn(tf.constant(5.))  # inputs to concrete functions must be Tensors.

从 TensorFlow 1 到 TensorFlow 2 的更改

本节列出了 TensorFlow 1 中的关键保存和加载术语、它们的 TensorFlow 2 等效项以及发生了哪些变化。

SavedModel

SavedModel 是一种格式，用于存储包含参数和计算的完整 TensorFlow 程序。它包含服务平台用于运行模型的签名。

文件格式本身没有发生重大变化，因此 SavedModel 可以使用 TensorFlow 1 或 TensorFlow 2 API 加载和服务。

TensorFlow 1 和 TensorFlow 2 之间的差异

除了 API 更改之外，服务和推理用例在 TensorFlow 2 中没有更新—改进是在能够重用和组合从 SavedModel 加载的模型中引入的。

在 TensorFlow 2 中，程序由对象表示，例如 tf.Variable、tf.Module 或更高级别的 Keras 模型 (tf.keras.Model) 和层 (tf.keras.layers)。不再有存储在会话中的全局变量，并且图现在存在于不同的 tf.function 中。因此，在模型导出期间，SavedModel 会分别保存每个组件和函数图。

当您使用 TensorFlow Python API 编写 TensorFlow 程序时，您必须构建一个对象来管理变量、函数和其他资源。通常，这是通过使用 Keras API 完成的，但您也可以通过创建或子类化 tf.Module 来构建对象。

Keras 模型 (tf.keras.Model) 和 tf.Module 会自动跟踪附加到它们的变量和函数。SavedModel 会保存模块、变量和函数之间的这些连接，以便它们可以在加载时恢复。

签名

签名是 SavedModel 的端点—它们告诉用户如何运行模型以及需要哪些输入。

在 TensorFlow 1 中，签名是通过列出输入和输出张量创建的。在 TensorFlow 2 中，签名是通过传入具体函数生成的。（在 图和 tf.function 简介 指南中阅读有关 TensorFlow 函数的更多信息，特别是多态性：一个函数，多个图部分。）简而言之，具体函数是从 tf.function 生成的

# Option 1: Specify an input signature.
@tf.function(input_signature=[...])
def fn(...):
  ...
  return outputs

tf.saved_model.save(model, path, signatures={
    'name': fn
})

# Option 2: Call `get_concrete_function`
@tf.function
def fn(...):
  ...
  return outputs

tf.saved_model.save(model, path, signatures={
    'name': fn.get_concrete_function(...)
})

Session.run

在 TensorFlow 1 中，只要您已经知道张量名称，就可以使用导入的图调用 Session.run。这使您可以检索恢复的变量值，或运行未在签名中导出的模型部分。

在 TensorFlow 2 中，您可以直接访问变量，例如权重矩阵 (kernel)

model = tf.Module()
model.dense_layer = tf.keras.layers.Dense(...)
tf.saved_model.save('my_saved_model')
loaded = tf.saved_model.load('my_saved_model')
loaded.dense_layer.kernel

或调用附加到模型对象的 tf.function：例如，loaded.__call__。

与 TF1 不同，没有办法提取函数的一部分并访问中间值。您必须在保存的对象中导出所有需要的功能。

TensorFlow Serving 迁移说明

SavedModel 最初是为了与 TensorFlow Serving 一起使用而创建的。此平台提供不同类型的预测请求：分类、回归和预测。

TensorFlow 1 API 允许您使用以下实用程序创建这些类型的签名

• tf.compat.v1.saved_model.classification_signature_def
• tf.compat.v1.saved_model.regression_signature_def
• tf.compat.v1.saved_model.predict_signature_def

分类 (classification_signature_def) 和 回归 (regression_signature_def) 限制输入和输出，因此输入必须是 tf.Example，输出必须是 classes、scores 或 prediction。同时，预测签名 (predict_signature_def) 没有限制。

使用 TensorFlow 2 API 导出的 SavedModel 与 TensorFlow Serving 兼容，但只会包含预测签名。分类和回归签名已被删除。

如果您需要使用分类和回归签名，您可以使用 tf.compat.v1.saved_model.signature_def_utils.MethodNameUpdater 修改导出的 SavedModel。

下一步

要了解有关 TensorFlow 2 中 SavedModel 的更多信息，请查看以下指南

• 使用 SavedModel 格式
• 保存和加载 Keras 模型

如果您使用的是 TensorFlow Hub，您可能会发现以下指南很有用

• TensorFlow Hub：TensorFlow 1/TensorFlow 2 的模型兼容性
• 使用 TensorFlow Hub 从 TensorFlow 1 迁移到 TensorFlow 2