保存模型到单个.pb文件中
前面两篇介绍了使用Saver 和SavedModel保存模型:
Tensorflow 模型保存与恢复(1)使用tf.train.Saver()
Tensorflow 模型保存与恢复(2)使用SavedModel
1.保存模型
上面两种方法保存的模型的变量都是单独保存在一个文件中的,模型的图元数据则是保存在另一个文件中。有时候希望将模型的元数据和变量值保存到同一个文件中,即最后的模型只有一个.pb文件。本篇介绍如何将TensorFlow的模型保存到单个文件中,以及读取恢复模型。
基本思路主要是利用GraphDef
对象,使用convert_variables_to_constants
方法将变量转化为常量,然后序列化保存到磁盘。基本步骤如下:
首先获取GraphDef
对象:
input_graph_def = graph.as_graph_def()
设定需要导出的节点,可以将模型的输入输出节点导出,在恢复模型进行预测的时候可以据此获取输入输出的tensor:
output_node_name = 'x_input,conv'
使用convert_variables_to_constants
方法将GraphDef
对象中的变量转化为常量,并返回一个新的GraphDef
:
output_graph_def = tf.graph_util.convert_variables_to_c