分享TensorFlow Lite應(yīng)用案例

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

1、移動端深度學習的技術(shù)選型

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

輸入法引擎預測效果圖

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

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

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

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

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

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

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

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

包括控制流 （control flow） 的 op

相對于 TF mobile，TF Lite 的部分 op 只支持最簡單的 case

目前的一個好的消息就是 TensorFlow 項目組一直在持續(xù)的推進對 RNN 系列的支持。

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

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

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

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

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

組合

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

補充

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

補充的方式有兩種：

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

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

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

模型拆分

1） 原因

需要模型拆分的原因一般有 3 個：

訓練時用流程控制的方式（如 batch）一次性跑完多個樣本，但在 Inference 的過程中，需要用到單步運行；

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

有部分的冗余，但是重新設(shè)計 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， 。..）

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

可能不需要 `BATCH`，雖然可以每次都指定 batch 為 1，但對于 TF 來說，

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

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

對于后端的模型算法工作者來說，寫出上述的訓練代碼是一件非常自然的事情。如果我們既想保持后端代碼的普適和自然度，又想要快速實現(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 個需要被注意的地方：

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

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

另外一個方式是可以考慮采用 `tf.variable_scope（reuse=True）` 的方式重寫 inference 的過程，以解耦 training 和 inference 的代碼，代價就是整個 graph 會偏大，但是優(yōu)點使得進行 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）

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

python

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

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

python

# logits has shape ［1， VOCABS］

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

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

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

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

拆完以后的模型仍然是一個 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

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

第二個是一個包含 toco 的小啟動器，因為 toco 從命令列呼叫起來的話要填的參數(shù)比較多，所以這個啟動器會使用 tensorflow 查詢一些可以自動填的參數(shù)，來降低手動填的參數(shù)數(shù)量。

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

5. 效果分析： TF Lite 帶來的收益

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

TF Lite based model performance metrics

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

6. TensorFlow 與 Kika

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