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


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

原文发布于微信公众号 - 智能算法(AI_Algorithm)

原文发表时间:2018-07-13

本文参与腾讯云自媒体分享计划,欢迎正在阅读的你也加入,一起分享。

发表于

我来说两句

0 条评论
登录 后参与评论

相关文章

来自专栏玉树芝兰

如何对PDF文献做可视化分析?

看了大量文献后,你的硬盘上想必存下不少PDF文件。能否充分利用它们,挖掘出你独特的领域知识地图呢?本文为你提供一种简便易行的办法。

18720
来自专栏美图数据技术团队

美图AB Test实践:Meepo系统

感谢阅读「美图数据技术团队」的第 8 篇文章,关注我们持续获取美图最新数据技术动态。

27210
来自专栏吉浦迅科技

假如用多个CPU核同时控制一个GPU,并且进行运算与数据的传递,请问这能够实现么?...

问:假如用多个CPU核同时控制一个GPU,并且进行运算与数据的传递,请问这能够实现么?我的工作站是48个CPU核,两块GPU卡,打算把CPU和GPU都用起来,我...

39770
来自专栏新智元

【自动编译代码】陈天奇团队TVM重磅更新:直接在浏览器使用GPU

【新智元导读】华盛顿大学陈天奇团队的深度学习自动优化代码生成器TVM发布更新,不需要写一行Javascprit代码,直接就能将深度学习模型编译到WebGL,然后...

48450
来自专栏灯塔大数据

分析 | Python抓取婚恋网用户数据,原来这才是年轻人的择偶观

刚好在看决策树这一章,书里面的理论和例子让我觉得这个理论和选择对象简直不能再贴切。看完长相看学历,看完学历看收入。

27830
来自专栏CVer

【重磅】吴恩达最新的机器学习书籍更新啦!

Amusi曾于2018年4月5日推送了重磅:吴恩达最新的机器学习书籍《Machine Learning Yearning》,当时Amusi收到了一封邮件,如下所...

13020
来自专栏华章科技

使用python抓取婚恋网用户数据并用决策树生成自己择偶观

之前在世纪佳缘上爬取过类似的数据,总体的感觉是上面的用户数据要么基本不填要么一看就很假,周围的一些老司机建议可以在花田网上看下,数据质量确实高很多,唯一的缺点就...

9520
来自专栏编程

基于Python库Face_Recognition动手DIY人脸识别

0X00 前言 最近使用iPhone x,把人脸识别代入了我们的生活中。前段时间了解了一个Python的一个开元函数库,并对其进行了分析、学习和实践,那么今天我...

37050
来自专栏目标检测和深度学习

谷歌GPU资源免费啦

注:本文编译自medium,原英文链接:https://medium.com/@nickbourdakos/train-your-machine-learnin...

43170
来自专栏机器之心

专栏 | 百度PaddlePaddle的新特性与大规模稀疏数据分布式模型训练

百度深度学习框架 PaddlePaddle 自 2016 年开源以来,受到了业界的广泛关注,PaddlePaddle 社区更是汇集了一大批 AI 技术开发者。开...

14530

扫码关注云+社区

领取腾讯云代金券