TensorFlow Lite学习笔记2:生成TFLite模型文件

简介

在桌面PC或是服务器上使用TensorFlow训练出来的模型文件,不能直接用在TFLite上运行,需要使用离线工具先转成.tflite文件。笔者发现官方文档中很多细节介绍的都不太明确,在使用过程中需要不断尝试。我把自己的尝试过的步骤分享出来,希望能帮助大家节省时间。

具体说来,tflite文件的生成大致分为3步:

1. 在算法训练的脚本中保存图模型文件(GraphDef)和变量文件(CheckPoint)。

2. 利用freeze_graph工具生成frozen的graphdef文件。

3. 利用toco(Tensorflow Optimizing COnverter)工具,生成最终的tflite文件。

图1. 生成tflite文件的整个流程示意图

第1步:导出图模型文件和变量文件

在你的算法的训练或推理任务的脚本中,利用tensorflow.train中的write_graph和saver API来导出GraphDef及Checkpoint文件。

图2. TensorFlow中导出GraphDef文件和Checkpoint文件

其中,tf.train.write_graph一行将导出模型的GraphDef文件,实际上保存了训练的神经网络的结构图信息。存储格式为protobuffer,所以文件名后缀为pb。

图3. 导出的GraphDef文件

tf.train.saver.save一行导出的是模型的变量文件,实际上保存了整个图中所有变量目前的取值。

图4. 导出的checkpoint文件

如图4所示,实际上产生了4个文件。在后续步骤中需要用到的是nsfw_model.ckpt.data-00000-of-00001这个文件,保存了当前神经网络各参数的取值。

第2步:生成frozen的graphdef文件

在此步骤中,使用Tensorflow源代码中自带的freeze_graph工具,生成一个frozen的GraphDef文件。

bazel-bin/tensorflow/python/tools/freeze_graph --input_graph=/data/deep_learning/nsfw/model/nsfw-graph.pb --input_checkpoint=/data/deep_learning/nsfw/model/nsfw_model.ckpt --input_binary=true --output_graph=/data/deep_learning/nsfw/model/frozen_nsfw.pb --output_node_names=predictions

这里有两个地方容易搞错。第一个地方,input_checkpoint参数实际上用到的文件应该是nsfw_model.ckpt.data-00000-of-00001,但是在指定文件名的时候只需要指定nsfw_model.ckpt即可。第二个地方,是output_node_names参数,此处指定的是神经网络图中的输出节点的名字,是在训练阶段的Python脚本中定义的。如下图所示,在定义网络结构时,输出节点的名称为"predictions"。则最终output_node_names需要指定为“predictions”。

图5. output_node_names参数取值与网络模型定义时的名字要对应

当然,也可以利用summarize_graph打印出模型的输入和输出节点,如:

bazel-bin/tensorflow/tools/graph_transforms/summarize_graph --in_graph=/data/deep_learning/nsfw/model/frozen_nsfw.pb

图6. 输入节点名称为input
图7. 输出节点名称为predictions

第3步:生成最终的tflite文件

在此步骤中,使用Tensorflow源代码中自带的toco工具,生成一个可供TensorFlow Lite框架使用tflite文件。其中input_arrays和output_arrays的名称需要与定义网络类型时取的名称保持一致。

bazel run --config=opt tensorflow/contrib/lite/toco:toco --input_file=/data/deep_learning/nsfw/model/frozen_nsfw.pb --input_format=TENSORFLOW_GRAPHDEF --output_format=TFLITE --output_file=/data/deep_learning/nsfw/model/nsfw.lite --inference_type=FLOAT --input_type=FLOAT --input_arrays=input --output_arrays=predictions --input_shapes=1,224,224,3

生成的nsfw.lite文件即可用于TensorFlow Lite应用。

最后编辑于
©著作权归作者所有,转载或内容合作请联系作者
平台声明:文章内容(如有图片或视频亦包括在内)由作者上传并发布,文章内容仅代表作者本人观点,简书系信息发布平台,仅提供信息存储服务。

推荐阅读更多精彩内容

  • 前言 本文主要参考了几篇文章,搭建了一个在iOS上跑Tensorflow MNIST模型的demo,本文会给出一个...
    Thanatos_defy阅读 2,358评论 9 10
  • Spring Cloud为开发人员提供了快速构建分布式系统中一些常见模式的工具(例如配置管理,服务发现,断路器,智...
    卡卡罗2017阅读 135,010评论 19 139
  • 同样对待孩子。有人说“我要给他/她全世界最好的。”有人说“我要让他/她成为全世界最好的。” 种下什么因,就一定会得...
    sunfeng0912阅读 248评论 0 3
  • 桑克尔还在悲鸣。 我望着灰暗的天空,身上使不出一丝力气,直到雨落在眼睛里,眼睑才挣扎着恢复生机。此刻我该想些什么,...
    不才子阅读 238评论 0 1
  • 都说程序员上辈子是江湖中飞檐走壁的侠客,在江湖中长剑在手,快意恩仇,潇潇洒洒,仗义人生。所以这辈子再也不想过那种漂...
    非著名程序员阅读 1,284评论 4 8