TensorFlow Mobile模型压缩

前言

前文中我们把训练好的模型打包成GraphDef文件(PB文件)了,可是打包出来的文件还是有点大;移动设备的内存容量有限,而且我们需要下载模型到移动端去加载,所以一个大的模型要经过压缩之后才能放在Android或者IOS上跑的,不然会给RAM造成极大的负担,系统会自行kill掉一些进程,甚至你的APP自己crash了。所以模型压缩是十分必要的。

正文

在Google的优化方案中,主要有那么几种优化方式:删除未使用节点,删除training所需的ops,对权重进行四舍五入运算,对BN进行预处理,内存映射,以及终极伤敌1000自损800的方案—转换权重数据类型。本文将逐一介绍。

optimize_for_inference

调用optimize_for_inference脚本会删除输入和输出节点之间所有不需要的节点。同时该脚本还做了一些其他优化以提高运行速度。例如它把显式批处理标准化运算跟卷积权重进行了合并,从而降低了计算量。

使用方法:

bazel build tensorflow/python/tools:optimize_for_inference

bazel-bin/tensorflow/python/tools/optimize_for_inference \
–input=/tf_files/CTNModel.pb \
–output=/tf_files/optimized_graph.pb \
–input_names=inputs/X \
–output_names=output/predict

首先用bazel build出我们的optimize_for_inference脚本,再调用这个脚本,提供几个参数:输入的PB文件路径,输出的PB文件路径,输入节点名以及输出节点名。

经过这一次的优化,文件变小了一些,可是还不足以我们放到手机端去运行,所以我们要进一步的压缩模型,同时还要保证准确率。

我们可以使用之前的在Python端测试PB文件的代码测试优化后的PB文件是否可行。

quantize_graph

在IOS或者Android项目中,人们通常会把PB文件放在assets文件夹中加载,不管是一起打包进APP还是进入APP后再进行下载、解压、复制,在经过optimize_for_inference优化过后的模型依然是非常的大的。IOS在使用.ipa包发布APP的时候,所有内容都会经过zip压缩;Android从网络端下载文件也要经过压缩后解压的过程,所以有没有一种行之有效的方法在不过多的降低精确度的情况下压缩更大的空间呢?Google提供了这么一个脚本,经过这个脚本的PB文件原本的大小不会改变,但会有更多的可利用的重复性,所以压缩成zip包会缩小大约3~4倍的大小。使用方式很简单:

bazel build tensorflow/tools/quantization:quantize_graph

bazel-bin/tensorflow/tools/quantization/quantize_graph \
–input=/tf_files/optimized_graph.pb \
–output=/tf_files/rounded_graph.pb \
–output_node_names=output/predict \
–mode=weights_rounded

输入的参数依然是:输入的PB文件路径,输出的PB文件路径,输出节点名,这里还有个特别的参数mode,这个参数是告诉脚本我们选择哪种压缩方式,这里我们选择了对权重进行四舍五入。

graph_transforms

在2017年的Google I/O大会上,曾提到过的graph_transforms是一整套优化工具,基本上所有的优化方案都能在这里找到,包括之前的那两种优化方式。

Optimizing for Deployment

这个方式类似前面的optimize_for_inference,只是被整合到了transforms中,脚本进行了删除前向传播过程中未调用到的节点,通过预先乘卷积的权重来优化BN中的一些乘法运算。

示例代码:

bazel build tensorflow/tools/graph_transforms:transform_graph

bazel-bin/tensorflow/tools/graph_transforms/transform_graph \
--in_graph=CTNModel.pb \
--out_graph=optimized_graph.pb \
--inputs='inputs/X' \
--outputs='output/predict' \
--transforms='
  strip_unused_nodes(type=float, shape="256*64")
  remove_nodes(op=Identity, op=CheckNumerics)
  fold_constants(ignore_errors=true)
  fold_batch_norms
  fold_old_batch_norms'

transforms的参数代表着想要做的操作,这里主要就是删除未调用的节点,以及优化BN算法,而且必须先运行fold_constants。

Fixing Missing Kernel Errors

