查看: 86|回复: 0

扣丁学堂浅谈将TensorFlow的模型网络导出为单个文件的方法

发表于 2018-6-1 11:20:23

今天给大家分享的是将TensorFlow的模型网络导出为单个文件的方法,喜欢Python开发的小伙伴和扣丁学堂Python在线学习小编一块来看一下吧。


我们可以采用以下方式冻结权重并保存网络:
import tensorflow as tf
from tensorflow.python.framework.graph_util import convert_variables_to_constants
# 构造网络
a = tf.Variable([[3],[4]], dtype=tf.float32, name='a')
b = tf.Variable(4, dtype=tf.float32, name='b')
# 一定要给输出tensor取一个名字!!
output = tf.add(a, b, name='out')
# 转换Variable为constant,并将网络写入到文件
with tf.Session() as sess:
  sess.run(tf.global_variables_initializer())
  # 这里需要填入输出tensor的名字
  graph = convert_variables_to_constants(sess, sess.graph_def, ["out"])
  tf.train.write_graph(graph, '.', 'graph.pb', as_text=False)
当恢复网络时,可以使用如下方式:
import tensorflow as tf
with tf.Session() as sess:
  with open('./graph.pb', 'rb') as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())
    output = tf.import_graph_def(graph_def, return_elements=['out:0'])
    print(sess.run(output))
输出结果为:
[array([[ 7.],       [ 8.]], dtype=float32)]
可以看到之前的权重确实保存了下来。
问题来了,我们的网络需要能有一个输入自定义数据的接口啊!不然这玩意有什么用。。别急,当然有办法。
import tensorflow as tf
from tensorflow.python.framework.graph_util import convert_variables_to_constants
a = tf.Variable([[3],[4]], dtype=tf.float32, name='a')
b = tf.Variable(4, dtype=tf.float32, name='b')
input_tensor = tf.placeholder(tf.float32, name='input')
output = tf.add((a+b), input_tensor, name='out')
with tf.Session() as sess:
  sess.run(tf.global_variables_initializer())
  graph = convert_variables_to_constants(sess, sess.graph_def, ["out"])
  tf.train.write_graph(graph, '.', 'graph.pb', as_text=False)
用上述代码重新保存网络至graph.pb,这次我们有了一个输入placeholder,下面来看看怎么恢复网络并输入自定义数据。
import tensorflow as tf
with tf.Session() as sess:
  with open('./graph.pb', 'rb') as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())
    output = tf.import_graph_def(graph_def, input_map={'input:0':4.}, return_elements=['out:0'], name='a')
    print(sess.run(output))
以上就是扣丁学堂Java在线学习小编给大家分享的将TensorFlow的模型网络导出为单个文件的方法,希望对大家的学习有所帮助,想要了解更多Python方面内容的小伙伴可以登录扣丁学堂官网咨询。扣丁学堂是专业的Python培训机构,不仅有专业的老师和与时俱进的课程体系,还有大量的Python在线教程供学员观看学习,想要学好Python开发的小伙伴快快行动吧。扣丁学堂Python技术交流群:279521237。

文章转载自:扣丁学堂http://www.codingke.com/


回复

使用道具 举报