從2016年開始，Kika 技術(shù)團(tuán)隊(duì)一直致力于 AI 技術(shù)在移動(dòng)端落地，尤其是在 keyboard 輸入法引擎做了很多算法與工程上的探索工作。2017 年 5 月，Kika 技術(shù)團(tuán)隊(duì)基于 TensorFlow Mobile 研發(fā)了 Kika AI Engine，將其應(yīng)用于 Kika 的全系輸入法產(chǎn)品中。2017 年 11 月，Google 發(fā)布 TensorFlow Lite （TF Lite） 后，Kika 技術(shù)團(tuán)隊(duì)迅速進(jìn)行了跟進(jìn)，并于 2018 年 1 月成功地開發(fā)了基于 TF Lite 全新一代的 Kika AI Engine，同時(shí)進(jìn)行了線上產(chǎn)品的更新。

1、移動(dòng)端深度學(xué)習(xí)的技術(shù)選型

輸入法引擎的技術(shù)要求包括：快、準(zhǔn)、全。需要在客戶端環(huán)境下，根據(jù)用戶輸入的上文內(nèi)容以及當(dāng)前鍵入的鍵碼，實(shí)時(shí)進(jìn)行『預(yù)測(cè)』。預(yù)測(cè)的內(nèi)容包括：?jiǎn)卧~，詞組，emoji 等等一切可能通過(guò)輸入法發(fā)送的內(nèi)容。從算法的原理上來(lái)講，這是一個(gè)典型的 RNN 應(yīng)用場(chǎng)景。

輸入法引擎預(yù)測(cè)效果圖

作為輸入法這樣的一個(gè)重度使用的工具類 APP，在移動(dòng)端做輕量化部署非常重要，具體包括以下四個(gè)方面：模型壓縮、快速的響應(yīng)時(shí)間、較低的內(nèi)存占用以及 較小的 so 庫(kù)（shared object，共享庫(kù)）大小等。

在 Kika 將 TF Mobile 部署到移動(dòng)端的過(guò)程中，除了 CPU 占用偏高，還有由于 TF Mobile 內(nèi)存管理與內(nèi)存保護(hù)設(shè)計(jì)的問(wèn)題，導(dǎo)致：

內(nèi)存保護(hù)機(jī)制不完善，在實(shí)際內(nèi)存不是很充足的情況（尤其對(duì)于部分低端機(jī)型以及在內(nèi)存消耗較大的應(yīng)用，如大型手游中彈起輸入法），容易引發(fā)內(nèi)存非法操作。

內(nèi)存大小控制機(jī)制存在一定的問(wèn)題，例如模型本身在計(jì)算時(shí)只有 20MB，但加載到內(nèi)存之后的運(yùn)行時(shí)峰值可能會(huì)飆升 40 到 70MB。

TF Lite 對(duì)于 CNN 類的應(yīng)用支持較好，目前對(duì)于 RNN 的支持尚存在 op 支持不足的缺點(diǎn)。但是考慮到內(nèi)存消耗和性能方面的提升，Kika 仍然建議投入一部分的研發(fā)力量，在移動(dòng)端考慮采用 TF Lite 做為基于 RNN 深度學(xué)習(xí)模型的 inference 部署方案。

2. TensorFlow Lite 對(duì) RNN/LSTM based 模型的原生支持情況

相對(duì)于 CNN 而言，TF Lite 對(duì)于 RNN/LSTM 的支持程度稍顯不足。目前的情況是，RNN 相關(guān)的基本元素的 op 目前都已經(jīng)支持，最近也剛剛支持了 LSTM，但遺憾的是 beamSearch 支持暫時(shí)還沒(méi)有完成。

不支持的 op 主要集中有兩大類情況：

包括控制流 （control flow） 的 op

相對(duì)于 TF mobile，TF Lite 的部分 op 只支持最簡(jiǎn)單的 case

目前的一個(gè)好的消息就是 TensorFlow 項(xiàng)目組一直在持續(xù)的推進(jìn)對(duì) RNN 系列的支持。

3. 如何應(yīng)對(duì) op 缺失的情況

對(duì)于移動(dòng)端用 TF Lite 部署最友好的開發(fā)姿勢(shì)是在設(shè)計(jì)模型之處就了解當(dāng)前的 TF Lite版本哪些 op 是缺失或者功能不完整的，然后在模型設(shè)計(jì)過(guò)程中：

