将TensorFlow的模型网络导出为单个文件的方法

更新时间：2018年04月23日 14:44:53 作者：EncodeTS

本篇文章主要介绍了将TensorFlow的网络导出为单个文件的方法，小编觉得挺不错的，现在分享给大家，也给大家做个参考。一起跟随小编过来看看吧

有时候，我们需要将TensorFlow的模型导出为单个文件（同时包含模型架构定义与权重），方便在其他地方使用（如在c++中部署网络）。利用tf.train.write_graph()默认情况下只导出了网络的定义（没有权重），而利用tf.train.Saver().save()导出的文件graph_def与权重是分离的，因此需要采用别的方法。

我们知道，graph_def文件中没有包含网络中的Variable值（通常情况存储了权重），但是却包含了constant值，所以如果我们能把Variable转换为constant，即可达到使用一个文件同时存储网络架构与权重的目标。

我们可以采用以下方式冻结权重并保存网络：

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))

输出结果为：

[array([[ 11.],
[ 12.]], dtype=float32)]

可以看到结果没有问题，当然在input_map那里可以替换为新的自定义的placeholder，如下所示：

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了，则在导入网络的时候，需要做一点小修改。

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

以上就是本文的全部内容，希望对大家的学习有所帮助，也希望大家多多支持脚本之家。

您可能感兴趣的文章:

Python实现FTP文件传输的实例
在本篇文章里小编给各位分享的是关于Python实现FTP文件传输的实例以及相关代码，需要的朋友们学习下。
2019-07-07
Python零基础入门学习之输入与输出
在之前的编程中，我们的信息打印，数据的展示都是在控制台（命令行）直接输出的，信息都是一次性的没有办法复用和保存以便下次查看，今天我们将学习Python的输入输出，解决以上问题
2019-04-04
windows上彻底删除jupyter notebook的实现
这篇文章主要介绍了windows上彻底删除jupyter notebook的实现，具有很好的参考价值，希望对大家有所帮助。一起跟随小编过来看看吧
2020-04-04
Pygame游戏开发之太空射击实战精灵的使用上篇
相信大多数8090后都玩过太空射击游戏，在过去游戏不多的年代太空射击自然属于经典好玩的一款了，今天我们来自己动手实现它，在编写学习中回顾过往展望未来，下面开始讲解精灵的使用
2022-08-08
写一个Python脚本自动爬取Bilibili小视频
这篇文章主要介绍了写一个 Python 脚本自动爬取 Bilibili 小视频的示例代码，帮助大家更好的理解和学习使用python，感兴趣的朋友可以了解下
2021-04-04
使用Flask集成bootstrap的方法
这篇文章主要介绍了使用Flask集成bootstrap的方法，小编觉得挺不错的，现在分享给大家，也给大家做个参考。一起跟随小编过来看看吧
2018-07-07
用python代码做configure文件
在lua中，我一直用lua作为config文件，或者承载数据的文件 - 好处是lua本身就很好阅读，然后无需额外写解析的代码，还支持在configure文件中读环境变量，条件判断等，方便又强大！
2014-07-07
Python批量实现word中查找关键字的示例代码
本文主要介绍了Python批量实现word中查找关键字的示例代码，文中通过示例代码介绍的非常详细，对大家的学习或者工作具有一定的参考学习价值，需要的朋友们下面随着小编来一起学习学习吧
2023-07-07
Anaconda下安装mysql-python的包实例
今天小编就为大家分享一篇Anaconda下安装mysql-python的包实例，具有很好的参考价值，希望对大家有所帮助。一起跟随小编过来看看吧
2018-06-06
python批量修改图片大小的方法
这篇文章主要为大家详细介绍了python批量修改图片大小的方法，具有一定的参考价值，感兴趣的小伙伴们可以参考一下
2018-07-07