玩轉TensorFlow Lite：有道雲筆記實操案例分享

新智元推薦

作者：有道技術團隊

編輯：克雷格

【新智元導讀】這一兩年來，在移動端實現實時的人工智慧已經形成了一波潮流。去年，谷歌推出面向移動端和嵌入式的神經網路計算框架TensorFlow Lite，將這股潮流繼續往前推。TensorFlow Lite如何進行操作？本文將介紹TFLite在有道雲筆記中用於文檔識別的實踐過程，以及 TFLite 都有些哪些特性，供大家參考。

近年來，有道技術團隊在移動端實時 AI 能力的研究上，做了很多探索及應用的工作。2017 年 11 月 Google 發布 TensorFlow Lite (TFLlite) 後，有道技術團隊第一時間跟進 TFLite 框架，並很快將其用在了有道雲筆記產品中。

以下是TFLite在有道雲筆記中用於文檔識別的實踐過程。

文檔識別工作的介紹

1. 文檔識別的定義

文檔識別最初是開發有道雲筆記的文檔掃描功能時面對的一個問題。文檔掃描功能希望能在用戶拍攝的照片中，識別出文檔所在的區域，進行拉伸 (比例還原)，識別出其中的文字，最終得到一張乾淨的圖片或是一篇帶有格式的文字版筆記。實現這個功能需要以下這些步驟：

識別文檔區域：將文檔從背景中找出來，確定文檔的四個角；

拉伸文檔區域，還原寬高比：根據文檔四個角的坐標，根據透視原理，計算出文檔原始寬高比，並將文檔區域拉伸還原成矩形；

色彩增強：根據文檔的類型，選擇不同的色彩增強方法，將文檔圖片的色彩變得乾淨清潔；

布局識別：理解文檔圖片的布局，找出文檔的文字部分；

OCR：將圖片形式的「文字」識別成可編碼的文字；

生成筆記：根據文檔圖片的布局，從 OCR 的結果中生成帶有格式的筆記。

文檔識別就是文檔掃描功能的第一步，也是場景最複雜的一個部分

2. 文檔識別在有道 AI 技術矩陣中的角色

有道近年來基於深度神經網路演算法，在自然語言、圖像、語音等媒體數據的處理和理解方面做了一系列工作，產出了基於神經網路的多語言翻譯、OCR（光學字元識別）、語音識別等技術。在這些技術的合力之下，我們的產品有能力讓用戶以他們最自然最舒服的方式去記錄內容，用技術去理解這些內容，並將其統一轉化為文本以待下一步處理。從這個角度來看，我們的各種技術組成了以自然語言為中心，多種媒體形式相互轉換的網路結構。

文檔識別是從圖像轉化為文本的這條轉換鏈上，不起眼卻又不可缺少的一環。有了它的存在，我們可以在茫茫圖海中，準確找到需要處理的文檔，並將其抽取出來進行處理。

3. 文檔識別的演算法簡介

我們的文檔識別演算法基於 FCNN (Fully Convolutional Neural Network) ，這是一種特別的 CNN（卷積神經網路），其特點是對於輸入圖片的每一個像素點，都對應著一個輸出（相對的，普通的 CNN 網路則是每一張輸入圖片對應著一個輸出）。因此，我們可以標記一批包含文檔的圖片，將圖片中文檔邊緣附近的像素標註為正樣本，其他部分標註為副樣本。訓練時，以圖片作為 FCNN 的輸入，將輸出值與標註值作對比得到訓練懲罰，從而進行訓練。關於文檔識別演算法的更多細節，可以參見有道技術團隊的《文檔掃描：深度神經網路在移動端的實踐》這篇文章。

由於演算法的主體是 CNN，因此文檔掃描演算法中主要用到的運算元（Operator）包括卷積層、Depthwise 卷積層、全連接層、池化層、Relu 層這些 CNN 中常用的運算元。

4. 文檔識別與 TensorFlow

