使用 SavedModel 格式

在 TensorFlow.org 上查看 在 Google Colab 中运行 在 GitHub 上查看源码 下载笔记本

SavedModel 包含完整的 TensorFlow 程序，包括已训练的参数（即 tf.Variable）和计算逻辑。它在运行时不需要原始的模型构建代码，这使得它非常适合通过 TFLite、TensorFlow.js、TensorFlow Serving 或 TensorFlow Hub 进行共享或部署。

您可以使用以下 API 以 SavedModel 格式保存和加载模型：

从 Keras 创建 SavedModel

为了快速入门，本节将导出预训练的 Keras 模型并使用它处理图像分类请求。本指南的其余部分将填补细节并讨论创建 SavedModel 的其他方法。

import os
import tempfile

from matplotlib import pyplot as plt
import numpy as np
import tensorflow as tf

tmpdir = tempfile.mkdtemp()

physical_devices = tf.config.list_physical_devices('GPU')
for device in physical_devices:
  tf.config.experimental.set_memory_growth(device, True)

file = tf.keras.utils.get_file(
    "grace_hopper.jpg",
    "https://storage.googleapis.com/download.tensorflow.org/example_images/grace_hopper.jpg")
img = tf.keras.utils.load_img(file, target_size=[224, 224])
plt.imshow(img)
plt.axis('off')
x = tf.keras.utils.img_to_array(img)
x = tf.keras.applications.mobilenet.preprocess_input(
    x[tf.newaxis,...])

我们将以格蕾丝·霍珀（Grace Hopper）的图片作为运行示例，并使用一个易于使用的 Keras 预训练图像分类模型。自定义模型同样适用，稍后将详细介绍。

labels_path = tf.keras.utils.get_file(
    'ImageNetLabels.txt',
    'https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt')
imagenet_labels = np.array(open(labels_path).read().splitlines())

pretrained_model = tf.keras.applications.MobileNet()
result_before_save = pretrained_model(x)

decoded = imagenet_labels[np.argsort(result_before_save)[0,::-1][:5]+1]

print("Result before saving:\n", decoded)

该图像的最高预测结果是“军装（military uniform）”。

mobilenet_save_path = os.path.join(tmpdir, "mobilenet/1/")
tf.saved_model.save(pretrained_model, mobilenet_save_path)

保存路径遵循 TensorFlow Serving 使用的约定，其中路径的最后一个组件（此处为 1/）是模型的版本号——这允许像 TensorFlow Serving 这样的工具判断模型的相对时效性。

您可以使用 tf.saved_model.load 将 SavedModel 加载回 Python，并查看霍珀上将的图片是如何被分类的。

loaded = tf.saved_model.load(mobilenet_save_path)
print(list(loaded.signatures.keys()))  # ["serving_default"]

导入的签名总是返回字典。要自定义签名名称和输出字典键，请参阅 导出时指定签名。

infer = loaded.signatures["serving_default"]
print(infer.structured_outputs)

从 SavedModel 运行推理得到的结果与原始模型相同。

labeling = infer(tf.constant(x))[pretrained_model.output_names[0]]

decoded = imagenet_labels[np.argsort(labeling)[0,::-1][:5]+1]

print("Result after saving and loading:\n", decoded)

在 TensorFlow Serving 中运行 SavedModel

SavedModel 可用于 Python（详见下文），但生产环境通常使用专用服务进行推理，而无需运行 Python 代码。使用 TensorFlow Serving 可以轻松从 SavedModel 设置此项服务。

有关完整的 tensorflow-serving 示例，请参阅 TensorFlow Serving REST 教程。

磁盘上的 SavedModel 格式

SavedModel 是一个包含序列化签名以及运行它们所需状态（包括变量值和词汇表）的目录。

ls {mobilenet_save_path}

saved_model.pb 文件存储了实际的 TensorFlow 程序（或模型）以及一组命名签名，每个签名标识了一个接收张量输入并产生张量输出的函数。

SavedModel 可能包含模型的多个变体（多个 v1.MetaGraphDef，通过 saved_model_cli 的 --tag_set 标志标识），但这很少见。创建模型多个变体的 API 包括 tf.Estimator.experimental_export_all_saved_models 以及 TensorFlow 1.x 中的 tf.saved_model.Builder。

