将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
以上就是本文的全部内容,希望对大家的学习有所帮助,也希望大家多多支持脚本之家。
- 如何定位Obj-C野指针随机Crash(二):让非必现Crash变成必现
- 六种开发环境部署大全:基于Openshift
- 手把手教你由TensorFlow上手PyTorch(附代码)
- 如何定位Obj-C野指针随机Crash(三):如何让Crash自报家门
- 【MIG专项测试组】腾讯手机管家实战分析:内存突增是为神马?
- 全系统栈崩溃是什么鬼?手机管家高级工程师 jaylin,教你如何抓鬼!
- 【团队分享】苍翼之刃:论File Descriptor泄漏如何导致Crash?
- 【沙龙干货分享】你要知道的N个Android适配问题
- 优化安卓应用内存的神秘方法以及背后的原理,一般人我不告诉他
- 安全快速下载Xcode的姿势
- Android之实现妙趣横生的粘连布局
- 萝莉有话说:你的App真正适配了iOS 9吗?
- H5 缓存机制浅析 移动端 Web 加载性能优化
- Android 开发必备知识:我和 Gradle 有个约会
- JavaScript 教程
- JavaScript 编辑工具
- JavaScript 与HTML
- JavaScript 与Java
- JavaScript 数据结构
- JavaScript 基本数据类型
- JavaScript 特殊数据类型
- JavaScript 运算符
- JavaScript typeof 运算符
- JavaScript 表达式
- JavaScript 类型转换
- JavaScript 基本语法
- JavaScript 注释
- Javascript 基本处理流程
- Javascript 选择结构
- Javascript if 语句
- Javascript if 语句的嵌套
- Javascript switch 语句
- Javascript 循环结构
- Javascript 循环结构实例
- Javascript 跳转语句
- Javascript 控制语句总结
- Javascript 函数介绍
- Javascript 函数的定义
- Javascript 函数调用
- Javascript 几种特殊的函数
- JavaScript 内置函数简介
- Javascript eval() 函数
- Javascript isFinite() 函数
- Javascript isNaN() 函数
- parseInt() 与 parseFloat()
- escape() 与 unescape()
- Javascript 字符串介绍
- Javascript length属性
- javascript 字符串函数
- Javascript 日期对象简介
- Javascript 日期对象用途
- Date 对象属性和方法
- Javascript 数组是什么
- Javascript 创建数组
- Javascript 数组赋值与取值
- Javascript 数组属性和方法
- 【机器学习基础】数学推导+纯Python实现机器学习算法26:随机森林
- 【机器学习基础】Python数据预处理:彻底理解标准化和归一化
- Logback配置文件这么写,TPS提高10倍
- 重学数据结构之队列
- 太简单,springboo 使用自定义的 Tomcat
- 没想到,几行代码,你就可以实现图片压缩(springboot)!
- Go 语言学习之基础数据类型
- Go 语言学习之变量
- Go 泛型的括号选择:[ ] or ( )?
- 如何使用GitLab CI/CD 触发多项目管道
- 使用docker数据卷持久化容器数据
- 面向初学者的Docker快速入门指南
- TypeScript:React、拖拽、实践!
- 太慢不能忍!CPU又拿硬盘和网卡开刀了!
- 懂了!VMware/KVM/Docker原来是这么回事儿