tensorflow保存模型的几种方法

tensorflow保存模型有多种方法

第一种：saver.save(sess, "./hello_model") # 生成ckpt模型文件， hello_model.data-00000-of-00001  hello_model.index  hello_model.meta

第二种：tf.train.write_graph(sess.graph_def, ./,  'hello.pb') # 生成hello.pb， 再通过freeze_graph把hello.pb与ckpt固化成新的pb文件

第三种：用tf.graph_util.convert_variables_to_constants把变量转成常量之后写入PB文件中

第四种：使用tf.saved_model.builder.SavedModelBuilder

具体看代码，及示例

保存模型的文件 saver_hello.py

点击(此处)折叠或打开

1. import tensorflow as tf
   
2. import sys
   
3. import os
   
4. 
   
5. # 把变量转成常量之后写入PB文件中
   
6. def SaveFrozenPb(nodeNameList, pbFile):
   
7.     gd = tf.graph_util.convert_variables_to_constants(sess, tf.get_default_graph().as_graph_def(), nodeNameList)
   
8.     with tf.gfile.GFile(pbFile, 'wb') as f:
   
9.         f.write(gd.SerializeToString())
   
10. 
    
11. # 通过freeze_graph把tf.train.write_graph()生成的pb文件与tf.train.saver()生成的chkp文件固化之后重新生成一个pb文件
    
12. # freeze_graph --input_graph=./hello.pb --input_checkpoint=./hello_model --output_node_names=hello,y --input_node_names=x --output_graph=./hello_frozen.pb
    
13. # 如果不调用freeze_graph, 直接使用会报错‘google.protobuf.message.DecodeError: Error parsing message’
    
14. def SavePbForFreezeGraph(pbDir, pbName):
    
15.     tf.train.write_graph(sess.graph_def, pbDir, pbName)
    
16. 
    
17. def SaveBuilderPb(pbDir):
    
18.     if not os.path.exists(pbDir):
    
19.         os.makedirs(pbDir)
    
20.     builder = tf.saved_model.builder.SavedModelBuilder(pbDir)
    
21.     builder.add_meta_graph_and_variables(sess, [tf.saved_model.tag_constants.TRAINING], signature_def_map=None, assets_collection=None)
    
22.     builder.save()
    
23. 
    
24. if __name__ == '__main__':
    
25.     hello = tf.Variable(tf.constant('Hello World', name = "hello")) # 要save成功，需要tf.Variable, 否则会报错'ValueError: No variables to save'
    
26.     x = tf.placeholder(tf.float32, name="x")
    
27.     y = tf.multiply(x, 2, name="y")
    
28. 
    
29.     init = tf.global_variables_initializer()
    
30.     sess = tf.Session()
    
31.     sess.run(init)
    
32. 
    
33.     saver = tf.train.Saver()
    
34. 
    
35.     typeStr = sys.argv[1]
    
36.     if typeStr == 'ckpt' or typeStr == 'pbNotFrozen':
    
37.         saver.save(sess, "./hello_model", write_meta_graph=True) # hello_model.data-00000-of-00001 hello_model.index hello_model.meta
    
38.     if typeStr == 'pbNotFrozen':
    
39.         SavePbForFreezeGraph('./', 'hello.pb') # 需要经由freeze_graph工具处理
    
40.     elif typeStr == 'pbFrozen':
    
41.         SaveFrozenPb(['x', 'y', 'hello'], './hello_frozen.pb') # 无需再经由freeze_graph工具处理
    
42.     elif typeStr == 'builderPb':
    
43.         SaveBuilderPb('./save/')

加载模型文件restore_hello.py

点击(此处)折叠或打开

1. import tensorflow as tf
   
2. import sys
   
3. 
   
4. def RestoreMeta(sess, name):
   
5.     #ckpt = tf.train.get_checkpoint_state('./')
   
6.     #restore = tf.train.import_meta_graph(ckpt.model_checkpoint_path +'.meta')
   
7.     #restore.restore(sess, ckpt.model_checkpoint_path)
   
8.     restore = tf.train.import_meta_graph(name)
   
9.     restore.restore(sess, "hello_model")
   
10. 
    
11. def RestorePb(sess, name):
    
12.     # 二进制读取模型文件
    
13.     with tf.gfile.FastGFile(name, 'rb') as f:
    
14.        graph_def = tf.GraphDef()
    
15.        graph_def.ParseFromString(f.read())
    
16.        sess.graph.as_default()
    
17.        tf.import_graph_def(graph_def, name='') # 导入计算图
    
18. 
    
19. def RestoreBuilderPb(sess, pbDir):
    
20.     tf.saved_model.loader.load(sess, [tf.saved_model.tag_constants.TRAINING], pbDir)
    
21. 
    
22. if __name__ == '__main__':
    
23.     sess = tf.Session()
    
24.     typeStr = sys.argv[1]
    
25.     if typeStr == 'ckpt':
    
26.         RestoreMeta(sess, 'hello_model.meta')
    
27.     elif typeStr == 'pbFrozen':
    
28.         RestorePb(sess, './hello_frozen.pb')
    
29.     elif typeStr == 'builderPb':
    
30.         RestoreBuilderPb(sess, './save/')
    
31. 
    
32.     x = tf.get_default_graph().get_tensor_by_name("x:0")
    
33.     y = tf.get_default_graph().get_tensor_by_name("y:0")
    
34.     hello = tf.get_default_graph().get_tensor_by_name("hello:0")
    
35. 
    
36.     print(sess.run(y, feed_dict={x:5})) # 10.0
    
37.     print(sess.run(hello)) # b'Hello World'

第一种：ckpt

保存模型

python3 ./saver_hello.py ckpt

生成checkpoint hello_model.data-00000-of-00001  hello_model.index  hello_model.meta
加载模型

python3 ./restore_hello.py ckpt

运行结果

10.0
b'Hello World'

第二种：ckpt+pb+固化 

python3 ./saver_hello.py pbNotFrozen
生成checkpoint   hello_model.data-00000-of-00001  hello_model.index  hello_model.meta  hello.pb

固化

freeze_graph --input_graph=./hello.pb --input_checkpoint=./hello_model --output_node_names=hello,y --input_node_names=x --output_graph=./hello_frozen.pb

加载

python3 ./restore_hello.py pbFrozen

第三种：固化的pb

保存

python3 ./saver_hello.py pbFrozen

加载

python3 ./restore_hello.py pbFrozen

第四种：

python3 ./saver_hello.py builderPb
python3 ./restore_hello.py builderPb