Tensorflow将模型导出为一个文件及接口设置

在上一篇文章中《Tensorflow加载预训练模型和保存模型》,我们学习到如何使用预训练的模型。但注意到,在上一篇文章中使用预训练模型,必须至少的要4个文件:

checkpoint
MyModel.meta
MyModel.data-00000-of-00001
MyModel.index

这很不便于我们的使用。有没有办法导出为一个pb文件,然后直接使用呢?答案是肯定的。在文章《Tensorflow加载预训练模型和保存模型》中提到,meta文件保存图结构,weights等参数保存在data文件中。也就是说,图和参数数据时分开保存的。说的更直白一点,就是meta文件中没有weights等数据。但是,值得注意的是,meta文件会保存常量。我们只需将data文件中的参数转为meta文件中的常量即可!

1 模型导出为一个文件

1.1 有代码并且从头开始训练

Tensorflow提供了工具函数tf.graph_util.convert_variables_to_constants()用于将变量转为常量。看看官网的描述:

if you have a trained graph containing Variable ops, it can be convenient to convert them all to Const ops holding the same values. This makes it possible to describe the network fully with a single GraphDef file, and allows the removal of a lot of ops related to loading and saving the variables.

我们继续通过一个简单例子开始:

import tensorflow as tf

w1 = tf.Variable(20.0, name="w1")
w2 = tf.Variable(30.0, name="w2")
b1= tf.Variable(2.0,name="bias")
w3 = tf.add(w1,w2)

#记住要定义name,后面需要用到
out = tf.multiply(w3,b1,name="out")

# 转换Variable为constant,并将网络写入到文件
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    # 这里需要填入输出tensor的名字
    graph = tf.graph_util.convert_variables_to_constants(sess, sess.graph_def, ["out"])
    tf.train.write_graph(graph, '.', './checkpoint_dir/graph.pb', as_text=False)

执行可以看到如下日志:

Converted 3 variables to const ops.

可以看到通过tf.graph_util.convert_variables_to_constants()函数将变量转为了常量,并存储在graph.pb文件中,接下来看看如何使用这个模型。

import tensorflow as tf
with tf.Session() as sess:
    with open('./checkpoint_dir/graph.pb', 'rb') as graph:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(graph.read())
        output = tf.import_graph_def(graph_def, return_elements=['out:0'])
        print(sess.run(output))

运行结果如下:

[100.0]

回到tf.graph_util.convert_variables_to_constants()函数,可以看到,需要传入Session对象和图,这都可以理解。看看第三个参数["out"],它是指定这个模型的输出Tensor

1.2 有代码和模型,但是不想重新训练模型

有模型源码时,在导出模型时就可以通过tf.graph_util.convert_variables_to_constants()函数来将变量转为常量保存到图文件中。但是很多时候,我们拿到的是别人的checkpoint文件,即meta、index、data等文件。这种情况下,需要将data文件里面变量转为常量保存到meta文件中。思路也很简单,先将checkpoint文件加载,再重新保存一次即可。

假设训练和保存模型代码如下:

import tensorflow as tf

w1 = tf.Variable(20.0, name="w1")
w2 = tf.Variable(30.0, name="w2")
b1= tf.Variable(2.0,name="bias")
w3 = tf.add(w1,w2)

#记住要定义name,后面需要用到
out = tf.multiply(w3,b1,name="out")

# 转换Variable为constant,并将网络写入到文件
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    saver = tf.train.Saver()
    # 这里需要填入输出tensor的名字
    saver.save(sess, './checkpoint_dir/MyModel', global_step=1000)

此时,模型文件如下:

checkpoint
MyModel-1000.data-00000-of-00001
MyModel-1000.index
MyModel-1000.meta

如果我们只有以上4个模型文件,但是可以看到训练源码。那么,将这4个文件导出为一个pb文件方法如下:

import tensorflow as tf
with tf.Session() as sess:

    #初始化变量
    sess.run(tf.global_variables_initializer())

    #获取最新的checkpoint,其实就是解析了checkpoint文件
    latest_ckpt = tf.train.latest_checkpoint("./checkpoint_dir")

    #加载图
    restore_saver = tf.train.import_meta_graph('./checkpoint_dir/MyModel-1000.meta')

    #恢复图,即将weights等参数加入图对应位置中
    restore_saver.restore(sess, latest_ckpt)

    #将图中的变量转为常量
    output_graph_def = tf.graph_util.convert_variables_to_constants(
        sess, sess.graph_def , ["out"])
    #将新的图保存到"/pretrained/graph.pb"文件中
    tf.train.write_graph(output_graph_def, 'pretrained', "graph.pb", as_text=False)

执行后,会有如下日志:

Converted 3 variables to const ops.

接下来就是使用,使用方法跟前面一致:

import tensorflow as tf
with tf.Session() as sess:
    with open('./pretrained/graph.pb', 'rb') as graph:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(graph.read())
        output = tf.import_graph_def(graph_def, return_elements=['out:0'])
        print(sess.run(output))