saved_model_cli show --dir {mobilenet_save_path} --tag_set serve

variables 目录包含一个标准的训练检查点（请参阅 训练检查点指南）。

ls {mobilenet_save_path}/variables

assets 目录包含 TensorFlow 图使用的文件，例如用于初始化词汇表的文本文件。在本示例中未使用此目录。

SavedModel 可能有一个 assets.extra 目录，用于存放 TensorFlow 图未使用的任何文件，例如关于如何使用该 SavedModel 的说明信息。TensorFlow 本身不会使用此目录。

fingerprint.pb 文件包含 SavedModel 的指纹，由多个唯一标识 SavedModel 内容的 64 位哈希值组成。指纹 API 目前处于实验阶段，但可以使用 tf.saved_model.experimental.read_fingerprint 将 SavedModel 指纹读取为 tf.saved_model.experimental.Fingerprint 对象。

保存自定义模型

tf.saved_model.save 支持保存 tf.Module 对象及其子类，例如 tf.keras.Layer 和 tf.keras.Model。

让我们看一个保存和恢复 tf.Module 的示例。

class CustomModule(tf.Module):

  def __init__(self):
    super(CustomModule, self).__init__()
    self.v = tf.Variable(1.)

  @tf.function
  def __call__(self, x):
    print('Tracing with', x)
    return x * self.v

  @tf.function(input_signature=[tf.TensorSpec([], tf.float32)])
  def mutate(self, new_v):
    self.v.assign(new_v)

module = CustomModule()

当您保存 tf.Module 时，所有 tf.Variable 属性、tf.function 修饰的方法以及通过递归遍历找到的 tf.Module 都会被保存。（有关此递归遍历的更多信息，请参阅检查点教程。）但是，任何 Python 属性、函数和数据都会丢失。这意味着当保存 tf.function 时，不会保存任何 Python 代码。

如果没有保存 Python 代码，SavedModel 如何知道如何恢复该函数？

简而言之，tf.function 的工作原理是跟踪 Python 代码以生成 ConcreteFunction（tf.Graph 的可调用包装器）。当保存 tf.function 时，实际上保存的是 tf.function 的 ConcreteFunction 缓存。

要详细了解 tf.function 与 ConcreteFunctions 之间的关系，请参阅 tf.function 指南。

module_no_signatures_path = os.path.join(tmpdir, 'module_no_signatures')
module(tf.constant(0.))
print('Saving model...')
tf.saved_model.save(module, module_no_signatures_path)

加载并使用自定义模型

当您在 Python 中加载 SavedModel 时，所有的 tf.Variable 属性、tf.function 修饰的方法以及 tf.Module 都会以与原始保存的 tf.Module 相同的对象结构进行恢复。

imported = tf.saved_model.load(module_no_signatures_path)
assert imported(tf.constant(3.)).numpy() == 3
imported.mutate(tf.constant(2.))
assert imported(tf.constant(3.)).numpy() == 6

由于没有保存 Python 代码，使用新的输入签名调用 tf.function 将会失败。

imported(tf.constant([3.]))

ValueError: Could not find matching function to call for canonicalized inputs ((,), {}). Only existing signatures are [((TensorSpec(shape=(), dtype=tf.float32, name=u'x'),), {})].

基础微调

变量对象是可用的，并且您可以对导入的函数进行反向传播。这足以在简单情况下微调（即重新训练）一个 SavedModel。

optimizer = tf.keras.optimizers.SGD(0.05)

def train_step():
  with tf.GradientTape() as tape:
    loss = (10. - imported(tf.constant(2.))) ** 2
  variables = tape.watched_variables()
  grads = tape.gradient(loss, variables)
  optimizer.apply_gradients(zip(grads, variables))
  return loss

for _ in range(10):
  # "v" approaches 5, "loss" approaches 0
  print("loss={:.2f} v={:.2f}".format(train_step(), imported.v.numpy()))

通用微调

