记录:tf.saved_model 模块的简单使用(TensorFlow 模型存储与恢复)

时间:2024-01-22 07:49:39

虽然说 TensorFlow 2.0 即将问世,但是有一些模块的内容却是不大变化的。其中就有 tf.saved_model 模块,主要用于模型的存储和恢复。为了防止学习记录文件丢失或者蠢笨的脑子直接遗忘掉这部分内容,在此做点简单的记录,以便将来查阅。

最近为了一个课程作业,不得已涉及到关于图像超分辨率恢复的内容,不得不准备随时存储训练的模型,只好再回过头来瞄一眼 TensorFlow 文档,真是太痛苦了。

tf.saved_model 模块下面有很多文件和函数,精力有限,只好选择于自己有用的东西来看,可能并不全面,望日后补上。

其中最重要的就是该模块下的一个类:tf.saved_model.builder.SavedModelBuilder

tf.saved_model.builder.SavedModelBuilder:

# 构造函数
.__init__(export_dir)
"""
作用:
  创建一个保存模型的实例对象
参数:
    export_dir: 模型导出路径,由于 TensorFlow 会在你指定的路径上创建文件夹和文件,所以指定的路径最后不需要带 /,
   例如:export_dir='/home/***/saved_model' 即可,最后不需要加上 /
"""

# 方法
# 1
.add_meta_graph_and_variables(sess, tags, signature_def_map=None, assets_collection=None,
                              clear_devices=False, main_op=None, strip_default_attrs=False, saver=None)
"""
作用: 
  保存会话对象中的 graph 和所有变量,具体描述可参见文档
参数:
  sess: TensorFlow 会话对象,用于保存元图和变量
  tags: 用于保存元图的标记集(如果存在多个图对象,需要设置保证每个图标签不一样),是一个列表
  signature_def_map: 一个字典,保存模型时传入的参数,key 可以是字符串,也可以是 tf.saved_model.signature_constants 文件下预定义的变量,
                      值为 signatureDef protobuf(protobuf 是一种结构化的数据存储格式)
  assets_collection: 略
  clear_devices: 如果需要清除默认图上的设备信息,则设置为 true
  main_op: 这个参数包括后面一系列与其相关的东西没有弄明白
  strip_default_attrs: 如果设置为 True,将从 NodeDefs 中删除默认值属性
  saver: tf.train.Saver 的一个实例,用于导出元图并保存变量
"""

# 2
.add_meta_graph()
"""
作用:
  其除了没有 sess 参数以外,其他参数和 .add_meta_graph_and_variables() 一模一样
  调用此方法之前必须先调用 .add_meta_graph_and_variables() 方法
"""

# 3
.save(as_text=False)
"""
作用: 
  将内建的 savedModel protobuf 写入磁盘
"""

 

 

除了这个最重要的类以外,tf.saved_model 模块还提供了一些方便构建 builder 和加载模型的函数方法。

# 1
tf.saved_model.utils.build_tensor_info(tensor)
"""
作用:
    构建 TensorInfo protobuf,根据输入的 tensor 构建相应的 protobuf,返回的 TensorInfo 中包含输入 tensor 的 name,shape,dtype 信息
参数:
    tensor: Tensor 或 SparseTensor
"""

# 2
tf.saved_model.signature_def_utils.build_signature_def(inputs=None, outputs=None, method_name=None)
"""
作用:
    构建 SignatureDef protobuf,并返回 SignatureDef protobuf
参数:
    inputs: 一个字典,键为字符串类型,值为关于 tensor 的信息,也就是上述的 .build_tensor_info() 函数返回的 TensorInfo protobuf
    outputs: 一个字典,同上
    method_name: SignatureDef 名称
"""

# 3
tf.saved_model.utils.get_tensor_from_tensor_info(tensor_info, graph=None, import_scope=None)
"""
作用:
    根据一个 TensorInfo protobuf 解析出一个 tensor
参数:
    tensor_info: 一个 TensorInfo protobuf
    graph: tensor 所存在的 graph,参数为 None 时,使用默认图
    import_scope: 给 tensor 的 name 加上前缀
"""

# 4
tf.saved_model.loader.load(sess, tags, export_dir, import_scope=None, **saver_kwargs)
"""
作用:
    加载已存储的模型
参数:
    sess: 用于恢复模型的 tf.Session() 对象
    tags: 用于标识 MetaGraphDef 的标记,应该和存储模型时使用的此参数完全一致
    export_dir: 模型存储路径
    import_scope: 加前缀
"""

