Tensorflow模型的格式
tensorflow模型的格式通常支持多种,主要有CheckPoint(*.ckpt)、GraphDef(*.pb)、SavedModel。
1. CheckPoint(*.ckpt)
在训练 TensorFlow 模型时,每迭代若干轮需要保存一次权值到磁盘,称为“checkpoint”,如下图所示:
文件中的meta为网络结构,data为权重。注意加载时的模型文件路径为:model.ckpt。
这种格式文件是由 tf.train.Saver() 对象调用 saver.save() 生成的,只包含若干 Variables 对象序列化后的数据,不包含图结构,所以只给 checkpoint 模型不提供代码是无法重新构建计算图的。载入 checkpoint 时,调用 saver.restore(session, checkpoint_path)。新近版本的tensorflow.keras也支持加载model.load_weights()。
缺点:首先模型文件是依赖 TensorFlow 的,只能在其框架下使用;其次,在恢复模型之前还需要再定义一遍网络结构,然后才能把变量的值恢复到网络中。
2. GraphDef(*.pb)
这种格式文件包含 protobuf 对象序列化后的数据,包含了计算图,可以从中得到所有运算符(operators)的细节,也包含张量(tensors)和 Variables 定义,但不包含 Variable 的值,因此只能从中恢复计算图,但一些训练的权值仍需要从 checkpoint 中恢复。下面代码实现了利用 *.pb 文件构建计算图:
TensorFlow 一些例程中用到 *.pb 文件作为预训练模型,这和上面 GraphDef 格式稍有不同,属于冻结(Frozen)后的 GraphDef 文件,简称 FrozenGraphDef 格式。这种文件格式不包含 Variables 节点。将 GraphDef 中所有 Variable 节点转换为常量(其值从 checkpoint 获取),就变为 FrozenGraphDef 格式。代码可以参考 tensorflow/python/tools/freeze_graph.py
*.pb 为二进制文件,实际上 protobuf 也支持文本格式(*.pbtxt),但包含权值时文本格式会占用大量磁盘空间,一般不用。
3. SavedModel
https://juejin.im/post/5bbfedd65188255c9b13d964
https://zhuanlan.zhihu.com/p/31417693
这是谷歌推荐的模型保存方式,它具有语言独立性,可独立运行,封闭的序列化格式,任何语言都可以解析它,它允许其他语言和深度学习框架读取、继续训练和迁移 TensorFlow 的模型。该格式为 GraphDef 和 CheckPoint 的结合体,另外还有标记模型输入和输出参数的 SignatureDef。从 SavedModel 中可以提取 GraphDef 和 CheckPoint 对象。
SavedModel 目录结构如下:
其中 saved_model.pb(或 saved_model.pbtxt)包含使用 MetaGraphDef protobuf 对象定义的计算图;assets 包含附加文件;variables 目录包含 tf.train.Saver() 对象调用 save() API 生成的文件。
以下代码实现了保存 SavedModel:
#创建signature
def signature_def(self):
inputs = {'char_inputs': tf.saved_model.utils.build_tensor_info(self.char_inputs),
'seg_inputs': tf.saved_model.utils.build_tensor_info(self.seg_inputs),
'dropout': tf.saved_model.utils.build_tensor_info(self.dropout)}
outputs = {'decode_tags': tf.saved_model.utils.build_tensor_info(self.decode_tags)}
return tf.saved_model.signature_def_utils.build_signature_def(inputs=inputs
,outputs=outputs
,method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME)
#保存模型
def save_model(self, sess, signature, save_path):
builder = tf.saved_model.builder.SavedModelBuilder(save_path)
builder.add_meta_graph_and_variables(sess, [tf.saved_model.tag_constants.SERVING], {'predict': signature}, clear_devices=True)
builder.save()
载入 SavedModel:
model = tf.saved_model.loader.load(sess, [tf.saved_model.tag_constants.SERVING], checkpoint_path)
signature = model.signature_def
char_inputs_ = signature['predict'].inputs['char_inputs'].name
seg_inputs_ = signature['predict'].inputs['seg_inputs'].name
dropout_ = signature['predict'].inputs['dropout'].name
decode_tags_ = signature['predict'].outputs['decode_tags'].name
# get tensor
char_inputs = sess.graph.get_tensor_by_name(char_inputs_)
seg_inputs = sess.graph.get_tensor_by_name(seg_inputs_)
dropout = sess.graph.get_tensor_by_name(dropout_)
decode_tags = sess.graph.get_tensor_by_name(decode_tags_)
decode_tags_ = sess.run([decode_tags], feed_dict={char_inputs: inputs[1], seg_inputs:inputs[2], dropout:1.0 })
更多细节可以参考 tensorflow/python/saved_model/README.md。
4. 各模式之间的转换
https://zhuanlan.zhihu.com/p/47649285
5. 小结
本文总结了 TensorFlow 常见模型格式和载入、保存方法。部署在线服务(Serving)时官方推荐使用 SavedModel 格式,而部署到手机等移动端的模型一般使用 FrozenGraphDef 格式(最近推出的 TensorFlow Lite 也有专门的轻量级模型格式 *.lite,和 FrozenGraphDef 十分类似)。这些格式之间关系密切,可以使用 TensorFlow 提供的 API 来互相转换。