TensorFlow之SavedModel模型保存和加载

一、参考资料

TensorFlow 2.0 - tf.saved_model.save 模型导出

TensorFlow 模型导出

磁盘上的 SavedModel 格式

Tensorflow SavedModel模型的保存与加载

二、SavedModel 格式

1. 引言

为了将训练好的机器学习模型部署到各个目标平台（如服务器、移动端、嵌入式设备和浏览器等），第一步往往是将训练好的整个模型完整导出（序列化）为一系列标准格式的文件。在此基础上，我们才可以在不同的平台上使用相对应的部署工具来部署模型文件。

TensorFlow 提供了统一模型导出格式 SavedModel，使得训练好的模型可以在多种不同平台上部署。SavedModel与编程语言无关，比如可以使用python语言训练模型，然后在Java中非常方便的加载模型。另外，如果使用Tensorflow Serving server来部署模型，必须选择SavedModel格式。

2. SavedModel 简介

SavedModel为 Frozen GraphDef 和 CheckPoint 的结合体，一个 SavedModel 包含了一个完整的 TensorFlow 程序信息，包括模型权重参数和网络结构（计算图）。当模型导出为SavedModel格式时，无需建立模型的源代码，即可再次运行模型，使得SavedModel格式尤其适合分享和部署，所以很容易在 TensorFlow Lite（移动端部署模型），TensorFlow.js，TensorFlow Serving（服务器端部署模型），TensorFlow Hub 上部署。目前，建议TensorFlow和Keras模型都导出成SavedModel格式，这样就可以直接使用通用的TensorFlow Serving服务，模型导出即可上线不需要改任何代码，并且可以简单被python，java等调用。

SavedModel保存的是整个训练图，并且参数没有冻结，参数不冻结无法进行转TensorRT等极致优化。

通常，SavedModel 模型文件由以下几个部分组成：

.
├── assets  # 所需的外部文件，例如初始化的词汇表文件
├── keras_metadata.pb
├── saved_model.pb  # 保存MetaGraph的网络结构
└── variables  # 权重参数，包含所有模型的变量(tf.Variable objects)
    ├── variables.data-00000-of-00001
    └── variables.index

3. SavedModel CLI (Command-Line Interface)

SavedModel CLI

查看SavedModel模型

# 终端执行
saved_model_cli show --dir=./saved_model --all

输出结果

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

signature_def['__saved_model_init_op']:
  The given SavedModel SignatureDef contains the following input(s):
  The given SavedModel SignatureDef contains the following output(s):
    outputs['__saved_model_init_op'] tensor_info:
        dtype: DT_INVALID
        shape: unknown_rank
        name: NoOp
  Method name is: 

signature_def['serving_default']:
  The given SavedModel SignatureDef contains the following input(s):
    inputs['input_1'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 224, 224, 3)
        name: serving_default_input_1:0
  The given SavedModel SignatureDef contains the following output(s):
    outputs['convDepthF'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 56, 56, 1)
        name: StatefulPartitionedCall:0
  Method name is: tensorflow/serving/predict

Defined Functions:
  Function Name: '__call__'
    Option #1
      Callable with:
        Argument #1
          inputs: TensorSpec(shape=(None, 224, 224, 3), dtype=tf.float32, name='inputs')
        Argument #2
          DType: bool
          Value: True
        Argument #3
          DType: NoneType
          Value: None
    Option #2
      Callable with:
        Argument #1
          input_1: TensorSpec(shape=(None, 224, 224, 3), dtype=tf.float32, name='input_1')
        Argument #2
          DType: bool
          Value: True
        Argument #3
          DType: NoneType
          Value: None
    Option #3
      Callable with:
        Argument #1
          input_1: TensorSpec(shape=(None, 224, 224, 3), dtype=tf.float32, name='input_1')
        Argument #2
          DType: bool
          Value: False
        Argument #3
          DType: NoneType
          Value: None
    Option #4
      Callable with:
        Argument #1
          inputs: TensorSpec(shape=(None, 224, 224, 3), dtype=tf.float32, name='inputs')
        Argument #2
          DType: bool
          Value: False
        Argument #3
          DType: NoneType
          Value: None

  Function Name: '_default_save_signature'
    Option #1
      Callable with:
        Argument #1
          input_1: TensorSpec(shape=(None, 224, 224, 3), dtype=tf.float32, name='input_1')

  Function Name: 'call_and_return_all_conditional_losses'
    Option #1
      Callable with:
        Argument #1
          input_1: TensorSpec(shape=(None, 224, 224, 3), dtype=tf.float32, name='input_1')
        Argument #2
          DType: bool
          Value: False
        Argument #3
          DType: NoneType
          Value: None
    Option #2
      Callable with:
        Argument #1
          input_1: TensorSpec(shape=(None, 224, 224, 3), dtype=tf.float32, name='input_1')
        Argument #2
          DType: bool
          Value: True
        Argument #3
          DType: NoneType
          Value: None
    Option #3
      Callable with:
        Argument #1
          inputs: TensorSpec(shape=(None, 224, 224, 3), dtype=tf.float32, name='inputs')
        Argument #2
          DType: bool
          Value: False
        Argument #3
          DType: NoneType
          Value: None
    Option #4
      Callable with:
        Argument #1
          inputs: TensorSpec(shape=(None, 224, 224, 3), dtype=tf.float32, name='inputs')
        Argument #2
          DType: bool
          Value: True
        Argument #3
          DType: NoneType
          Value: None