来自 Keras 的 SavedModel 提供了更多细节，而不仅仅是简单的 __call__，以解决更复杂的微调案例。TensorFlow Hub 建议在为微调目的而共享的 SavedModel 中提供以下内容（如果适用）：

  • 如果模型在训练和推理期间的前向传播过程不同（例如使用 dropout 或批归一化），则 __call__ 方法接受一个可选的、Python 值的 training= 参数，默认为 False，但可以设置为 True。
  • 除了 __call__ 属性外，还有 .variable 和 .trainable_variable 属性，包含相应的变量列表。原本可训练但希望在微调期间冻结的变量会从 .trainable_variables 中省略。
  • 为了支持像 Keras 这样将权重正则化项表示为层或子模型属性的框架，还可以有一个 .regularization_losses 属性。它包含一个零参数函数列表，这些函数的值用于加到总损失中。

回到最初的 MobileNet 示例，您可以看到其中一些功能的作用：

loaded = tf.saved_model.load(mobilenet_save_path)
print("MobileNet has {} trainable variables: {}, ...".format(
          len(loaded.trainable_variables),
          ", ".join([v.name for v in loaded.trainable_variables[:5]])))

trainable_variable_ids = {id(v) for v in loaded.trainable_variables}
non_trainable_variables = [v for v in loaded.variables
                           if id(v) not in trainable_variable_ids]
print("MobileNet also has {} non-trainable variables: {}, ...".format(
          len(non_trainable_variables),
          ", ".join([v.name for v in non_trainable_variables[:3]])))

导出时指定签名

像 TensorFlow Serving 和 saved_model_cli 这样的工具可以与 SavedModel 进行交互。为了帮助这些工具确定使用哪些 ConcreteFunction，您需要指定服务签名。 tf.keras.Model 会自动指定服务签名，但您必须为我们的自定义模块显式声明服务签名。

默认情况下，自定义 tf.Module 中不会声明任何签名。

assert len(imported.signatures) == 0

要声明服务签名，请使用 signatures 关键字参数指定一个 ConcreteFunction。当指定单个签名时，其签名键将为 'serving_default'，它被保存为常量 tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY。

module_with_signature_path = os.path.join(tmpdir, 'module_with_signature')
call = module.__call__.get_concrete_function(tf.TensorSpec(None, tf.float32))
tf.saved_model.save(module, module_with_signature_path, signatures=call)

imported_with_signatures = tf.saved_model.load(module_with_signature_path)
list(imported_with_signatures.signatures.keys())

要导出多个签名，请传递一个签名键到 ConcreteFunctions 的字典。每个签名键对应一个 ConcreteFunction。

module_multiple_signatures_path = os.path.join(tmpdir, 'module_with_multiple_signatures')
signatures = {"serving_default": call,
              "array_input": module.__call__.get_concrete_function(tf.TensorSpec([None], tf.float32))}

tf.saved_model.save(module, module_multiple_signatures_path, signatures=signatures)

imported_with_multiple_signatures = tf.saved_model.load(module_multiple_signatures_path)
list(imported_with_multiple_signatures.signatures.keys())

默认情况下，输出张量的名称非常通用，例如 output_0。要控制输出的名称，请修改您的 tf.function，使其返回一个将输出名称映射到输出的字典。输入的名称源自 Python 函数的参数名称。

class CustomModuleWithOutputName(tf.Module):
  def __init__(self):
    super(CustomModuleWithOutputName, self).__init__()
    self.v = tf.Variable(1.)

  @tf.function(input_signature=[tf.TensorSpec(None, tf.float32)])
  def __call__(self, x):
    return {'custom_output_name': x * self.v}

module_output = CustomModuleWithOutputName()
call_output = module_output.__call__.get_concrete_function(tf.TensorSpec(None, tf.float32))
module_output_path = os.path.join(tmpdir, 'module_with_output_name')
tf.saved_model.save(module_output, module_output_path,
                    signatures={'serving_default': call_output})

imported_with_output_name = tf.saved_model.load(module_output_path)
imported_with_output_name.signatures['serving_default'].structured_outputs

Proto 分割

由于 protobuf 实现的限制，proto 大小不能超过 2GB。在尝试保存非常大的模型时，这可能会导致以下错误：

