Python tf.saved_model.save用法及代码示例

导出 tf.Module(和子类)obj到SavedModel 格式.

用法

tf.saved_model.save(
    obj, export_dir, signatures=None, options=None
)

obj 必须继承自 Trackable 类。

示例用法：

class Adder(tf.Module):
  @tf.function(input_signature=[tf.TensorSpec(shape=[], dtype=tf.float32)])
  def add(self, x):
    return x + x

model = Adder()
tf.saved_model.save(model, '/tmp/adder')

然后，生成的 SavedModel 可以使用名为 "x" 的输入进行服务，这是一个 dtype 为 float32 的标量。

签名

签名定义计算的输入和输出类型。可选的 save signatures 参数控制 obj 中的哪些方法可用于消耗 SavedModel 的程序，例如服务 API。 Python 函数可以用 @tf.function(input_signature=...) 装饰并作为签名直接传递，或者在用 @tf.function 装饰的方法上懒惰地调用 get_concrete_function。

例子：

class Adder(tf.Module):
  @tf.function
  def add(self, x):
    return x + x

model = Adder()
tf.saved_model.save(
  model, '/tmp/adder',signatures=model.add.get_concrete_function(
    tf.TensorSpec([], tf.float32)))

如果 @tf.function 没有输入签名并且没有在该方法上调用 get_concrete_function，则该函数将无法在恢复的 SavedModel 中直接调用。

例子：

class Adder(tf.Module):
  @tf.function
  def add(self, x):
    return x + x

model = Adder()
tf.saved_model.save(model, '/tmp/adder')
restored = tf.saved_model.load('/tmp/adder')
restored.add(1.)
Traceback (most recent call last):

ValueError:Found zero restored functions for caller function.

如果省略signatures 参数，则将在obj 中搜索@tf.function 修饰的方法。如果恰好找到一个跟踪的@tf.function，则该方法将用作 SavedModel 的默认签名。否则，任何附加到 obj 的 @tf.function 或其依赖项都将被导出以用于 tf.saved_model.load 。

在导出的 SavedModel 中调用签名时，Tensor 参数由名称标识。默认情况下，这些名称将来自 Python 函数的参数名称。它们可以通过在相应的 tf.TensorSpec 对象中指定 name=... 参数来覆盖。如果多个 Tensor 通过单个参数传递给 Python 函数，则需要显式命名。

用作 signatures 的函数的输出必须是平面列表，在这种情况下输出将被编号，或者是映射字符串键到 Tensor 的字典，在这种情况下，键将用于命名输出。

签名在 tf.saved_model.load 作为 .signatures 属性返回的对象中可用。这是一个保留属性：具有自定义 .signatures 属性的对象上的 tf.saved_model.save 将引发异常。

_将 tf.savedmodel.save 与 Keras 模型一起使用

虽然 Keras 有自己的保存和加载 API，但此函数可用于导出 Keras 模型。例如，使用指定的签名导出：

class Adder(tf.keras.Model):
  @tf.function(input_signature=[tf.TensorSpec(shape=[], dtype=tf.string)])
  def concat(self, x):
     return x + x

model = Adder()
tf.saved_model.save(model, '/tmp/adder')

从没有固定签名的函数导出：

class Adder(tf.keras.Model):
  @tf.function
  def concat(self, x):
     return x + x

model = Adder()
tf.saved_model.save(
  model, '/tmp/adder',
  signatures=model.concat.get_concrete_function(
    tf.TensorSpec(shape=[], dtype=tf.string, name="string_input")))

从输入和输出构造的 tf.keras.Model 实例已经具有签名，因此不需要 @tf.function 装饰器或 signatures 参数。如果两者均未指定，则导出模型的正向传递。

x = tf.keras.layers.Input((4,), name="x")
y = tf.keras.layers.Dense(5, name="out")(x)
model = tf.keras.Model(x, y)
tf.saved_model.save(model, '/tmp/saved_model/')

导出的 SavedModel 采用形状为 [None, 4] 的 "x" 并返回形状为 [None, 5] 的 "out"

变量和检查点

必须通过将变量分配给被跟踪对象的属性或直接分配给obj 的属性来跟踪变量。 TensorFlow 对象(例如来自 tf.keras.layers 的层，来自 tf.train 的优化器)自动跟踪它们的变量。这与tf.train.Checkpoint 使用的跟踪方案相同，导出的Checkpoint 对象可以通过将tf.train.Checkpoint.restore 指向SavedModel 的"variables/" 子目录来恢复为训练检查点。

tf.function 不使用函数体外部的hard-code 设备注释，而不是使用调用上下文的设备。这意味着例如导出在 GPU 上运行的模型并在 CPU 上提供服务通常可以工作，但有一些异常：

• 函数体内的tf.device注解在导出的模型中将是hard-coded；不鼓励使用这种类型的注释。
• Device-specific 操作，例如名称中带有 "cuDNN" 或带有 device-specific 布局可能会导致问题。
• 对于 ConcreteFunctions ，主动分布策略将导致设备放置在函数中为 hard-coded。

导出的 SavedModelstf.saved_model.save 剥离 default-valued 属性自动，当 SavedModel 的消费者运行的 TensorFlow 版本比生产者更旧时，它会消除一个不兼容的来源。然而，还有其他无法自动处理的不兼容来源，例如当导出的模型包含消费者没有定义的操作时。

eager模式兼容性

图形构建时没有得到很好的支持。从 TensorFlow 1.x 开始，tf.compat.v1.enable_eager_execution() 应该首先运行。从 TensorFlow 1.x 构建图形时，在循环中调用 tf.saved_model.save 将在每次迭代时向默认图形添加新的保存操作。

不能从函数体内调用。