打印信息如下:

[100.0]

2 模型接口设置

我们注意到,前面只是简单的获取一个输出接口,但是很明显,我们使用的时候,不可能只有一个输出,还需要有输入,接下来我们看看,如何设置输入和输出。同样我们分为有代码并且从头开始训练,和有代码和模型,但是不想重新训练模型两种情况。

2.1 有代码并且从头开始训练

相比1.1中的代码略作修改即可,第6行代码处做了修改:

import tensorflow as tf

w1 = tf.Variable(20.0, name="w1")
w2 = tf.Variable(30.0, name="w2")

#这里将b1改为placeholder,让用户输入,而不是写死
#b1= tf.Variable(2.0,name="bias")
b1= tf.placeholder(tf.float32, name='bias')

w3 = tf.add(w1,w2)

#记住要定义name,后面需要用到
out = tf.multiply(w3,b1,name="out")

# 转换Variable为constant,并将网络写入到文件
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    # 这里需要填入输出tensor的名字
    graph = tf.graph_util.convert_variables_to_constants(sess, sess.graph_def, ["out"])
    tf.train.write_graph(graph, '.', './checkpoint_dir/graph.pb', as_text=False)

日志如下:

Converted 2 variables to const ops.

接下来看看如何使用:

import tensorflow as tf
with tf.Session() as sess:
    with open('./checkpoint_dir/graph.pb', 'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
        output = tf.import_graph_def(graph_def, input_map={'bias:0':4.}, return_elements=['out:0'])
        print(sess.run(output))

打印信息如下:

[200.0]

也就是说,在设置输入时,首先将需要输入的数据作为placeholdler,然后在导入图tf.import_graph_def()时,通过参数input_map={}来指定输入。输出通过return_elements=[]直接引用tensor的name即可。

2.2 有代码和模型,但是不想重新训练模型

在有代码和模型,但是不想重新训练模型情况下,意味着我们不能直接修改导出模型的代码。但是我们可以通过graph.get_tensor_by_name()函数取得图中的某些中间结果,然后再加入一些逻辑。其实这种情况在上一篇文章已经讲了。可以参考上一篇文件解决,相比“有代码并且从头开始训练”情况局限比较大,大部分情况只能是获取模型的一些中间结果,但是也满足我们大多数情况使用了。

本文参与腾讯云自媒体分享计划,欢迎正在阅读的你也加入,一起分享。

发表于

我来说两句

0 条评论
登录 后参与评论

相关文章

来自专栏10km的专栏

SSD(Single Shot MultiBox Detector):ubuntu16安装及训练自己的数据集(VOC2007格式)过程记录

安装SSD # SSD代码clone到 caffe-ssd文件夹下 git clone --recursive https://github.com/weili...

52460
来自专栏CNN

从Tensorflow模型文件中解析并显示网络结构图(CKPT模型篇)

本文介绍如何从CKPT模型文件中提取网络结构图并实现可视化。

1.8K30
来自专栏PaddlePaddle

【FAQ】本地训练与预测相关问题汇总

导语 在使用指南的最后一部分,我们汇总了使用PaddlePaddle过程中的常见问题,本部分推文目录如下: 2.22:【FAQ】模型配置相关问题汇总 2.23:...

377100
来自专栏蘑菇先生的技术笔记

探索c#之一致性Hash详解

26250
来自专栏fangyangcoder

tensorflow笔记(一)之基础知识

http://www.cnblogs.com/fydeblog/p/7399701.html

12020
来自专栏wOw的Android小站

[Tensorflow] Faster R-CNN 和自定义 VOC 数据集

看了pascal_voc.py代码,可以把代码的jpg拼接改成png,这样可以不做上一步.

1.8K20
来自专栏PaddlePaddle

【进阶篇】C-API 使用流程

编写|PaddlePaddle 排版|wangp 这篇推文介绍 PaddlePaddle C-API 整体使用流程 1 使用流程 使用 C-API 的工作流...

37730
来自专栏目标检测和深度学习

深度学习工程模板:简化加载数据、构建网络、训练模型和预测样本的流程

注意:支持在训练中调用callbacks,额外添加模型存储、TensorBoard、FPR度量等。

14640
来自专栏机器之心

教程 | TensorFlow 官方解读:如何在多系统和网络拓扑中构建高性能模型

选自Tensorflow 机器之心编译 参与:黄玉胜、黄小天 这个文档和附带的脚本详细介绍了如何构建针对各种系统和网络拓扑的高性能可拓展模型。这个技术在本文档中...

393110
来自专栏流柯技术学院

jmeter参数化随机取值实现

jmeter能用来做参数化的组件有几个,但是都没有随机取值的功能,遇到随机取值的需求怎么办呢?

27120

扫码关注云+社区

领取腾讯云代金券