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应用。

最后编辑于
©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 204,189评论 6 478
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 85,577评论 2 381
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 150,857评论 0 337
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 54,703评论 1 276
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 63,705评论 5 366
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 48,620评论 1 281
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 37,995评论 3 396
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 36,656评论 0 258
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 40,898评论 1 298
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 35,639评论 2 321
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 37,720评论 1 330
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 33,395评论 4 319
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 38,982评论 3 307
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 29,953评论 0 19
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 31,195评论 1 260
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 44,907评论 2 349
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 42,472评论 2 342

推荐阅读更多精彩内容

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