ValueError: Message tensorflow.SavedModel exceeds maximum protobuf size of 2GB: ...

google.protobuf.message.DecodeError: Error parsing message as the message exceeded the protobuf limit with type 'tensorflow.GraphDef'

如果您希望保存超过 2GB 限制的模型，则需要使用新的 proto 分割选项进行保存：

tf.saved_model.save(
  ...,
  options=tf.saved_model.SaveOptions(experimental_image_format=True)
)

更多信息可以在 Proto 分割/合并库指南 中找到。

在 C++ 中加载 SavedModel

SavedModel 的 C++ 加载器 提供了一个从路径加载 SavedModel 的 API，同时允许使用 SessionOptions 和 RunOptions。您必须指定与要加载的图相关的标签。加载的 SavedModel 版本称为 SavedModelBundle，包含 MetaGraphDef 以及加载它的会话。

const string export_dir = ...
SavedModelBundle bundle;
...
LoadSavedModel(session_options, run_options, export_dir, {kSavedModelTagTrain},
               &bundle);

SavedModel 命令行界面详细信息

您可以使用 SavedModel 命令行界面 (CLI) 来检查和执行 SavedModel。例如，您可以使用 CLI 检查模型的 SignatureDef。CLI 使您能够快速确认输入张量的数据类型 (dtype) 和形状是否与模型匹配。此外，如果您想测试模型，可以使用 CLI 通过传入各种格式（例如 Python 表达式）的样本输入并获取输出来进行完整性检查。

安装 SavedModel CLI

通常情况下，您可以通过以下两种方式之一安装 TensorFlow：

  • 安装预构建的 TensorFlow 二进制文件。
  • 从源代码构建 TensorFlow。

如果您通过预构建的 TensorFlow 二进制文件安装了 TensorFlow，则 SavedModel CLI 已经安装在您系统中的 bin/saved_model_cli 路径下。

如果您是从源代码构建 TensorFlow 的，则必须运行以下附加命令来构建 saved_model_cli：

$ bazel build //tensorflow/python/tools:saved_model_cli

命令概览

SavedModel CLI 在 SavedModel 上支持以下两个命令：

  • show，用于显示 SavedModel 中可用的计算。
  • run，用于运行 SavedModel 中的计算。

show 命令

一个 SavedModel 包含一个或多个模型变体（技术上为 v1.MetaGraphDef），由其标签集（tag-sets）标识。要部署模型，您可能想知道每个模型变体中有哪些类型的 SignatureDef，以及它们的输入和输出是什么。 show 命令让您可以按层次结构检查 SavedModel 的内容。语法如下：

usage: saved_model_cli show [-h] --dir DIR [--all]
[--tag_set TAG_SET] [--signature_def SIGNATURE_DEF_KEY]

例如，以下命令显示 SavedModel 中所有可用的标签集：

$ saved_model_cli show --dir /tmp/saved_model_dir
The given SavedModel contains the following tag-sets:
serve
serve, gpu

以下命令显示标签集的所有可用 SignatureDef 键：

$ saved_model_cli show --dir /tmp/saved_model_dir --tag_set serve
The given SavedModel `MetaGraphDef` contains `SignatureDefs` with the
following keys:
SignatureDef key: "classify_x2_to_y3"
SignatureDef key: "classify_x_to_y"
SignatureDef key: "regress_x2_to_y3"
SignatureDef key: "regress_x_to_y"
SignatureDef key: "regress_x_to_y2"
SignatureDef key: "serving_default"

如果标签集中有多个标签，则必须指定所有标签，每个标签用逗号分隔。例如：

$ saved_model_cli show --dir /tmp/saved_model_dir --tag_set serve,gpu

要显示特定 SignatureDef 的所有输入和输出 TensorInfo，请将 SignatureDef 键传递给 signature_def 选项。这对于想要了解执行计算图所需的输入张量的张量键值、数据类型和形状非常有用。例如：

$ saved_model_cli show --dir \
/tmp/saved_model_dir --tag_set serve --signature_def serving_default
The given SavedModel SignatureDef contains the following input(s):
  inputs['x'] tensor_info:
      dtype: DT_FLOAT
      shape: (-1, 1)
      name: x:0