由于TensorFlow在Mobile中使用的时候,默认情况下是只能推理预测,可是在build so文件、jar文件或者.a文件的时候,依赖文件是写入在tensorflow / contrib / makefile / tf_op_files.txt中的,里面还包含了一些training相关的ops,这些可能会导致我们在加载PB文件的时候报错—No OpKernel was registered to support Op,这时候我们可以通过这个脚本来修复

实际上,经过上面的脚本,已经删除了这部分节点,不需要重复使用。

示例代码:

bazel build tensorflow/tools/graph_transforms:transform_graph

bazel-bin/tensorflow/tools/graph_transforms/transform_graph \
--in_graph=CTNModel.pb \
--out_graph=fixed_kernel_graph.pb \
--inputs='inputs/X' \
--outputs='output/predict' \
--transforms='
  strip_unused_nodes(type=float, shape="256*64")
  fold_constants(ignore_errors=true)
  fold_batch_norms
  fold_old_batch_norms'

round_weights

上文中提及的quantize_graph,其实在graph_transforms也有,在graph_transforms中我们应该如何使用呢?

示例代码:

bazel build tensorflow/tools/graph_transforms:transform_graph

bazel-bin/tensorflow/tools/graph_transforms/transform_graph \
--in_graph=CTNModel.pb \
--out_graph=rounded_grapg.pb \
--inputs='inputs/X' \
--outputs='output/predict' \
--transforms='
  strip_unused_nodes(type=float, shape="256*64")
  fold_constants(ignore_errors=true)
  fold_batch_norms
  fold_old_batch_norms
  round_weights(num_steps=256)'

Eight-bit

在graph_transforms中有一种更加残暴的文件大小压缩方式,就是把权重的数据类型由32位的float 32转成8位的int,在数据类型层面为模型减少3~4倍的占用大小。示例图如下:

image.png

原本80多M的文件经过这个转换缩减成20M左右,成效还是挺明显的,不过缺点就是相比之前的几种压缩方式,这种压缩方式损失精度较大,非必要情况,建议别使用。

使用方法:

bazel build tensorflow/tools/graph_transforms:transform_graph

bazel-bin/tensorflow/tools/graph_transforms/transform_graph \
--in_graph=CTNModel.pb \
--out_graph=eight_bit_graph.pb \
--inputs='inputs/X' \
--outputs='output/predict' \
--transforms='
  strip_unused_nodes(type=float, shape="256*64")
  fold_constants(ignore_errors=true)
  fold_batch_norms
  fold_old_batch_norms
  quantize_weights'

再次提醒,非极端情况请勿使用!

由上面几种grapg转换功能来讲,graph_transforms确实是一个非常强大的优化脚本,他能够一步到位的优化我们的PB文件,不像之前需要使用多种脚本之后才能得到想要的结果。上文提供了几种常用的脚本代码示例,每次使用只需要选其一即可。

memmapped_format

在前面几种压缩之后,如果文件还是较大,那该怎么办呢,这时候可以采用内存映射的方式,这一种方式是在运行时控制内存占用的一种有效方式,只是使用起来与原本的PB文件调用方式有些不同(暂时找不到Android的资料,只有IOS的,代码详见example)

示例图如下:

image.png

从整个文件读到内存中变成内存的映射,这能大量的节省内存带宽以及占用量。

使用方法:

bazel build tensorflow/contrib/util:convert_graphdef_memmapped_format

bazel-bin/tensorflow/contrib/util/convert_graphdef_memmapped_format \
–in_graph=/tf_files/rounded_graph.pb \
–out_graph=/tf_files/mmapped_graph.pb

经过这个脚本输出的PB文件不能再Python中直接调用,所以我们暂时也无法检测其可行性,需要在IOS端调用。

后记

以上是PB文件常见的压缩方法,在TensorFlow移植Mobile的过程中还有其他的优化方案,比如优化TensorFlow依赖库的大小,对ops进行高级定制,以删除不需要的ops,此方法需要针对不同的model进行不同的定制,较为高端。还有就是对模型进行优化比如尝试减少层数、减少节点数、对卷积运行做优化(点卷积降维(inception)、卷积拆分(MobileNets))这些都是较为高级的Mobile端压缩方式,本文暂时不做介绍了。

欢迎各位指错讨论。

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

推荐阅读更多精彩内容