将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实现CNN-GRU故障诊断的代码示例

    使用python实现CNN-GRU故障诊断的代码示例

    这篇文章主要给大家详细介绍了如何使用python实现CNN-GRU故障诊断,文章中有详细的代码示例,具有一定的参考价值,需要的朋友可以参考下
    2023-07-07
  • python中的netCDF4批量处理NC文件的操作方法

    python中的netCDF4批量处理NC文件的操作方法

    这篇文章主要介绍了python的netCDF4批量处理NC格式文件的操作方法,使用python批量提取所有数据,查看数据属性,本文通过实例代码给大家介绍的非常详细,对大家的学习或工作具有一定的参考借鉴价值,需要的朋友可以参考下
    2022-03-03
  • Python实现文件复制删除

    Python实现文件复制删除

    本文通过2个具体的实例,给大家展示了如何使用Python实现文件的复制与删除,非常的简单实用,有需要的小伙伴可以参考下
    2016-04-04
  • 基于MATLAB和Python实现MFCC特征参数提取

    基于MATLAB和Python实现MFCC特征参数提取

    这篇文章主要介绍了基于MATLAB和Python实现MFCC特征参数提取,文中通过示例代码介绍的非常详细,对大家的学习或者工作具有一定的参考学习价值,需要的朋友们下面随着小编来一起学习学习吧
    2019-08-08
  • 解决Django连接db遇到的问题

    解决Django连接db遇到的问题

    今天小编大家分享一篇解决Django连接db遇到的问题,具有很好的参考价值,希望对大家有所帮助。一起跟随小编过来看看吧
    2019-08-08
  • Python中一些自然语言工具的使用的入门教程

    Python中一些自然语言工具的使用的入门教程

    这篇文章主要介绍了Python中一些自然语言工具的使用的入门教程,本文来自于IBM官方网站技术文档,需要的朋友可以参考下
    2015-04-04
  • Python与数据库的交互问题小结

    Python与数据库的交互问题小结

    这篇文章主要介绍了Python与数据库的交互,本文给大家介绍的非常详细,对大家的学习或工作具有一定的参考借鉴价值,需要的朋友可以参考下
    2021-12-12
  • Python实现感知器模型、两层神经网络

    Python实现感知器模型、两层神经网络

    这篇文章主要为大家详细介绍了Python实现感知器模型、两层神经网络,文中示例代码介绍的非常详细,具有一定的参考价值,感兴趣的小伙伴们可以参考一下
    2017-12-12
  • Python与MongoDB交互的代码实现

    Python与MongoDB交互的代码实现

    Python与MongoDB的交互通常通过pymongo库来实现,pymongo是一个官方的Python驱动程序,用于与MongoDB数据库进行交互,以下是一个简单的示例,具有一定的参考价值,需要的朋友可以参考下
    2024-10-10
  • python使用xlrd实现检索excel中某列含有指定字符串记录的方法

    python使用xlrd实现检索excel中某列含有指定字符串记录的方法

    这篇文章主要介绍了python使用xlrd实现检索excel中某列含有指定字符串记录的方法,涉及Python使用xlrd模块检索Excel的技巧,非常具有实用价值,需要的朋友可以参考下
    2015-05-05

最新评论