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


Posted in Python onApril 23, 2018

有时候,我们需要将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 相关文章推荐
Python实现高效求解素数代码实例
Jun 30 Python
Python3 模块、包调用&路径详解
Oct 25 Python
numpy中实现ndarray数组返回符合特定条件的索引方法
Apr 17 Python
详解python持久化文件读写
Apr 06 Python
python日期相关操作实例小结
Jun 24 Python
Python-copy()与deepcopy()区别详解
Jul 12 Python
python匿名函数用法实例分析
Aug 03 Python
Python 3 判断2个字典相同
Aug 06 Python
python 解决tqdm模块不能单行显示的问题
Feb 19 Python
Django 5种类型Session使用方法解析
Apr 29 Python
Python 类,对象,数据分类,函数参数传递详解
Sep 25 Python
python中tkinter复选框使用操作
Nov 11 Python
tensorflow1.0学习之模型的保存与恢复(Saver)
Apr 23 #Python
tensorflow 使用flags定义命令行参数的方法
Apr 23 #Python
Tensorflow之Saver的用法详解
Apr 23 #Python
python获取文件路径、文件名、后缀名的实例
Apr 23 #Python
Python基于FTP模块实现ftp文件上传操作示例
Apr 23 #Python
Python基于whois模块简单识别网站域名及所有者的方法
Apr 23 #Python
Python实现自定义顺序、排列写入数据到Excel的方法
Apr 23 #Python
You might like
分页显示Oracle数据库记录的类之一
2006/10/09 PHP
mysql From_unixtime及UNIX_TIMESTAMP及DATE_FORMAT日期函数
2010/03/21 PHP
php使用mkdir创建多级目录入门例子
2014/05/10 PHP
php图片裁剪函数
2018/10/31 PHP
Yii2.0框架模型多表关联查询示例
2019/07/18 PHP
js 处理URL实用技巧
2010/11/23 Javascript
实现网页页面跳转的几种方法(meta标签、js实现、php实现)
2014/05/20 Javascript
Clipboard.js 无需Flash的JavaScript复制粘贴库
2015/10/02 Javascript
详解JavaScript设计模式开发中的桥接模式使用
2016/05/18 Javascript
微信公众号  提示:Unauthorized API function 问题解决方法
2016/12/05 Javascript
JS无缝滚动效果实现方法分析
2016/12/21 Javascript
jQuery动态移除与增加onclick属性的方法详解
2018/06/07 jQuery
判断js数据类型的函数实例详解
2019/05/23 Javascript
详解vue beforeRouteEnter 异步获取数据给实例问题
2019/08/09 Javascript
jquery实现商品sku多属性选择功能(商品详情页)
2019/12/20 jQuery
JavaScript中this的学习笔记及用法整理
2020/02/17 Javascript
python实现连接mongodb的方法
2015/05/08 Python
python实现将内容分行输出
2015/11/05 Python
讲解Python的Scrapy爬虫框架使用代理进行采集的方法
2016/02/18 Python
Python入门_学会创建并调用函数的方法
2017/05/16 Python
Python实现的本地文件搜索功能示例【测试可用】
2018/05/30 Python
Python 带有参数的装饰器实例代码详解
2018/12/06 Python
pyqt5 获取显示器的分辨率的方法
2019/06/18 Python
python找出因数与质因数的方法
2019/07/25 Python
解决Django加载静态资源失败的问题
2019/07/28 Python
原生python实现knn分类算法
2019/10/24 Python
Python异步编程之协程任务的调度操作实例分析
2020/02/01 Python
Python map及filter函数使用方法解析
2020/08/06 Python
UNIONBAY官网:美国青少年服装品牌
2019/03/26 全球购物
工作睡觉检讨书
2014/02/25 职场文书
节水口号标语
2014/06/19 职场文书
2015年高一班主任工作总结
2015/05/13 职场文书
《认识钟表》教学反思
2016/02/16 职场文书
技术转让协议书
2016/03/19 职场文书
读《茶花女》有感:山茶花的盛开与凋零
2020/01/17 职场文书
MySQL Shell的介绍以及安装
2021/04/24 MySQL