浅谈tensorflow模型保存为pb的各种姿势资源-CSDN文库

19 浏览量 2020-09-16 20:04:04 上传评论收藏 105KB PDF 举报

资源推荐

资源详情

资源评论

浅谈浅谈tensorflow模型保存为模型保存为pb的各种姿势的各种姿势

主要介绍了浅谈tensorflow模型保存为pb的各种姿势，具有很好的参考价值，希望对大家有所帮助。一起跟随小

编过来看看吧

一，直接保存一，直接保存pb

1, 首先我们当然可以直接在tensorflow训练中直接保存为pb为格式，保存pb的好处就是使用场景是实现创建模型与使用模型的

解耦，使得创建模型与使用模型的解耦，使得前向推导inference代码统一。另外的好处就是保存为pb的时候，模型的变量会

变成固定的，导致模型的大小会大大减小。

这里稍稍解释下pb：是MetaGraph的protocol buffer格式的文件，MetaGraph包括计算图，数据流，以及相关的变量和输入输

出

主要使用tf.SavedModelBuilder来完成这个工作，并且可以把多个计算图保存到一个pb文件中，如果有多个MetaGraph，那么

只会保留第一个MetaGraph的版本号。

保持pb的文件代码：

import tensorflow as tf

import os

from tensorflow.python.framework import graph_util

pb_file_path = os.getcwd()

with tf.Session(graph=tf.Graph()) as sess:

x = tf.placeholder(tf.int32, name='x')

y = tf.placeholder(tf.int32, name='y')

b = tf.Variable(1, name='b')

xy = tf.multiply(x, y)

# 这里的输出需要加上name属性

op = tf.add(xy, b, name='op_to_store')

sess.run(tf.global_variables_initializer())

# convert_variables_to_constants 需要指定output_node_names，list()，可以多个

constant_graph = graph_util.convert_variables_to_constants(sess, sess.graph_def, ['op_to_store'])

# 测试 OP

feed_dict = {x: 10, y: 3}

print(sess.run(op, feed_dict))

# 写入序列化的 PB 文件

with tf.gfile.FastGFile(pb_file_path+'model.pb', mode='wb') as f:

f.write(constant_graph.SerializeToString())

# 输出

# INFO:tensorflow:Froze 1 variables.

# Converted 1 variables to const ops.

# 31

其实主要是：

# convert_variables_to_constants 需要指定output_node_names，list()，可以多个

constant_graph = graph_util.convert_variables_to_constants(sess, sess.graph_def, ['op_to_store'])

# 写入序列化的 PB 文件

with tf.gfile.FastGFile(pb_file_path+'model.pb', mode='wb') as f:

f.write(constant_graph.SerializeToString())

1.1 加载测试代码

from tensorflow.python.platform import gfile

sess = tf.Session()

with gfile.FastGFile(pb_file_path+'model.pb', 'rb') as f:

graph_def = tf.GraphDef()

graph_def.ParseFromString(f.read())

sess.graph.as_default()

tf.import_graph_def(graph_def, name='') # 导入计算图

# 需要有一个初始化的过程

sess.run(tf.global_variables_initializer())

# 需要先复原变量

print(sess.run('b:0'))

# 1

本内容试读结束，登录后可阅读更多

下载后可阅读完整内容，剩余6页未读，立即下载

内容反馈

weixin_38545959

粉丝: 1
资源: 928

最新资源

资源上传下载、课程学习等过程中有任何疑问或建议，欢迎提出宝贵意见哦~我们会及时处理！点击此处反馈

feedback-tip