两种Tensorflow模型保存的方法

提示：扫一扫查出行
【扫一扫了解最新限行尾号】
复制提示

在Ｔｅｎｓｏｒｆｌｏｗ中，有两种保存模型的方法：一种是Ｃｈｅｃｋｐｏｉｎｔ，另一种是Ｐｒｏｔｏｂｕｆ，也就是ＰＢ格式；

一．Ｃｈｅｃｋｐｏｉｎｔ方法：

１．保存时使用方法：

tf.train.Saver()

生成四个文件：

checkpoint 检查点文件

model.ckpt.data-xxx参数值

model.ckpt.index各个参数

model.ckpt.meta图的结构

２．恢复时使用方法：

saver.restore()：模型文件依赖Ｔｅｎｓｏｒｆｌｏｗ，只能在其框架下使用，恢复模型之前需要定义下网络结构

saver=tf.train.import_meta_graph('./ckpt/mode..ckpt.meta') :直接加载网络结构，不需要重新定义网络

二．ＰＢ方法：

１．保存模型为ＰＢ文件（谷歌推荐），具有语言独立性，可独立运行，序列化的格式，任何语言可解析它，允许其他语言和框架读取，训练和迁移；模型变量是固定的，模型大小会大大减少，适合在手机端运行；

２．实现创建模型与使用模型的解耦，使得前向推导Ｉｎｆｅｒｅｎｃｅ代码统一；

３．ＰＢ文件表示ＭｅｔａＧｒａｐｈ的protocol buffer格式的文件；

４．ＧｒａｐｈＤｅｆ不保存任何Ｖａｒｉａｂｌｅ信息，不能从graph_def 来构建图并恢复训练．

一般情况下，ＰＢ可直接生成；

当然也可以从checkpoint文件中生成，代码如下：

output_graph = os.path.join('./checkpoint/','frozen_graph.pb')
input_checkpoint = os.path.join('./checkpoint/','model.ckpt-xxxxx') ＃［xxxxxx为训练生成的ｓｔｅｐ号］
saver = tf.train.import_meta_graph(input_checkpoint+'.meta',clear_devices=True)
graph = tf.get_default_graph()
input_graph_def = graph.as_graph_def for op in graph.get_operations():
print("checkpoint2pb",op.name,op.values()) variable_names = [v.name for v in tf.trainable_variables()］
pirnt("trainalbe_variables:",variable_names) output_node_name=['fc2/add'] #fc2/add 上面的列表里需要存在该操作 with tf.Session() as sess:
saver.restore(sess,input_checkpoint) output_graph_def = graph_util.convert_variables_to_constants(sess=sess,
input_graph_def = input_graph_def,
output_node_names = output_node_name) with tf.gfile.GFile(output_graph,"wb") as f:
f.write(output_graph_def.SerializeToString())