能夠訓練和部署 CNN 模型的框架非常多。我們選擇使用 TensorFlow 框架，是基於以下幾方面的考慮的：

TensorFlow 提供的運算元全面且數量眾多，自己創建新的運算元也並不麻煩。在演算法研發的初期會需要嘗試各種不同的模型網路結構，用到各種奇奇怪怪的運算元。此時一個提供全面運算元的框架能夠節省大量的精力；

TensorFlow 能夠較好的覆蓋伺服器端、Android 端、iOS 端等多個平台，並在各個平台上都有完整的運算元支持；

TensorFlow 是一個比較主流的選擇，這意味著當遇到困難時，更容易在互聯網上找到現成的解決辦法。

5. 為什麼想在文檔識別中用 TFLite

在 TFLite 發布之前，有道雲筆記中的文檔識別功能是基於移動端 TensorFlow 庫 (TensorFlow Mobile) 的。當 TFLite 發布後，我們希望遷移到 TFLite 上。促使我們遷移的主要動力是鏈接庫的體積。

經過壓縮後，Android 上的 TensorFlow 動態庫的體積大約是 4.5M 左右。如果希望滿足 Android 平台下的多種處理器架構，可能需要打包 4 個左右的動態庫，加起來體積達到 18M 左右；而 tflite 庫的體積在 600K 左右，即便是打包 4 個平台下的鏈接庫，也只需要佔用 2.5M 左右的體積。這在寸土寸金的移動 App 上，價值是很大的。

TFLite的介紹

1. TFLite 是什麼

TFLite 是 Google I/O 2017 推出的面向移動端和嵌入式的神經網路計算框架，於2017年11月5日發布開發者預覽版本 (developer preview)。相比與 TensorFlow，它有著這樣一些優勢：

輕量級。如上所述，通過 TFLite 生成的鏈接庫體積很小；

沒有太多依賴。TensorFlow Mobile 的編譯依賴於 protobuf 等庫，而 tflite 則不需要大的依賴庫；

可以用上移動端硬體加速。TFLite 可以通過 Android Neural Networks API (NNAPI) 進行硬體加速，只要加速晶元支持 NNAPI，就能夠為 TFLite 加速。不過目前在大多數 Android 手機上，Tflite 還是運行在 CPU 上的。

TensorFlow Lite的架構設計

2. TFLite 的代碼結構

作為 TFLite 的使用者，我們也探索了一下 TFLite 的代碼結構，這裡分享一下。

目前，TFLite 的代碼位於 TensorFlow 工程中 "tensorflow/contrib/lite" 文件夾下。文件夾下有若干頭/源文件和一些子文件夾。

其中，一些比較重要的頭文件有：

model.h:和模型文件相關的一些類和方法。其中 FlatBufferModel 這個類是用來讀取並存儲模型內容的，InterpreterBuilder 則可以解析模型內容；

Interpreter.h:提供了用以推斷的類 Interpreter，這是我們最常打交道的類；

context.h:提供了存儲 Tensors 和一些狀態的 struct TfLiteContext。實際使用時一般會被包裝在 Interpreter 中；

此外，有一些比較重要的子文件夾：

kernels:運算元就是在這裡被定義和實現的。其中 regester.cc 文件定義了哪些運算元被支持，這個是可以自定義的。

downloads:一些第三方的庫，主要包括：

abseil: Google 對 c++ 標準庫的擴展；

eigen: 一個矩陣運算庫；

farmhash: 做 hash 的庫；

flatbuffers: TFLite 所使用的 FlatBuffers 模型格式的庫；

gemmlowp: Google 開源的一個低精度矩陣運算庫；

neon_2_sse: 把 arm 上的 neon 指令映射到相對應的 sse 指令。

java:主要是 Android 平台相關的一些代碼；

nnapi:提供了 nnapi 的調用介面。如果想自己實現 nnapi 可以看一看；

schema:TFLite 所使用的 FlatBuffers 模型格式的具體定義；

toco:protobuf 模型轉換到 FlatBuffers 模型格式的相關代碼。

