一、参考资料
TensorFlow 2.0 - tf.saved_model.save 模型导出
TensorFlow 模型导出
磁盘上的 SavedModel 格式
Tensorflow SavedModel模型的保存与加载
二、SavedModel格式相关介绍
1. 问题引入
为了将训练好的机器学习模型部署到各个目标平台（如服务器、移动端、嵌入式设备和浏览器等），我们的第一步往往是将训练好的整个模型完整导出（序列化）为一系列标准格式的文件。在此基础上，我们才可以在不同的平台上使用相对应的部署工具来部署模型文件。
TensorFlow 提供了统一模型导出格式 SavedModel，使得训练好的模型可以在多种不同平台上部署。SavedModel与编程语言无关，比如可以使用python语言训练模型，然后在Java中非常方便的加载模型。另外，如果使用Tensorflow Serving server来部署模型，必须选择SavedModel格式。
2. SavedModel 简介
目前建议TensorFlow和Keras模型都导出成SavedModel格式，这样就可以直接使用通用的TensorFlow Serving服务，模型导出即可上线不需要改任何代码，并且可以简单被python，java等调用。
SavedModel为 Frozen GraphDef 和 CheckPoint 的结合体，一个 SavedModel 包含了一个完整的 TensorFlow 程序信息，包括模型权重参数和网络结构（计算图）。当模型导出为SavedModel格式时，无需建立模型的源代码，即可再次运行模型，使得SavedModel格式尤其适合分享和部署，所以很容易在 TensorFlow Lite（移动端部署模型），TensorFlow.js，TensorFlow Serving（服务器端部署模型），TensorFlow Hub 上部署。
SavedModel保存的是一整个训练图，并且参数没有冻结，参数不冻结无法进行转TensorRT等极致优化。
通常SavedModel由以下几个部分组成：
.
├── assets # 所需的外部文件，例如初始化的词汇表文件
├── keras_metadata.pb
├── saved_model.pb # 保存MetaGraph的网络结构
└── variables # 权重参数，包含所有模型的变量(tf.Variable objects)参数
├── variables.data-00000-of-00001
└── variables.index
2 directories, 4 files
3. SavedModel CLI (Command-Line Interface)
查看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格式保存和加载
TensorFlow 1.x
"""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})
对于普通的tf 模型导出，simple_save 是最简单的方式，只需要补充简单的必要参数，有很多参数被省略。
TensorFlow 2.x
4.1 保存/导出
# 保存模型
# 方式一
tf.saved_model.save(mymodel, "./saved_model")
# 方式二
model.save("./saved_model", save_format="tf") # 保存SavedModel格式
model.save("SPEED_model.h5", save_format="h5") # 保存h5格式
4.2 加载/导入
# 加载模型
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. 示例代码
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)
2万+
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
