将TensorFlow的模型网络导出为单个文件的方法
有时候，我们需要将TensorFlow的模型导出为单个文件（同时包含模型架构定义与权重），方便在其他地方使用（如在c++中部署网络）。利用tf.train.write_graph()默认情况下只导出了网络的定义（没有权重），而利用tf.train.Saver().save()导出的文件graph_def与权重是分离的，因此需要采用别的方法。
我们知道，graph_def文件中没有包含网络中的Variable值（通常情况存储了权重），但是却包含了constant值，所以如果我们能把Variable转换为constant，即可达到使用一个文件同时存储网络架构与权重的目标。
我们可以采用以下方式冻结权重并保存网络：
importtensorflowastf fromtensorflow.python.framework.graph_utilimportconvert_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，并将网络写入到文件 withtf.Session()assess: 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)
当恢复网络时，可以使用如下方式：
importtensorflowastf
withtf.Session()assess:
withopen('./graph.pb','rb')asf:
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)]
可以看到之前的权重确实保存了下来!!
问题来了，我们的网络需要能有一个输入自定义数据的接口啊！不然这玩意有什么用。。别急，当然有办法。
importtensorflowastf fromtensorflow.python.framework.graph_utilimportconvert_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') withtf.Session()assess: 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，下面来看看怎么恢复网络并输入自定义数据。
importtensorflowastf
withtf.Session()assess:
withopen('./graph.pb','rb')asf:
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，如下所示：
importtensorflowastf
new_input=tf.placeholder(tf.float32,shape=())
withtf.Session()assess:
withopen('./graph.pb','rb')asf:
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了，则在导入网络的时候，需要做一点小修改。
importtensorflowastf
fromgoogle.protobufimporttext_format
withtf.Session()assess:
#不使用'rb'模式
withopen('./graph.pb','r')asf:
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))
参考资料
IsthereanexampleonhowtogenerateprotobuffilesholdingtrainedTensorflowgraphs
以上就是本文的全部内容，希望对大家的学习有所帮助，也希望大家多多支持毛票票。