盡量避免使用這些 TF Lite 不支持的 op；

對(duì)于不得不使用的情況，也需要結(jié)合具體的業(yè)務(wù)邏輯，優(yōu)化設(shè)計(jì)，使得在移動(dòng)端部署的二次開發(fā)的工作量盡可能的小。

以下是應(yīng)對(duì) op 缺失的一些常見做法。

組合

最為常見的處理方式，例如在早期的 TF Lite 版本中，tf.tile 和 tf.range 都不支持，這個(gè)時(shí)候建議采用 broadcast_add 來(lái)組合代替實(shí)現(xiàn)。

補(bǔ)充

TF mobile 的 op 相當(dāng)于完整版的 TensorFlow，于此相比，TF Lite 缺失最嚴(yán)重的是包含控制流的部分。例如 seq2seq 模型中常用的 beam search。

補(bǔ)充的方式有兩種：

直接開發(fā)一個(gè)全新的 op；

在 TF Lite 之外的上層 api 中實(shí)現(xiàn) （此時(shí)可能需要拆解模型）。

兩種方式各有優(yōu)劣，具體的需要根據(jù)功能的復(fù)雜度和業(yè)務(wù)邏輯決定。

模型拆分

1） 原因

需要模型拆分的原因一般有 3 個(gè)：

訓(xùn)練時(shí)用流程控制的方式（如 batch）一次性跑完多個(gè)樣本，但在 Inference 的過(guò)程中，需要用到單步運(yùn)行；

某些 op 不支持，需要在 TF Lite 的上層『手動(dòng)』實(shí)現(xiàn)，可能需要將原有的模型拆分為若干的子模型 （sub graph）；

有部分的冗余，但是重新設(shè)計(jì) graph 再訓(xùn)練的時(shí)間代價(jià)較大。

2） 方法與坑

以下通過(guò)一個(gè)實(shí)例來(lái)描述如何進(jìn)行模型的拆分。

將 variable 共享給不同的 op，甚至于不同的 sub graph，通用做法是 采用 `placeholder` 的方式將輸入輸出分開，然后在導(dǎo)出 freeze graph 的時(shí)候用 `tf.graph_util.convert_variables_to_constants` 只抓取需要的部分。

代碼實(shí)例：

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， 。..）

實(shí)際整合進(jìn)入客戶端產(chǎn)品 inference 的時(shí)候，可能存在的坑：

可能不需要 `BATCH`，雖然可以每次都指定 batch 為 1，但對(duì)于 TF 來(lái)說(shuō)，

batch = 1 跟直接沒(méi)有這個(gè)維度的模型結(jié)構(gòu)并不同；

如果都需要單步運(yùn)行的話，`dynamic_rnn` 也不需要，而且這里有大量流程控制 （最新的 TF Lite 開始逐步的對(duì) dynamic rnn 進(jìn)行了支持）。

對(duì)于后端的模型算法工作者來(lái)說(shuō)，寫出上述的訓(xùn)練代碼是一件非常自然的事情。如果我們既想保持后端代碼的普適和自然度，又想要快速實(shí)現(xiàn)能夠在客戶端部署，需要作出如下的事情：

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 個(gè)需要被注意的地方：

RNN cell 本身可以被調(diào)用。同一個(gè) cell 如果想讓多個(gè)地方同時(shí)調(diào)用，內(nèi)部 variable 只會(huì)產(chǎn)生一次。

一般聲明的 variables 如果是用 `tf.get_variable（）` 出來(lái)的，直接用即可。

另外一個(gè)方式是可以考慮采用 `tf.variable_scope（reuse=True）` 的方式重寫 inference 的過(guò)程，以解耦 training 和 inference 的代碼，代價(jià)就是整個(gè) graph 會(huì)偏大，但是優(yōu)點(diǎn)使得進(jìn)行 sub graph 切分的工作變得更加簡(jiǎn)單。

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， 。..）

在進(jìn)行這些『切分』操作的時(shí)候需要注意到幾個(gè)問(wèn)題：

1. `tf.Variable（）` 和 `tf.get_variable（）`

