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