我們是怎麼用TFLite的？

1. TFLite 的編譯

TFLite 可以運行在 Android 和 iOS 上，官方給出了不同的編譯流程。

在 Android 上，我們可以使用 bazel 構建工具進行編譯。bazel 工具的安裝和配置就不再贅述了，有過 TensorFlow 編譯經驗的同學應該都熟悉。依照官方文檔，bazel 編譯的 target 是 "//tensorflow/contrib/lite/java/demo/app/src/main:TfLiteCameraDemo"，這樣得到的是一個 demo app。如果只想編譯庫文件，可以編譯 "//tensorflow/contrib/lite/java:tensorflowlite" 這個 target，得到的是 libtensorflowlite_jni.so 庫和相應的 java 層介面。

更多細節見官方文檔：

https://github.com/tensorflow/tensorflow/blob/master/tensorflow/docs_src/mobile/tflite/demo_android.md

在 iOS 上，則需要使用 Makefile 編譯。在 mac 平台上運行 build_ios_universal_lib.sh,會編譯生成 tensorflow/contrib/lite/gen/lib/libtensorflow-lite.a 這個庫文件。這是個 fat library，打包了 x86_64, i386, armv7, armv7s, arm64 這些平台上的庫。

更多細節見官方文檔：

https://github.com/tensorflow/tensorflow/blob/master/tensorflow/docs_src/mobile/tflite/demo_ios.md

兩個平台上 TFLite 庫的調用介面也有所不同：Android 上提供了 Java 層的調用介面，而 iOS 上則是 c++ 層的調用介面。

當然，TFLite 的工程結構是比較簡單的，如果你熟悉了 TFLite 的結構，也可以用自己熟悉的編譯工具來編譯 TFLite。

2. 模型轉換

TFLite 不再使用舊的 protobuf 格式（可能是為了減少依賴庫），而是改用 FlatBuffers 。因此需要把訓練好的 protobuf 模型文件轉換成 FlatBuffers 格式。

TensorFlow 官方給出了模型轉化的指導。首先，由於 TFLite 支持的運算元比較少，更不支持訓練相關的運算元，因此需要提前把不需要的運算元從模型中移除，即 Freeze Graph ；接著就可以做模型格式轉換了，使用的工具是 tensorflow toco。這兩個工具也是通過 bazel 編譯得到。

更多細節見官方文檔：

https://github.com/tensorflow/tensorflow/blob/master/tensorflow/docs_src/mobile/tflite/devguide.md

3. 缺失的運算元

TFLite 目前僅提供有限的運算元，主要以 CNN 中使用到的運算元為主，如卷積、池化等。我們的模型是全卷積神經網路，大部分運算元 TFLite 都有提供，但 conv2d_transpose（反向卷積）運算元並沒有被提供。幸運的該運算元出現在網路模型的末端，因此我們可以將反向卷積之前的計算結果取出，自己用 c++ 實現一個反向卷積，從而計算出最終的結果。由於反向卷積的運算量並不大，所以基本沒有影響到運行速度。

如果不巧，你的模型需要但 TFLite 缺少的運算元並非出現在網路的末端，該怎麼辦呢？你可以自定義一個 TFLite 運算元，將其註冊在 TFLite 的 kernels 列表中，這樣編譯得到的 TFLite 庫就可以處理該運算元了。同時，在模型轉換時，還需要加上 --allow_custom_ops 選項，將 TFLite 默認不支持的運算元也保留在模型中。

更多細節見官方文檔：

https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/lite/g3doc/custom_operators.md

TFLite 優缺點

優點：在庫的大小、開發方便程度、跨平台性、性能之間達成一個平衡

作為對比，有道技術團隊選取了一些其他的移動端深度學習框架，分別分析其在「開發方便程度、跨平台性、庫的大小、性能」四個方面的表現：

TensorFlow Mobile，由於和 server 上的 TensorFlow 是同一套代碼，所以可以直接使用 server 上訓練得到的模型，開發非常方便；能支持 Android, iOS, 跨平台性沒問題；如前所述，庫的大小比較大；性能主流。

