近年來,有道技術團隊在移動端實時 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 方面的技術探索以及實際産品應用。
原文釋出時間為:2018-04-19
本文來自雲栖社群合作夥伴新智元,了解相關資訊可以關注“AI_era”。
原文連結:
玩轉TensorFlow Lite:有道雲筆記實操案例分享