分享TensorFlow Lite应用案例

从2016年开始,kika 技术团队一直致力于 ai 技术在移动端落地,尤其是在 keyboard 输入法引擎做了很多算法与工程上的探索工作。2017 年 5 月,kika 技术团队基于 tensorflow mobile 研发了 kika ai engine,将其应用于 kika 的全系输入法产品中。2017 年 11 月,google 发布 tensorflow lite (tf lite) 后,kika 技术团队迅速进行了跟进,并于 2018 年 1 月成功地开发了基于 tf lite 全新一代的 kika ai engine,同时进行了线上产品的更新。
1、移动端深度学习的技术选型
输入法引擎的技术要求包括:快、准、全。需要在客户端环境下,根据用户输入的上文内容以及当前键入的键码,实时进行『预测』。预测的内容包括:单词,词组,emoji 等等一切可能通过输入法发送的内容。从算法的原理上来讲,这是一个典型的 rnn 应用场景。
输入法引擎预测效果图
作为输入法这样的一个重度使用的工具类 app,在移动端做轻量化部署非常重要,具体包括以下四个方面:模型压缩、快速的响应时间、较低的内存占用以及 较小的 so 库(shared object,共享库)大小等。
在 kika 将 tf mobile 部署到移动端的过程中,除了 cpu 占用偏高,还有由于 tf mobile 内存管理与内存保护设计的问题,导致:
内存保护机制不完善,在实际内存不是很充足的情况(尤其对于部分低端机型以及在内存消耗较大的应用,如大型手游中弹起输入法),容易引发内存非法操作。
内存大小控制机制存在一定的问题,例如模型本身在计算时只有 20mb,但加载到内存之后的运行时峰值可能会飙升 40 到 70mb。
tf lite 对于 cnn 类的应用支持较好,目前对于 rnn 的支持尚存在 op 支持不足的缺点。但是考虑到内存消耗和性能方面的提升,kika 仍然建议投入一部分的研发力量,在移动端考虑采用 tf lite 做为基于 rnn 深度学习模型的 inference 部署方案。
2. tensorflow lite 对 rnn/lstm based 模型的原生支持情况
相对于 cnn 而言,tf lite 对于 rnn/lstm 的支持程度稍显不足。目前的情况是,rnn 相关的基本元素的 op 目前都已经支持,最近也刚刚支持了 lstm,但遗憾的是 beamsearch 支持暂时还没有完成。
不支持的 op 主要集中有两大类情况:
包括控制流 (control flow) 的 op
相对于 tf mobile,tf lite 的部分 op 只支持最简单的 case
目前的一个好的消息就是 tensorflow 项目组一直在持续的推进对 rnn 系列的支持。
3. 如何应对 op 缺失的情况
对于移动端用 tf lite 部署最友好的开发姿势是在设计模型之处就了解当前的 tf lite版本哪些 op 是缺失或者功能不完整的,然后在模型设计过程中:
尽量避免使用这些 tf lite 不支持的 op;
对于不得不使用的情况,也需要结合具体的业务逻辑,优化设计,使得在移动端部署的二次开发的工作量尽可能的小。
以下是应对 op 缺失的一些常见做法。
组合
最为常见的处理方式,例如在早期的 tf lite 版本中,tf.tile 和 tf.range 都不支持,这个时候建议采用 broadcast_add 来组合代替实现。
补充
tf mobile 的 op 相当于完整版的 tensorflow,于此相比,tf lite 缺失最严重的是包含控制流的部分。例如 seq2seq 模型中常用的 beam search。
补充的方式有两种:
直接开发一个全新的 op;
在 tf lite 之外的上层 api 中实现 (此时可能需要拆解模型)。
两种方式各有优劣,具体的需要根据功能的复杂度和业务逻辑决定。
模型拆分
1) 原因
需要模型拆分的原因一般有 3 个:
训练时用流程控制的方式(如 batch)一次性跑完多个样本,但在 inference 的过程中,需要用到单步运行;
某些 op 不支持,需要在 tf lite 的上层『手动』实现,可能需要将原有的模型拆分为若干的子模型 (sub graph);
有部分的冗余,但是重新设计 graph 再训练的时间代价较大。
2) 方法与坑
以下通过一个实例来描述如何进行模型的拆分。
将 variable 共享给不同的 op,甚至于不同的 sub graph,通用做法是 采用 `placeholder` 的方式将输入输出分开,然后在导出 freeze graph 的时候用 `tf.graph_util.convert_variables_to_constants` 只抓取需要的部分。
代码实例:
python
vars = tf.get_variable(。..)
inputs = tf.placeholder(‘inputids’, shape=[batch, none], 。..)
embs = tf.nn.embedding_lookup(vars, inputs)
cells = tf.nn.rnn_cell.multirnncell(。..)
output, state = tf.nn.dynamic_rnn(cells, embs, 。..)
实际整合进入客户端产品 inference 的时候,可能存在的坑:
可能不需要 `batch`,虽然可以每次都指定 batch 为 1,但对于 tf 来说,
batch = 1 跟直接沒有这个维度的模型结构并不同;
如果都需要单步运行的话,`dynamic_rnn` 也不需要,而且这里有大量流程控制 (最新的 tf lite 开始逐步的对 dynamic rnn 进行了支持)。
对于后端的模型算法工作者来说,写出上述的训练代码是一件非常自然的事情。如果我们既想保持后端代码的普适和自然度,又想要快速实现能够在客户端部署,需要作出如下的事情:
python
prod_inputs = tf.placeholder(‘prod_inputids’, shape=[none], 。..)
prod_embs = tf.nn.embedding_lookup(vars, prod_inputs)
prod_output, prod_state = cells(prod_embs, 。..)
其中有 3 个需要被注意的地方:
rnn cell 本身可以被调用。同一个 cell 如果想让多个地方同时调用,內部 variable 只会产生一次。
一般声明的 variables 如果是用 `tf.get_variable()` 出來的,直接用即可。
另外一个方式是可以考虑采用 `tf.variable_scope(reuse=true)` 的方式重写 inference 的过程,以解耦 training 和 inference 的代码,代价就是整个 graph 会偏大,但是优点使得进行 sub graph 切分的工作变得更加简单。
python
with tf.variable_scope(‘my_network’):
vars = tf.get_variable(。..)
inputs = tf.placeholder(‘inputids’, shape=[batch, none], 。..)
embs = tf.nn.embedding_lookup(vars, inputs)
cells = tf.nn.rnn_cell.multirnncell(。..)
output, state = tf.nn.dynamic_rnn(cells, embs, 。..)
# 。..
with tf.variable_scope(‘my_network’, reuse=true):
vars = tf.get_variable(。..)
prod_inputs = tf.placeholder(‘prod_inputids’, shape=[none], 。..)
prod_embs = tf.nn.embedding_lookup(vars, prod_inputs)
prod_cells = tf.nn.rnn_cell.multirnncell(。..)
prod_output, prod_state = prod_cells(prod_embs, 。..)
在进行这些『切分』操作的时候需要注意到几个问题:
1. `tf.variable()` 和 `tf.get_variable()`
尽量用后者,因为`tf.variable()`对 variable scope 无效。
2. 部分 op 有隐藏的 optional argument
有些 op 有 optional argument,如果不指定的话,可能会自动引入一些额外的 op 來代入默认值。这样偶尔会引入一些 tf lite 不支持的 op。例如:
python
softmax = tf.nn.softmax(logits)
其实有个参数 axis 默认是 -1 ,也就是最后一个维度。不写明的话 tf 会『默认』插入一些 op 在运行时帮你计算:
python
axis = tf.sub(tf.shape(logits), tf.constant(1))
`tf.shape()` 在 tf lite 一直到最近才支持,而且只要调用的时候直接写明,并不需要在运行时算:
python
# logits has shape [1, vocabs]
softmax = tf.nn.softmax(logits, axis=1)
这类 op 暂时没有系统性的方式可以辨认 (spec 上没写),只能等到试错的时候才会被发现。
因此,在实际操作的时候对于默认参数,需要特别的注意。
4. toolchain -- 模型转换与整合
拆完以后的模型仍然是一个 protobuffer 格式,要先把它转换成 tflite 的 flatbuffers 格式才能用。
转换工具可以直接采用 tf 官方的转换工具。比如在kika 我们的 toolchain 是这样的:
bash
git clone -b tflite https://github.com/kikatech/tensorflow.git
cd tensorflow/kika
bazel build -s -c dbg
@org_tensorflow//tensorflow/contrib/lite/toco:toco
//graph_tools/python:tf2lite
//graph_tools/python:tfecho
//graph_tools/python:quantize
第一个就是模型转换工具 toco,建议采用独立的命令行版本,而不是采用 python api,目前对于 osx 这样的系统,会有一些编译上的问题,同时编译的耗时也比较长。
第二个是一个包含 toco 的小启动器,因为 toco 从命令列呼叫起来的话要填的参数比较多,所以这个启动器会使用 tensorflow 查询一些可以自动填的参数,来降低手动填的参数数量。
第三个就是量化工具。如果只是要验证 graph 能否在 tf lite 上运行,不需要用到。如果要整合进客户端产品的话,还会经过量化把模型体积压缩后才推送至用户手机 (或打包进安装包),在用户手机上做一次性的还原后才能运行。
5. 效果分析: tf lite 带来的收益
在客户端实现基于 tf lite 模型的部署之后,我们分别测试了同一模型在 tf 完全版(tf mobile)和 tf lite 10, 000 次 inference 的资源消耗情况,如下图所示。主要的 metrics 包括内存占用 (memory),运行时间(speed)和静态链接库的大小 (image size)。
tf lite based model performance metrics
可以看到,各项 metrics 都得到的大幅的优化,这对于提升产品的整体性能与稳定度都是十分有利的。
6. tensorflow 与 kika
除了输入法引擎之外,kika 技术团队近年来也一直在致力于采用 ai 技术解决内容推荐,语音识别和自然语义理解方面等方面的诸多实际问题,在客户端和服务端部署分别采用 tf lite 和 tf serving 这两个基于 tensorflow 的优秀框架。后续 kika 技术团队将持续带来关于 kika 在 tf lite 和 tf serving 实践中的经验分享。

诺基亚通过N系列商标审批 重拾诺基亚N9后续机型?
电容的作用及其特点
HUAWEI Mate RS保时捷设计,独一无二的奢华体验
中国科学院苏州纳米所:研发柔性水伏离子传感器用于可穿戴电子
智能家电电磁铁在智能家电的应用
分享TensorFlow Lite应用案例
贝叶斯分类器原理及应用分析
天马LCD CUP全面屏解决方案为用户带来极致全面屏的视觉冲击
基于FPGA的极化码的SC译码算法结构的改进方法
如何才能轻松有效的解决ESD静电问题
远景维珍车队正式宣布与现任车手Robin Frijns续约
搭载麒麟710处理器的荣耀8X跑分曝光
西数若加入东芝半导体股权收购,可扩大NAND Flash市占
总线控制与运控芯片的结合
如何正确使用热保护器?
化解融资难融资贵和信息不对称,江苏银行试水物联网管控风险
TCL率先占领电视行业的风口趋势
中国移动发布了2020-2021年镀锌钢绞线产品集中采购中标候选人结果
Fluke50II系列测温仪助力“欧洲之星”事故调查
如何使用Domino接口结合VB6工具来开发一个邮件扩展工具