caffe2，可以比較方便的從 caffe 訓練出的模型轉換到 caffe2 ，但缺少一些運算元, 開發方便程度一般；能支持 Android, iOS，跨平台性沒問題；庫編譯出來比較大，但是是靜態庫可以壓縮；性能主流。

Mental/Accelerate，這兩個都是 iOS 上的框架。比較底層，需要模型轉換&自己寫 inference 代碼，開發比較痛苦；僅支持 iOS；庫是系統自帶，不涉及庫大小問題；速度很快。

CoreML，這個是 WWDC17 發布的 iOS 11 上的框架。有一些模型轉換工具，只涉及通用運算元時開發不算痛苦，涉及自定義運算元時就很難辦了；僅支持 iOS 11 以上；庫是系統自帶，不涉及庫大小問題；速度很快。

最後是 TFLite:

TFLite，其模型可以由 TensorFlow 訓練得到的模型轉換而來，但缺少一些運算元, 開發方便程度一般；能支持 Android, iOS，跨平台性沒問題；庫編譯出來很小；就我們的實驗來看，速度比TensorFlow 快一點。

可以看到，TensorFlow Mobile 開發方便，通用性好，但鏈接庫大，性能主流（其他 server 端神經網路框架的 mobile 版也都有類似的特點）；Mental/Accelerate 這些比較底層的庫速度很快，但不能跨平台，開發比較痛苦；caffe2、TFLite 這類有為移動端優化過的神經網路框架則比較平衡，雖然初時會有運算元不全的問題，但只要背後的團隊不斷支持推進框架的開發，這個問題未來會得到解決。

優點：相對容易擴展

由於 TFLite 的代碼（相對於 TensorFlow）比較簡單，結構比較容易理清，所以可以相對容易的去擴展。如果你想增加一個 TFLite 上沒有而 TensorFlow 上有的運算元，你可以增加一個自定義的類；如果你想增加一個 TensorFlow 上也沒有的運算元，你也可以直接去修改 FlatBuffers 模型文件。

缺點：ops 不夠全面

如前所述，TFLite 目前主要支持 CNN 相關的運算元 ，對其他網路中的運算元還沒有很好的支持。因此，如果你想遷移 rnn 模型到移動端，TFLite 目前是不 OK 的。

不過根據最新的 Google TensorFlow 開發者峰會，Google 和 TensorFlow 社區正在努力增加 ops 的覆蓋面，相信隨著更多開發者的相似需求， 更多的模型會被很好的支持。這也是我們選擇 TensorFlow 這樣的主流社區的原因之一。

缺點：目前還不能支持各種運算晶元

雖然 TFLite 基於 NNAPI，理論上是可以利用上各種運算晶元的，但目前還沒有很多運算晶元支持 NNAPI。期待未來 TFLite 能夠支持更多的運算晶元，畢竟在 CPU 上優化神經網路運行速度是有上限的，用上定製晶元才是新世界的大門。

總結

這一兩年來，在移動端實現實時的人工智慧似乎已經形成了一波潮流。有道技術團隊在移動端 AI 演算法的研究上，也做了諸多嘗試，推出了離線神經網路翻譯 (離線 NMT) 、離線文字識別 (離線 OCR) 以及離線文檔掃描等移動端實時 AI 能力，並在有道詞典、有道翻譯官、有道雲筆記中進行產品化應用。由於目前移動端 AI 尚處在蓬勃發展階段，各種框架、計算平台等都尚不完善。

在這裡，我們以有道雲筆記中的離線文檔識別功能作為實踐案例，看到了 TFLite 作為一個優秀的移動端AI框架，能夠幫助開發者相對輕鬆地在移動端實現常見的神經網路。後續我們也會為大家帶來更多有道技術團隊結合 TFLite 在移動端實時 AI 方面的技術探索以及實際產品應用。

本文授權轉載自TensorFlow（ID：tensorflowers）

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！