TensorFlow Lite在Kika Keyboard中的应用案例分享

文 / Kika 技术团队

Kika keyboard 与 TensorFlow Lite

业务背景

『基于 AI 技术变革沟通,让世界沟通更简单』一直是 Kika keyboard 最重要的使命。从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,同时进行了线上产品的更新。

移动端深度学习的技术选型

输入法引擎的技术要求包括:快、准、全。需要在客户端环境下,根据用户输入的上文内容以及当前键入的键码,实时进行『预测』。预测的内容包括:单词,词组,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], ...)

实际整合进入客户端产品 inference 的时候,可能存在的坑:

可能不需要 `BATCH`,虽然可以每次都指定 batch 为 1,但对于 TF 来说,

batch = 1 跟直接沒有这个维度的模型结构并不同;

如果都需要单步运行的话,`dynamic_rnn` 也不需要,而且这里有大量流程控制 (最新的TF Lite开始逐步的对 dynamic rnn 进行了支持)。

对于后端的模型算法工作者来说,写出上述的训练代码是一件非常自然的事情。如果我们既想保持后端代码的普适和自然度,又想要快速实现能够在客户端部署,需要作出如下的事情:

python

prod_inputs = tf.placeholder('prod_inputids', shape=[None], ...)

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], ...)

# ...

with tf.variable_scope('my_network', reuse=True):

vars = tf.get_variable(...)

prod_inputs = tf.placeholder('prod_inputids', shape=[None], ...)

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

其实有个参数 axis 默认是 -1 ,也就是最后一个维度。不写明的话 TF 会『默认』插入一些 op 在运行时帮你计算:

python

axis = tf.sub(tf.shape(logits), tf.constant(1))

`tf.shape()` 在TF Lite一直到最近才支持,而且只要调用的时候直接写明,并不需要在运行时算:

python

# logits has shape [1, VOCABS]

这类 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 Lite10, 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 实践中的经验分享。

声明:本文系网络转载,版权归原作者所有。如涉及版权,请联系删除!

  • 发表于:
  • 原文链接https://kuaibao.qq.com/s/20180713B0G2MN00?refer=cp_1026
  • 腾讯「云+社区」是腾讯内容开放平台帐号(企鹅号)传播渠道之一,根据《腾讯内容开放平台服务协议》转载发布内容。
  • 如有侵权,请联系 yunjia_community@tencent.com 删除。

扫码关注云+社区

领取腾讯云代金券