由以上信息可知：

• MetaGraphDef with tag-set: ‘serve’：tag_constants.SERVING
• 输出节点名称：StatefulPartitionedCall
• signature_def['serving_default']：signature签名为 serving_default；

4. SavedModel格式保存和加载

4.1 TensorFlow 1.x

对于 TensorFlow 1.x 的模型导出，simple_save 是最简单的方式，只需要补充简单的必要参数，有很多参数被省略。

"""tf1.x"""
x = tf.placeholder(tf.float32, [None, 784], name="myInput")
y = tf.nn.softmax(tf.matmul(x, W) + b, name="myOutput")
tf.saved_model.simple_save(
                sess,
                export_dir,
                inputs={"myInput": x},
                outputs={"myOutput": y})

4.2 TensorFlow 2.x

保存/导出模型：

# 方式一
tf.saved_model.save(mymodel, "./saved_model")

# 方式二
model.save("./saved_model", save_format="tf")  # 保存SavedModel格式
model.save("SPEED_model.h5", save_format="h5")  # 保存h5格式

加载/导入：

# 加载模型
mymodel = tf.saved_model.load("./saved_model")

5. SavedModel模型推理

import tensorflow as tf
import os


os.environ['CUDA_VISIBLE_DEVICES'] = '-1'  # 启用cpu推理
frozen_out_path = './models'
frozen_graph_filename = "SPEED_model.pb"
saved_model_path = './saved_model/'

model = tf.keras.models.load_model(saved_model_path)
model.summary()
images = tf.random.uniform((1, 224, 224, 3))
result = model.predict(images)
print(result.shape)  # (1, 56, 56, 1)

6. SavedModel转pb

SavedModel格式转成Frozen GraphDef格式。

示例1

import tensorflow as tf
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
import os


os.environ['CUDA_VISIBLE_DEVICES'] = '-1'  # 启用cpu推理
frozen_out_path = './models'
frozen_graph_filename = "SPEED_model.pb"
saved_model_path = './saved_model/'

# 定义输入格式
img = tf.random.uniform((1, 224, 224, 3))

# 加载模型
network = tf.keras.models.load_model(saved_model_path)

# Convert Keras model to ConcreteFunction
full_model = tf.function(lambda x: network(x))

full_model = full_model.get_concrete_function(tf.TensorSpec(img.shape, img.dtype))  # (1, 224, 224, 3) <dtype: 'uint8'>

# Get frozen ConcreteFunction
frozen_func = convert_variables_to_constants_v2(full_model)
frozen_func.graph.as_graph_def()

layers = [op.name for op in frozen_func.graph.get_operations()]
print("-" * 50)
print("Frozen model layers: ")
for layer in layers:
    print(layer)

print("-" * 50)
print("Frozen model inputs: ")  # [<tf.Tensor 'x:0' shape=(1, 224, 224, 3) dtype=float32>]
print(frozen_func.inputs)
print("Frozen model outputs: ")  # [<tf.Tensor 'Identity:0' shape=(1, 56, 56, 1) dtype=float32>]
print(frozen_func.outputs)

# Save frozen graph from frozen ConcreteFunction to hard drive
tf.io.write_graph(graph_or_graph_def=frozen_func.graph,
                  logdir=frozen_out_path,
                  name=frozen_graph_filename,
                  as_text=False)

示例2

from tensorflow.python.tools import freeze_graph
from tensorflow.python.saved_model import tag_constants


input_saved_model_dir = "./saved_model/"
output_node_names = "StatefulPartitionedCall"
input_binary = False
input_saver_def_path = False
restore_op_name = None
filename_tensor_name = None
clear_devices = False
input_meta_graph = False
checkpoint_path = None
input_graph_filename = None
saved_model_tags = tag_constants.SERVING
output_graph_filename = 'frozen_graph.pb'

freeze_graph.freeze_graph(input_graph_filename,
  input_saver_def_path,
  input_binary,
  checkpoint_path,
  output_node_names,
  restore_op_name,
  filename_tensor_name,
  output_graph_filename,
  clear_devices,
  "", "", "",
  input_meta_graph,
  input_saved_model_dir,
  saved_model_tags)

7. 签名Signatures

A Tour of SavedModel Signatures
tf.keras.Model使用saved_model,自定义输入输出signature

不同的模型导出时只要指定输入和输出的signature即可。

保存模型时指定签名

tf.keras.Model 会自动指定服务上线签名，但是，对于自定义模块，我们必须明确声明服务上线签名。

当指定单个签名时，签名键为 'serving_default'，并将保存为常量 tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY。

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

assert len(imported.signatures) == 0

如果导出时指定签名，需要使用 signatures 关键字参数指定 ConcreteFunction。

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)

