SavedModel 签名之旅
2021 年 3 月 3 日

由 TensorFlow 工程师 Daniel Ellis 撰写

注意:这篇博文针对希望了解图形和模型存储细节的 TensorFlow 开发人员。如果您是 TensorFlow 新手,请在阅读本文之前查看 TensorFlow 基础指南

TensorFlow 可以运行模型,而无需原始的 Python 对象,如 TensorFlow ServingTensorFlow Lite 所示,或者当您从 TensorFlow Hub 下载训练好的模型时。

模型和层可以从这种表示中加载,而无需实际创建创建它的 Python 类的实例。这在您没有(或不需要)Python 解释器的情况下是可取的,例如大规模服务或在边缘设备上服务,或者在原始 Python 代码不可用时。

保存的模型由两个独立但同样重要的部分表示:图形,它描述了代码中描述的固定计算,以及权重,它是您在训练期间训练的动态参数。如果您不熟悉这一点和 @tf.function,您应该查看 图形和函数入门指南 以及 关于保存的部分,该部分位于 模块、层和模型指南 中。

从代码的角度来看,用 @tf.function 装饰的函数创建了一个 Python 可调用对象;在 文档 中,我们将其称为多态函数,因为它们是 Python 可调用对象,可以接受各种参数签名。每次您使用新的参数签名调用 @tf.function 时,TensorFlow 都会为该参数集追踪一个新的图形。然后,这个新的图形将作为“具体函数”添加到可调用对象中。因此,保存的模型可以是一个或多个子图,每个子图都有不同的签名。

当您调用 tf.saved_model.save() 时,您将获得一个 SavedModel。保存的模型存储为磁盘上的目录。该目录中的文件 saved_model.pb 是一个 协议缓冲区,它描述了函数 tf.Graph

在这篇博文中,我们将深入了解这个 protobuf,看看函数签名序列化和反序列化是如何在幕后工作的。阅读完本文后,您将更加了解函数和签名,这可以帮助您加载、修改或优化保存的模型。

背景

在保存的模型 protobuf 中,总共有五个地方定义了函数的输入。理解和记住每个定义的作用可能很困难。这篇文章旨在列出每个定义以及它们的使用目的。它还将介绍一个基本示例,说明一个简单的模型在序列化后的样子。

您实际使用的 API 将始终经过仔细的版本控制(自 2016 年以来一直如此),并且模型本身将符合 版本兼容性指南。但是,本文中的内容介绍了现有状态的快照。任何指向代码的链接都将包含时间点修订,以避免过时。与所有未记录的实现细节一样,这些细节可能会在将来发生变化。

我们有时会使用“签名”一词来谈论描述函数输入的一般概念(例如本文的标题)。从这个意义上说,我们将不仅仅指的是 TensorFlow 的特定签名概念,而是 TensorFlow 定义和验证函数输入的所有方法。上下文应该使含义清晰。

本文不涉及的内容

本文并非旨在描述签名或函数从用户角度的工作方式。它旨在供从事 TensorFlow 内部工作的 TensorFlow 开发人员使用。同样,本文没有说明“应该”如何做。它的目标只是记录现状。

签名定义概述

总共有五个 proto 以某种方式存储函数输入的定义。它们的名字和代码位置,以及它们在保存的模型 proto 中的路径,如下所示

Proto 消息,以及它们在 SavedModel 中的位置

FunctionDef

在这篇文档中讨论的五个定义中,FunctionsDefs 是执行中最核心的定义。加载保存的模型时,这些函数定义会注册在运行时的函数库中,并用于创建 ConcreteFunctions。然后,这些函数可以通过 PartitionedCallTFE_Py_Execute 执行。

这是定义实际描述执行的节点的地方,以及函数的输入和输出是什么。

SignatureDef

SignatureDefs从传递到 @tf.function 的签名生成的。但是,我们不会直接保存签名的 TensorSpecs。相反,在保存时,我们会使用 TensorSpecs 调用底层函数,以生成一个具体函数。然后,我们检查生成的具体函数以获取输入和输出,并将它们存储在 SignatureDef 上。

在加载端,SignatureDefs基本上会被忽略。它们主要在 v1 或 C++ 中使用,在这些环境中,加载模型的开发人员可以直接检查返回的 SignatureDef proto。这允许他们使用他们想要的签名名称来查找执行所需的占位符和输出名称。

然后,这些输入和输出名称可以作为馈送和提取传递给 TensorFlow V1 代码中调用 Session.run

SavedFunction

SavedFunctionSavedObjects 的众多类型之一,位于 ObjectGraphDef 的节点列表中。SavedFunctions 在加载时恢复到 RestoredFunctions。与该列表中的所有节点一样,它们然后通过子级 ObjectReference 字段定义的层次结构附加到返回的模型。