盡量用后者，因?yàn)閌tf.Variable（）`對(duì) variable scope 無(wú)效。

2. 部分 op 有隱藏的 optional argument

有些 op 有 optional argument，如果不指定的話，可能會(huì)自動(dòng)引入一些額外的 op 來(lái)代入默認(rèn)值。這樣偶爾會(huì)引入一些 TF Lite 不支持的 op。例如：

python

softmax = tf.nn.softmax（logits）

其實(shí)有個(gè)參數(shù) axis 默認(rèn)是 -1 ，也就是最后一個(gè)維度。不寫明的話 TF 會(huì)『默認(rèn)』插入一些 op 在運(yùn)行時(shí)幫你計(jì)算：

python

axis = tf.sub（tf.shape（logits）， tf.constant（1））

`tf.shape（）` 在 TF Lite 一直到最近才支持，而且只要調(diào)用的時(shí)候直接寫明，并不需要在運(yùn)行時(shí)算：

python

# logits has shape ［1， VOCABS］

softmax = tf.nn.softmax（logits， axis=1）

這類 op 暫時(shí)沒(méi)有系統(tǒng)性的方式可以辨認(rèn) （spec 上沒(méi)寫），只能等到試錯(cuò)的時(shí)候才會(huì)被發(fā)現(xiàn)。

因此，在實(shí)際操作的時(shí)候?qū)τ谀J(rèn)參數(shù)，需要特別的注意。

4. toolchain -- 模型轉(zhuǎn)換與整合

拆完以后的模型仍然是一個(gè) protobuffer 格式，要先把它轉(zhuǎn)換成 tflite 的 flatbuffers 格式才能用。

轉(zhuǎn)換工具可以直接采用 TF 官方的轉(zhuǎn)換工具。比如在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

第一個(gè)就是模型轉(zhuǎn)換工具 toco，建議采用獨(dú)立的命令行版本，而不是采用 python API，目前對(duì)于 OSX 這樣的系統(tǒng)，會(huì)有一些編譯上的問(wèn)題，同時(shí)編譯的耗時(shí)也比較長(zhǎng)。

第二個(gè)是一個(gè)包含 toco 的小啟動(dòng)器，因?yàn)?toco 從命令列呼叫起來(lái)的話要填的參數(shù)比較多，所以這個(gè)啟動(dòng)器會(huì)使用 tensorflow 查詢一些可以自動(dòng)填的參數(shù)，來(lái)降低手動(dòng)填的參數(shù)數(shù)量。

第三個(gè)就是量化工具。如果只是要驗(yàn)證 graph 能否在 TF Lite 上運(yùn)行，不需要用到。如果要整合進(jìn)客戶端產(chǎn)品的話，還會(huì)經(jīng)過(guò)量化把模型體積壓縮后才推送至用戶手機(jī) （或打包進(jìn)安裝包），在用戶手機(jī)上做一次性的還原后才能運(yùn)行。

5. 效果分析： TF Lite 帶來(lái)的收益

在客戶端實(shí)現(xiàn)基于 TF Lite 模型的部署之后，我們分別測(cè)試了同一模型在 TF 完全版（TF Mobile）和 TF Lite 10， 000 次 Inference 的資源消耗情況，如下圖所示。主要的 Metrics 包括內(nèi)存占用 （memory），運(yùn)行時(shí)間（speed）和靜態(tài)鏈接庫(kù)的大小 （image size）。

TF Lite based model performance metrics

可以看到，各項(xiàng) Metrics 都得到的大幅的優(yōu)化，這對(duì)于提升產(chǎn)品的整體性能與穩(wěn)定度都是十分有利的。

6. TensorFlow 與 Kika

除了輸入法引擎之外，Kika 技術(shù)團(tuán)隊(duì)近年來(lái)也一直在致力于采用 AI 技術(shù)解決內(nèi)容推薦，語(yǔ)音識(shí)別和自然語(yǔ)義理解方面等方面的諸多實(shí)際問(wèn)題，在客戶端和服務(wù)端部署分別采用 TF Lite 和 TF Serving 這兩個(gè)基于 TensorFlow 的優(yōu)秀框架。后續(xù) Kika 技術(shù)團(tuán)隊(duì)將持續(xù)帶來(lái)關(guān)于 Kika 在 TF Lite 和 TF Serving 實(shí)踐中的經(jīng)驗(yàn)分享。