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], ...)
實際整合進入客戶端產品 inference 的時候，可能存在的坑：
可能不需要 `BATCH`，雖然可以每次都指定 batch 為 1，但對於 TF 來說，
batch = 1 跟直接沒有這個維度的模型結構並不同；
如果都需要單步運行的話，`dynamic_rnn` 也不需要，而且這裡有大量流程式控制制 （最新的TF Lite開始逐步的對 dynamic rnn 進行了支持）。
對於後端的模型演算法工作者來說，寫出上述的訓練代碼是一件非常自然的事情。如果我們既想保持後端代碼的普適和自然度，又想要快速實現能夠在客戶端部署，需要作出如下的事情：
python
prod_inputs = tf.placeholder("prod_inputids", shape=[None], ...)
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], ...)
# ...
with tf.variable_scope("my_network", reuse=True):
vars = tf.get_variable(...)
prod_inputs = tf.placeholder("prod_inputids", shape=[None], ...)
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
其實有個參數 axis 默認是 -1 ，也就是最後一個維度。不寫明的話 TF 會『默認』插入一些 op 在運行時幫你計算：
python
axis = tf.sub(tf.shape(logits), tf.constant(1))
`tf.shape()` 在TF Lite一直到最近才支持，而且只要調用的時候直接寫明，並不需要在運行時算：
python
# logits has shape [1, VOCABS]
這類 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 Lite10, 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 實踐中的經驗分享。
聲明：本文系網路轉載，版權歸原作者所有。如涉及版權，請聯繫刪除！
TAG:智能演算法 |