SavedFunction的主要目的是多态性SavedFunctions通过指定函数库中定义的多个具体函数名称(通过FunctionDef)来支持多态性。在调用时,我们遍历具体函数名称以找到第一个与签名匹配的函数。如果我们找到匹配项,则调用它;否则,我们将抛出异常。

还有一个复杂之处。当用一组特定的参数调用RestoredFunction时,会创建一个新的具体函数,其唯一目的是调用匹配的具体函数。这是使用restored_function_body在后台完成的,并且是查找适当的具体函数的逻辑所在。

这在SavedModel协议缓冲区中是不可见的,但这些额外的具体函数在运行时在运行时的函数库中注册,就像其他函数库函数一样。

SavedFunction的第二个目的是使用存储在SavedFunction上的FunctionSpec更新所有关联的ConcreteFunctionsFunctionSpec。这个函数规范在调用时用于

  1. 验证传入的结构化参数,以及
  2. 将结构化参数转换为调用底层具体函数所需的扁平参数

SavedBareConcreteFunction

类似于SavedFunctionsSavedBareConcreteFunctions用于更新一个

特定的具体函数的参数和函数规范。这是通过这里完成的。与SavedFunctions不同,它们只引用一个特定的具体函数

在实践中,SavedBareConcreteFunctions通常附加到签名映射并通过签名映射访问(即加载对象上的签名属性)。它们修改的底层具体函数在这种情况下是signature_wrapper函数。这种包装是为了以v1 期望的方式(即张量字典)格式化输出。类似于restored_function_body具体函数,除了重构输出外,这些具体函数除了调用其关联的具体函数之外什么也不做。

SavedConcreteFunction

SavedConcreteFunction对象不是SavedObjectGraph节点。它们存储在映射中,直接位于SavedObjectGraph上。这些对象引用一个特定的、已经注册的具体函数 - 映射中的键是该具体函数的注册名称。

这些对象有两个目的。第一个是通过

bound_inputs字段处理函数“捕获”。捕获的变量是函数读取或修改但调用函数时未显式传入的变量。由于函数库中的函数没有捕获变量的概念,因此函数使用的任何变量都必须作为参数传入。bound_inputs存储一个节点 ID 列表,这些节点 ID 应在调用时传入底层ConcreteFunction。我们在这里设置了它。

第二个目的,与SavedFunctionSavedBareConcreteFunction类似,是修改现有具体函数的FuncGraph结构化输入和输出。这也用于参数验证。此设置是在这里完成的。

示例演练

一个简单的示例可能有助于更清楚地说明所有这些内容。让我们创建一个基本模型并查看随后生成的协议缓冲区,以便更好地了解正在发生的事情。

基本模型

class ExampleModel(tf.Module):

  @tf.function(input_signature=[tf.TensorSpec(shape=(), dtype=tf.float32)])
  def capture_fn(self, x):
    if not hasattr(self, 'weight'):
      self.weight = tf.Variable(5.0, name='weight')
    self.weight.assign_add(x * self.weight)
    return self.weight

  @tf.function
  def polymorphic_fn(self, x):
    return tf.constant(3.0) * x

model = ExampleModel()
model.polymorphic_fn(tf.constant(4.0))
model.polymorphic_fn(tf.constant([1.0, 2.0, 3.0]))
tf.saved_model.save(
    model, "/tmp/example-model", signatures={'capture_fn': model.capture_fn})

此模型包含我们探索保存和签名的复杂性的基础。这将允许我们查看有和没有签名的函数,有和没有捕获的函数,以及有和没有多态性的函数。

带有捕获的函数

让我们从查看带有捕获的函数capture_fn开始。我们可以看到我们如期在函数库中定义了一个具体函数

Image of concrete function defined in the function library
一个位于MetaGraphDef.graph_defFunctionDefLibrary中的FunctionDef

请注意预期的浮点输入"x",以及额外的捕获参数"mul_readvariableop_resource"。由于此函数具有捕获,因此我们应该看到一个变量在SavedConcreteFunctions之一的bound_inputs字段中被引用

SavedConcreteFunctions
一个位于ObjectGraphDefconcrete_functions映射中的SavedConcreteFunction

事实上,我们可以看到bound_inputs引用节点1,它是一个名为SavedVariable,具有我们期望的名称和数据类型的变量

A `SavedVariable` located in `ObjectGraphDef.nodes`
一个位于ObjectGraphDef.nodes中的SavedVariable

请注意,我们还在canonicalized_input_signature上存储了将在修改具体函数时使用的数据。此对象的键"__inference_capture_fn_59"与我们函数库中注册的具体函数的名称相同。

