天天看點

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

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

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

文檔識别工作的介紹

1. 文檔識别的定義

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

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

識别文檔區域: 将文檔從背景中找出來,确定文檔的四個角;

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

拉伸文檔區域,還原寬高比: 根據文檔四個角的坐标,根據透視原理,計算出文檔原始寬高比,并将文檔區域拉伸還原成矩形;

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

色彩增強: 根據文檔的類型,選擇不同的色彩增強方法,将文檔圖檔的色彩變得幹淨清潔;

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

布局識别: 了解文檔圖檔的布局,找出文檔的文字部分;

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

OCR: 将圖檔形式的“文字”識别成可編碼的文字;

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

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

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

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

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

有道近年來基于深度神經網絡算法,在自然語言、圖像、語音等媒體資料的處理和了解方面做了一系列工作,産出了基于神經網絡的多語言翻譯、OCR(光學字元識别)、語音識别等技術。在這些技術的合力之下,我們的産品有能力讓使用者以他們最自然最舒服的方式去記錄内容,用技術去了解這些内容,并将其統一轉化為文本以待下一步處理。從這個角度來看,我們的各種技術組成了以自然語言為中心,多種媒體形式互相轉換的網絡結構。

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

3. 文檔識别的算法簡介

我們的文檔識别算法基于 FCNN (Fully Convolutional Neural Network) ,這是一種特别的 CNN(卷積神經網絡),其特點是對于輸入圖檔的每一個像素點,都對應着一個輸出(相對的,普通的 CNN 網絡則是每一張輸入圖檔對應着一個輸出)。是以,我們可以标記一批包含文檔的圖檔,将圖檔中文檔邊緣附近的像素标注為正樣本,其他部分标注為副樣本。訓練時,以圖檔作為 FCNN 的輸入,将輸出值與标注值作對比得到訓練懲罰,進而進行訓練。關于文檔識别算法的更多細節,可以參見有道技術團隊的《

文檔掃描:深度神經網絡在移動端的實踐

》這篇文章。

由于算法的主體是 CNN,是以文檔掃描算法中主要用到的算子(Operator)包括卷積層、Depthwise 卷積層、全連接配接層、池化層、Relu 層這些 CNN 中常用的算子。

4. 文檔識别與 TensorFlow

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

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

TensorFlow 提供的算子全面且數量衆多,自己建立新的算子也并不麻煩。在算法研發的初期會需要嘗試各種不同的模型網絡結構,用到各種奇奇怪怪的算子。此時一個提供全面算子的架構能夠節省大量的精力;

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

TensorFlow 能夠較好的覆寫伺服器端、Android 端、iOS 端等多個平台,并在各個平台上都有完整的算子支援;

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

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,它有着這樣一些優勢:

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

輕量級。如上所述,通過 TFLite 生成的連結庫體積很小;

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

沒有太多依賴。TensorFlow Mobile 的編譯依賴于 protobuf 等庫,而 tflite 則不需要大的依賴庫;

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

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

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

TensorFlow Lite的架構設計

2. TFLite 的代碼結構

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

目前,TFLite 的代碼位于 TensorFlow 工程中 "tensorflow/contrib/lite" 檔案夾下。檔案夾下有若幹頭/源檔案和一些子檔案夾。

其中,一些比較重要的頭檔案有:

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

model.h: 和模型檔案相關的一些類和方法。其中 FlatBufferModel 這個類是用來讀取并存儲模型内容的,InterpreterBuilder 則可以解析模型内容;

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

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

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

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

此外,有一些比較重要的子檔案夾:

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

kernels: 算子就是在這裡被定義和實作的。其中 regester.cc 檔案定義了哪些算子被支援,這個是可以自定義的。

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

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

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

abseil: Google 對 c++ 标準庫的擴充;

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

eigen: 一個矩陣運算庫;

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

farmhash: 做 hash 的庫;

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

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

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

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

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

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

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

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

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

nnapi: 提供了 nnapi 的調用接口。如果想自己實作 nnapi 可以看一看;

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

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

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

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 Lite:有道雲筆記實操案例分享

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

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

caffe2,可以比較友善的從 caffe 訓練出的模型轉換到 caffe2 ,但缺少一些算子, 開發友善程度一般;能支援 Android, iOS,跨平台性沒問題;庫編譯出來比較大,但是是靜态庫可以壓縮;性能主流。

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

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

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

CoreML,這個是 WWDC17 釋出的 iOS 11 上的架構。有一些模型轉換工具,隻涉及通用算子時開發不算痛苦,涉及自定義算子時就很難辦了;僅支援 iOS 11 以上;庫是系統自帶,不涉及庫大小問題;速度很快。

最後是 TFLite:

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

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:有道雲筆記實操案例分享

繼續閱讀