谷歌推荐的保存模型的方式是保存模型为 PB 文件,它具有语言独立性,可独立运行,封闭的序列化格式,任何语言都可以解析它,它允许其他语言和深度学习框架读取、继续训练和迁移 TensorFlow 的模型
它的主要使用场景是实现创建模型与使用模型的解耦, 使得前向推导inference的代码统一
另外的好处是保存为PB文件时候,模型的变量都会变成固定的,导致模型的大小会大大减小,适合在手机端运行
还有一个就是,真正离线测试使用的时候,PB格式的数据能够保证数据不会更新变动,就是不会进行反馈调节
保存称为这种PB文件主要使用tf.SaveModelBuilder类来完成这个工作,并且可以把多个计算图保存到一个PB文件中,如果有多个MetaGraph,那么只会保留第一个的版本号,并且为每个MetaGraph指定特殊的名称tag用以区分,通常这个名称tag以该计算图的功能和使用到的设备命名,比如serving or training, cpu or gpu
保存为PB文件的代码示例如下
01.import tensorflow as tf02.import os03.from tensorflow.python.framework import graph_util04. 05.pb_file_path = os.getcwd()06. 07.with tf.Session(graph=tf.Graph()) as sess:08. x = tf.placeholder(tf.int32, name='x')09. y = tf.placeholder(tf.int32, name='y')10. b = tf.Variable(1, name='b')11. xy = tf.multiply(x, y)12. # 这里的输出需要加上name属性13. op = tf.add(xy, b, name='op_to_store')14. 15. sess.run(tf.global_variables_initializer())16. 17. # convert_variables_to_constants 需要指定output_node_names,list(),可以多个18. constant_graph = graph_util.convert_variables_to_constants(sess, sess.graph_def, ['op_to_store'])19. 20. # 测试 OP21. feed_dict = {x: 10, y: 3}22. print(sess.run(op, feed_dict))23. 24. # 写入序列化的 PB 文件25. with tf.gfile.FastGFile(pb_file_path+'model.pb', mode='wb') as f:26. f.write(constant_graph.SerializeToString())27. 28.#打印输出29.INFO:tensorflow:Froze 1 variables.30.Converted 1 variables to const ops.31.31
加载PB模型文件的代码示例如下
01.from tensorflow.python.platform import gfile02. 03.sess = tf.Session()04.with gfile.FastGFile(pb_file_path+'model.pb', 'rb') as f:05. graph_def = tf.GraphDef()06. graph_def.ParseFromString(f.read())07. sess.graph.as_default()08. tf.import_graph_def(graph_def, name='') # 导入计算图09. 10.# 需要有一个初始化的过程 11.sess.run(tf.global_variables_initializer())12. 13.# 需要先复原变量14.print(sess.run('b:0'))15.# 116. 17.# 输入18.input_x = sess.graph.get_tensor_by_name('x:0')19.input_y = sess.graph.get_tensor_by_name('y:0')20. 21.op = sess.graph.get_tensor_by_name('op_to_store:0')22. 23.ret = sess.run(op, feed_dict={input_x: 5, input_y: 5})24.print(ret)25.#打印输出26.26
保存为 save model 格式也可以生成模型的 PB 文件,并且更加简单。
保存好以后到saved_model_dir目录下,会有一个saved_model.pb文件以及variables文件夹。顾名思义,variables保存所有变量,saved_model.pb用于保存模型结构等信息
01.import tensorflow as tf02.import os03.from tensorflow.python.framework import graph_util04. 05.pb_file_path = os.getcwd()06. 07.with tf.Session(graph=tf.Graph()) as sess:08. x = tf.placeholder(tf.int32, name='x')09. y = tf.placeholder(tf.int32, name='y')10. b = tf.Variable(1, name='b')11. xy = tf.multiply(x, y)12. # 这里的输出需要加上name属性13. op = tf.add(xy, b, name='op_to_store')14. 15. sess.run(tf.global_variables_initializer())16. 17. # convert_variables_to_constants 需要指定output_node_names,list(),可以多个18. constant_graph = graph_util.convert_variables_to_constants(sess, sess.graph_def, ['op_to_store'])19. 20. # 测试 OP21. feed_dict = {x: 10, y: 3}22. print(sess.run(op, feed_dict))23. 24. # 写入序列化的 PB 文件25. with tf.gfile.FastGFile(pb_file_path+'model.pb', mode='wb') as f:26. f.write(constant_graph.SerializeToString())27. 28. # INFO:tensorflow:Froze 1 variables.29. # Converted 1 variables to const ops.30. # 3131. 32. 33. # 官网有误,写成了 saved_model_builder 34. builder = tf.saved_model.builder.SavedModelBuilder(pb_file_path+'savemodel')35. # 构造模型保存的内容,指定要保存的 session,特定的 tag, 36. # 输入输出信息字典,额外的信息37. builder.add_meta_graph_and_variables(sess,['cpu_server_1'])38. 39. 40.# 添加第二个 MetaGraphDef 41.#with tf.Session(graph=tf.Graph()) as sess:42.# ...43.# builder.add_meta_graph([tag_constants.SERVING])44.#...45. 46.builder.save() # 保存 PB 模型
上述方法对应的导入模型的方法
01.with tf.Session(graph=tf.Graph()) as sess:02. tf.saved_model.loader.load(sess, ['cpu_1'], pb_file_path+'savemodel')03. sess.run(tf.global_variables_initializer())04. 05. input_x = sess.graph.get_tensor_by_name('x:0')06. input_y = sess.graph.get_tensor_by_name('y:0')07. 08. op = sess.graph.get_tensor_by_name('op_to_store:0')09. 10. ret = sess.run(op, feed_dict={input_x: 5, input_y: 5})11. print(ret)
只需要指定要恢复模型的 session,模型的 tag,模型的保存路径即可,使用起来更加简单
这样和之前的导入 PB 模型一样,也是要知道tensor的name。那么如何可以在不知道tensor name的情况下使用呢,实现彻底的解耦呢? 给add_meta_graph_and_variables方法传入第三个参数,signature_def_map即可
- 标签
- tensorflow
- pb
内容版权声明:除非注明,否则皆为本站原创文章。
转载注明出处:https://sulao.cn/post/861
评论列表