由于我们已经指定了签名,因此我们还应该看到一个SavedBareConcreteFunction

SavedBareConcreteFunction
一个位于ObjectGraphDef.nodes中的SavedBareConcreteFunction

如上所述,我们使用函数规范和参数信息来修改底层具体函数。但是"__inference_signature_wrapper_68"这个名字是怎么回事?它如何与代码的其他部分相适应?

首先,请注意这是节点列表中的第五个(5)节点。这一点稍后会再次出现。

让我们从查看节点列表开始。如果从节点列表中的第一个节点开始,我们将看到一个名为"signatures"的节点作为子节点附加

SavedUserObject
一个位于ObjectGraphDef.nodes中的SavedUserObject

如果我们查看节点2,我们将看到此节点是一个签名映射,它引用了最后一个节点:节点5,即我们的BareConcreteSavedFunction

Node5
一个位于ObjectGraphDef.nodes中的SavedUserObject

因此,当我们通过model.signatures["capture_fn"]访问此函数时,实际上将首先调用此中间签名包装函数。

那么那个函数"__inference_signature_wrapper_68"是什么样子的呢?

FunctionDef
一个位于MetaGraphDef.graph_defFunctionDefLibrary中的FunctionDef

它接受我们期望的参数,并调用…"__inference_capture_fn_59",即我们的原始函数!正如我们所预期的那样。

但是等等…如果我们不通过model.signatures["capture_fn"]访问我们的函数怎么办?毕竟,我们应该能够通过model.capture_fn直接调用它。

请注意上面,我们在顶级对象上有一个名为"capture_fn"的子节点,其node_id3。如果我们查看节点3,我们将看到一个SavedFunction对象,它引用了我们的原始具体函数,没有签名包装中介

Node 3
一个位于ObjectGraphDef.nodes中的SavedFunction

同样,函数规范用于修改我们具体函数"__inference_capture_fn_59"的函数规范。另外请注意,这里concrete_functions是一个列表。我们现在只有一项,但这将在我们查看多态函数示例时再次出现。

现在,我们已经基本映射了执行此函数所需的一切,但还有一件事需要查看:SignatureDef。我们已经定义了一个签名,因此我们期望定义一个SignatureDef

SignatureDef
一个位于MetaObjectGraph.signature_def映射中的SignatureDef

这对在 v1 和 C++ 中加载以进行服务非常重要。请注意那些奇怪的名称:"capture_fn_x:0""StatefulPartitionedCall:0"。为了在 v1 中调用此函数,我们需要一种方法将我们友好的参数名称映射到实际的图形占位符名称,以便作为馈送和提取传递(以及进行验证,如果我们愿意)。查看此SignatureDef使我们能够做到这一点。

多态函数

我们还没有完全完成。让我们看看我们的多态函数。我们不会重复所有内容,因为很多内容都是一样的。由于我们在其中跳过了签名,因此我们将没有任何签名包装函数或签名定义。让我们看看哪些不同。

A Polymorphic FunctionDef
一个位于MetaGraphDef.graph_defFunctionDefLibrary中的FunctionDef

一方面,我们现在在函数库中注册了两个具体函数,每个函数的输入形状略有不同。

我们还有两个SavedConcreteFunction修饰符

Two SavedConcreteFunctions
两个位于ObjectGraphDefconcrete_functions映射中的SavedConcreteFunctions

最后,我们可以看到我们的SavedFunction引用了两个底层具体函数,而不是一个

SavedFunction
一个位于ObjectGraphDef.nodes中的SavedFunction

这里的函数规范将在加载时附加到这两个具体函数。当我们调用SavedFunction时,它将使用我们传入的参数来找到正确的具体函数并执行它。

下一步

你现在应该是函数及其签名如何在代码级别保存的专家。请记住,本文中描述的是当前代码的工作方式。有关将来更新的代码和示例,请参阅tensorflow.org上的官方文档。

说到文档,如果你想快速了解保存模型的基本 API,你应该查看有关函数模块的 API 如何被追踪和保存的介绍性文章。对于专家来说,不要错过这篇关于SavedModel本身的详细指南,以及一篇关于自动图的完整讨论。

最后,如果你做了任何令人兴奋或有用的协议缓冲区手术,请在 Twitter 上与我们分享。感谢你阅读到这里!

下一篇文章
A Tour of SavedModel Signatures

作者:Daniel Ellis,TensorFlow 工程师注意:这篇文章针对的是想要了解图形和模型如何存储的 TensorFlow 开发人员。如果你不熟悉 TensorFlow,在阅读本文之前,你应该查看TensorFlow 基础指南 TensorFlow 可以运行没有原始 Python 对象的模型,正如TensorFlow ServingTensorFlow Lite所证明的那样,o…