從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 實踐中的經驗分享。