除了这些以外,还有一些 TensorFlow 为了方便而预定义的一些变量,这些变量完全可以使用自定义字符串代替,不再赘述。详情:https://tensorflow.google.cn/api_docs/python/tf/saved_model

如果只看这些内容的话,确实会使人产生巨大的疑惑,下面是具体实践的例子:

import tensorflow as tf
from tensorflow import saved_model as sm


# 首先定义一个极其简单的计算图
X = tf.placeholder(tf.float32, shape=(3, ))
scale = tf.Variable([10, 11, 12], dtype=tf.float32)
y = tf.multiply(X, scale)

# 在会话中运行
with tf.Session() as sess:
    sess.run(tf.initializers.global_variables())
    value = sess.run(y, feed_dict={X: [1., 2., 3.]})
    print(value)
    
    # 准备存储模型
    path = '/home/×××/tf_model/model_1'
    builder = sm.builder.SavedModelBuilder(path)
    
    # 构建需要在新会话中恢复的变量的 TensorInfo protobuf
    X_TensorInfo = sm.utils.build_tensor_info(X)
    scale_TensorInfo = sm.utils.build_tensor_info(scale)
    y_TensorInfo = sm.utils.build_tensor_info(y)

    # 构建 SignatureDef protobuf
    SignatureDef = sm.signature_def_utils.build_signature_def(
                                inputs={'input_1': X_TensorInfo, 'input_2': scale_TensorInfo},
                                outputs={'output': y_TensorInfo},
                                method_name='what'
    )

    # 将 graph 和变量等信息写入 MetaGraphDef protobuf
    # 这里的 tags 里面的参数和 signature_def_map 字典里面的键都可以是自定义字符串,TensorFlow 为了方便使用,不在新地方将自定义的字符串忘记,可以使用预定义的这些值
    builder.add_meta_graph_and_variables(sess, tags=[sm.tag_constants.TRAINING], 
                                         signature_def_map={sm.signature_constants.CLASSIFY_INPUTS: SignatureDef}
  ) 

 # 将 MetaGraphDef 写入磁盘
    builder.save()

 

 

 

这样我们就把模型整体存储到了磁盘中,而且我们将三个变量 X, scale, y 全部序列化后存储到了其中,所以恢复模型时便可以将他们完全解析出来:

import tensorflow as tf
from tensorflow import saved_model as sm


# 需要建立一个会话对象,将模型恢复到其中
with tf.Session() as sess:
    path = '/home/×××/tf_model/model_1'
    MetaGraphDef = sm.loader.load(sess, tags=[sm.tag_constants.TRAINING], export_dir=path)

    # 解析得到 SignatureDef protobuf
    SignatureDef_d = MetaGraphDef.signature_def
    SignatureDef = SignatureDef_d[sm.signature_constants.CLASSIFY_INPUTS]

    # 解析得到 3 个变量对应的 TensorInfo protobuf
    X_TensorInfo = SignatureDef.inputs['input_1']
    scale_TensorInfo = SignatureDef.inputs['input_2']
    y_TensorInfo = SignatureDef.outputs['output']

    # 解析得到具体 Tensor
    # .get_tensor_from_tensor_info() 函数中可以不传入 graph 参数,TensorFlow 自动使用默认图
    X = sm.utils.get_tensor_from_tensor_info(X_TensorInfo, sess.graph)
    scale = sm.utils.get_tensor_from_tensor_info(scale_TensorInfo, sess.graph)
    y = sm.utils.get_tensor_from_tensor_info(y_TensorInfo, sess.graph)

    print(sess.run(scale))
    print(sess.run(y, feed_dict={X: [3., 2., 1.]}))

# 输出
[10. 11. 12.]
[30. 22. 12.]

可以看出模型整体和变量个体都被完整地保存了下来。其中涉及的关于 protobuf 的知识,需要补习,在 TensorFlow 中好多地方都用到了相关的知识。上述恢复模型的代码中对具体的 TensorInfo protobuf 解析时,还可以使用另一种方式得到相应的 Tensor:

# 已知 X_TensorInfo, scale_TensorInfo, y_TensorInfo
X = sess.graph.get_tensor_by_name(X_TensorInfo.name)
scale = sess.grpah.get_tensor_by_name(scale_TensorInfo.name)
y = sess.graph.get_tensor_by_name(y_TensorInfo.name)

# 因为 TensorFlow 构建 TensorInfo protobuf 时,使用了 Tensor 的 name 信息,所以可以直接读出来使用