有时候,我们需要将TensorFlow的模型导出为单个文件(同时包含模型架构定义与权重),方便在其他地方使用(如在c++中部署网络)。利用tf.train.write_graph()默认情况下只导出了网络的定义(没有权重),而利用tf.train.Saver().save()导出的文件graph_def与权重是分离的,因此需要采用别的方法。
我们知道,graph_def文件中没有包含网络中的Variable值(通常情况存储了权重),但是却包含了constant值,所以如果我们能把Variable转换为constant,即可达到使用一个文件同时存储网络架构与权重的目标。
我们可以采用以下方式冻结权重并保存网络:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
|
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 )
|
当恢复网络时,可以使用如下方式:
1
2
3
4
5
6
7
|
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)]
可以看到之前的权重确实保存了下来!!
问题来了,我们的网络需要能有一个输入自定义数据的接口啊!不然这玩意有什么用。。别急,当然有办法。
1
2
3
4
5
6
7
8
9
10
11
|
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,下面来看看怎么恢复网络并输入自定义数据。
1
2
3
4
5
6
7
8
|
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))
|
输出结果为:
[array([[ 11.],
[ 12.]], dtype=float32)]
可以看到结果没有问题,当然在input_map那里可以替换为新的自定义的placeholder,如下所示:
1
2
3
4
5
6
7
8
9
10
|
import tensorflow as tf
new_input = tf.placeholder(tf.float32, shape = ())
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' :new_input}, return_elements = [ 'out:0' ], name = 'a' )
print (sess.run(output, feed_dict = {new_input: 4 }))
|
看看输出,同样没有问题。
[array([[ 11.],
[ 12.]], dtype=float32)]
另外需要说明的一点是,在利用tf.train.write_graph写网络架构的时候,如果令as_text=True了,则在导入网络的时候,需要做一点小修改。
1
2
3
4
5
6
7
8
9
10
11
|
import tensorflow as tf
from google.protobuf import text_format
with tf.Session() as sess:
# 不使用'rb'模式
with open ( './graph.pb' , 'r' ) as f:
graph_def = tf.GraphDef()
# 不使用graph_def.ParseFromString(f.read())
text_format.Merge(f.read(), graph_def)
output = tf.import_graph_def(graph_def, return_elements = [ 'out:0' ])
print (sess.run(output))
|
参考资料
Is there an example on how to generate protobuf files holding trained Tensorflow graphs
以上就是本文的全部内容,希望对大家的学习有所帮助,也希望大家多多支持服务器之家。
原文链接:https://blog.csdn.net/encodets/article/details/54428456