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], ...) 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 实践中的经验分享。

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