# 输出
Tracing with Tensor("x:0", dtype=float32)
Tracing with Tensor("x:0", dtype=float32)
INFO:tensorflow:Assets written to: /tmp/tmp0ya0f8kf/module_with_signature/assets

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

# 输出
['serving_default']

要导出多个签名，请将签名键的字典传递给 ConcreteFunction。每个签名键对应一个 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)

# 输出
Tracing with Tensor("x:0", shape=(None,), dtype=float32)
Tracing with Tensor("x:0", shape=(None,), dtype=float32)
INFO:tensorflow:Assets written to: /tmp/tmp0ya0f8kf/module_with_multiple_signatures/assets

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

# 输出
['serving_default', 'array_input']

默认情况下，输出张量名称非常通用，如 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([], 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})

# 输出
INFO:tensorflow:Assets written to: /tmp/tmp0ya0f8kf/module_with_output_name/assets

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

# 输出
{'custom_output_name': TensorSpec(shape=(), dtype=tf.float32, name='custom_output_name')}

8. 示例代码

github样例

8.1保存SavedModel模型

import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()


with tf.Session() as sess:
    # 定义计算图结构 x*y+b
    x = tf.placeholder(tf.int32, name='x_input')
    y = tf.placeholder(tf.int32, name='y_input')
    b = tf.Variable(1, name='b')
    xy = tf.multiply(x, y)
    output = tf.add(xy, b, name='output')  # 指定输出节点的名称name

    sess.run(tf.global_variables_initializer())  # 初始化过程

    pb_file_path = "./models/"
    builder = tf.saved_model.builder.SavedModelBuilder(pb_file_path+'saved_model')
    # 构造模型保存的内容，指定要保存的 session 和 tags,
    builder.add_meta_graph_and_variables(sess, tags=['cpu_server_1'])
    builder.save()

8.2加载SavedModel模型

import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()


with tf.Session() as sess:
    # 定义计算图结构 x*y+b
    x = tf.placeholder(tf.int32, name='x_input')
    y = tf.placeholder(tf.int32, name='y_input')
    b = tf.Variable(1, name='b')
    xy = tf.multiply(x, y)
    output = tf.add(xy, b, name='output')  # 指定输出节点的名称name

    sess.run(tf.global_variables_initializer())  # 初始化过程

    pb_file_path = "./models/"

    with tf.Session(graph=tf.Graph()) as sess:
    	# tag标签与保存时的一致
        tf.saved_model.loader.load(sess, ['cpu_server_1'], pb_file_path+'saved_model')
        sess.run(tf.global_variables_initializer())

        x_input = sess.graph.get_tensor_by_name('x_input:0')
        y_input = sess.graph.get_tensor_by_name('y_input:0')

        op = sess.graph.get_tensor_by_name('output:0')

        ret = sess.run(op, feed_dict={x_input: 10, y_input: 3})
        print(ret)

三、自定义模型的保存与加载

1. 保存自定义模型

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 代码。

简单地说，tf.function 的工作原理是，通过跟踪 Python 代码来生成 ConcreteFunction（一个可调用的 tf.Graph 包装器）。当保存 tf.function 时，实际上保存的是 tf.function 的 ConcreteFunction 缓存。

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)

# 输出
Tracing with Tensor("x:0", shape=(), dtype=float32)
Saving model...
Tracing with Tensor("x:0", shape=(), dtype=float32)
INFO:tensorflow:Assets written to: /tmp/tmp0ya0f8kf/module_no_signatures/assets

2. 加载和使用自定义模型

在 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'),), {})]. 

可以使用 tf.saved_model.load 将 SavedModel 加载回 Python。

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

# 输出
["serving_default"]

导入的签名总是会返回字典。

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

# 输出
{'predictions': TensorSpec(shape=(None, 1000), dtype=tf.float32, name='predictions')}

四、示例代码

Keras API模型导出

# 模型导出
model.save('catdog.h5')

# 模型载入
model = tf.keras.models.load_model('catdog.h5')

"""
由于 SavedModel 基于计算图，所以对于使用继承 tf.keras.Model 类建立的 Keras 模型，其需要导出到 SavedModel 格式的方法（比如 call ）都需要使用 @tf.function 修饰。
"""
class MLPmodel(tf.keras.Model):
    def __init__(self):
        super().__init__()
        # 除第一维以外的维度展平
        self.flatten = tf.keras.layers.Flatten()
        self.dense1 = tf.keras.layers.Dense(units=100, activation='relu')
        self.dense2 = tf.keras.layers.Dense(units=10)
    
    @tf.function # 计算图模式，导出模型，必须写
    def call(self, input):
        x = self.flatten(input)
        x = self.dense1(x)
        x = self.dense2(x)
        output = tf.nn.softmax(x)
        return output

# 导出模型， 模型目录
tf.saved_model.save(mymodel, "./my_model_path")
# 载入模型
mymodel = tf.saved_model.load('./my_model_path')

# 对于使用继承 tf.keras.Model 类建立的 Keras 模型 model ，使用 SavedModel 载入后将无法使用 evaluate，predict 直接进行推理，而需要使用 model.call() 。

res = mymodel.call(data_loader.test_data)
print(res)