The given SavedModel SignatureDef contains the following output(s):
  outputs['y'] tensor_info:
      dtype: DT_FLOAT
      shape: (-1, 1)
      name: y:0
Method name is: tensorflow/serving/predict

要显示 SavedModel 中的所有可用信息，请使用 --all 选项。例如：

$ saved_model_cli show --dir /tmp/saved_model_dir --all
MetaGraphDef with tag-set: 'serve' contains the following SignatureDefs:

signature_def['classify_x2_to_y3']:
  The given SavedModel SignatureDef contains the following input(s):
    inputs['inputs'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 1)
        name: x2:0
  The given SavedModel SignatureDef contains the following output(s):
    outputs['scores'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 1)
        name: y3:0
  Method name is: tensorflow/serving/classify

...

signature_def['serving_default']:
  The given SavedModel SignatureDef contains the following input(s):
    inputs['x'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 1)
        name: x:0
  The given SavedModel SignatureDef contains the following output(s):
    outputs['y'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 1)
        name: y:0
  Method name is: tensorflow/serving/predict

run 命令

调用 run 命令以运行图计算，传入输入并显示（也可选择保存）输出。语法如下：

usage: saved_model_cli run [-h] --dir DIR --tag_set TAG_SET --signature_def
                           SIGNATURE_DEF_KEY [--inputs INPUTS]
                           [--input_exprs INPUT_EXPRS]
                           [--input_examples INPUT_EXAMPLES] [--outdir OUTDIR]
                           [--overwrite] [--tf_debug]

run 命令提供了以下三种向模型传递输入的方式：

  • --inputs 选项使您可以传递文件中的 numpy ndarray。
  • --input_exprs 选项使您可以传递 Python 表达式。
  • --input_examples 选项使您可以传递 tf.train.Example。

--inputs

要传递文件中的输入数据，请指定 --inputs 选项，其一般格式如下：

--inputs <INPUTS>

其中 INPUTS 是以下任一格式：

  • <input_key>=<filename>
  • <input_key>=<filename>[<variable_name>]

您可以传递多个 INPUTS。如果传递多个输入，请使用分号分隔每个 INPUTS。

saved_model_cli 使用 numpy.load 加载 filename。filename 可以采用以下任一格式：

  • .npy
  • .npz
  • pickle 格式

.npy 文件总是包含一个 numpy ndarray。因此，从 .npy 文件加载时，内容将直接分配给指定的输入张量。如果您为该 .npy 文件指定了 variable_name，则 variable_name 将被忽略并会发出警告。

从 .npz (zip) 文件加载时，您可以选择指定 variable_name 以标识 zip 文件中要为输入张量键加载的变量。如果您不指定 variable_name，SavedModel CLI 将检查 zip 文件中是否只包含一个文件，并将其加载到指定的输入张量键。

从 pickle 文件加载时，如果方括号中未指定 variable_name，则 pickle 文件中的任何内容都将传递给指定的输入张量键。否则，SavedModel CLI 将假定 pickle 文件中存储的是一个字典，并使用与 variable_name 对应的值。

--input_exprs

要通过 Python 表达式传递输入，请指定 --input_exprs 选项。当您手头没有数据文件，但仍想用一些简单的输入（匹配模型 SignatureDef 的数据类型和形状）对模型进行完整性检查时，这非常有用。例如：

`<input_key>=[[1],[2],[3]]`

除了 Python 表达式，您还可以传递 numpy 函数。例如：

`<input_key>=np.ones((32,32,3))`

（请注意，numpy 模块已经以 np 的形式提供给您。）

--input_examples

要传递 tf.train.Example 作为输入，请指定 --input_examples 选项。对于每个输入键，它接受一个字典列表，其中每个字典都是一个 tf.train.Example 实例。字典键是特征，值是每个特征的值列表。例如：

`<input_key>=[{"age":[22,24],"education":["BS","MS"]}]`

保存输出

默认情况下，SavedModel CLI 将输出写入 stdout。如果将目录传递给 --outdir 选项，输出将保存为给定目录下以输出张量键命名的 .npy 文件。

使用 --overwrite 来覆盖现有的输出文件。