AIoT 深度學習模型的編譯與 Edge 部署

Kevin Huang
6 min readFeb 21, 2021

--

深度學習 (Deep Learning) 的熱潮

最近這幾年深度學習的熱潮引發了工程師和學生學習這項技術的熱潮,學習大概都是由深度學習理論與使用前端深度學習的套件,例如 TensorFlow / Caffe / PyTorch 或 Keras 這些套件來練習深度學習相關的算法與應用。也由於這些好用的前端與後端工具,讓深度學習這個領域蓬勃發展。(”深度學習” 這四個字打請來很累,後文用 “DL” 來代表)。

在去年 (2020 年) 開始,各式各樣內建 DL 的晶片陸續推出適合各個運用場景的深度學習的硬體算力引擎,AI 技術往 Edge 部署的方向非常明確。我們用特大算力的 GPU 電腦,或是在雲端用 Microsoft / AWS / Google 的超大雲算力訓練出我們的 DL 模型,要怎麼部署的特定 IC 硬體架構的 Edge 上,變成一個苦惱的工程問題。

問題來了 — 深度學習模型如何部署到 Edge 端 ?

平常 AI 工程師使用的這些的 DL 工具大部分僅支援 CPU 或 GPU (尤其只支援 nVidia GPU 居多),而面對這麼多新的各式各樣的深度學習硬體,架構橫跨 CPU / GPU / DSP / FPGA / TPU / 還有新創 AI 公司的自定義硬體… 那麼 AI 工程師辛苦訓練好的模型要怎麼部署到這些硬體上面呢? 先別說部署了,模型在工具間轉換,比如從 TensorFlow 轉移到 PyTorch 就要花很多時間了… 更遑論要部署到這些不同種類的硬體上。

部署困難變成 AIoT 概念要落地的時候一個主要的工程問題。

深度學習模型的交換格式

在起初,深度學習的模型所使用的框架,每一種都有自己的格式,這裏再次舉 TensorFlow 為例,它自己就有多種不同格式:

  1. Checkpoint → 用於儲存模型的權重,主要用於模型訓練過程中引數的備份和模型訓練熱啟動
  2. GraphDef → 用於儲存模型的Graph,不包含模型權重,加上checkpoint後就有模型上線的全部資訊。
  3. SavedModel / FrozenGraph / TFLite: 使用在產品環境 (Production Environment) 服務上面的格式,使用 saved_model 介面匯出的模型檔案,包含模型 Graph 和許可權可直接用於上線,TensorFlow 和 Keras 模型推薦使用 SavedModel 模型格式。FrozenGraph / TFLite 格式可以直接部署到 Android/ iOS。

前面所提的兩種 TensorFlow 格式 (CheckPoint 和 GraphDef) 主要是訓練深度學習模型的時候用的,而關於部署到不同的硬體,我們主要關心 SavedMode / FrozenGraph / TFLite 這三種格式。更細緻地說,SavedMode 是一個把權重 (Weight) 與連接圖 (Graph) 都記錄下來的格式,而 FrozenGraph / TFLite 是針對部署再優化過的格式。這感覺有點像 “gcc -g” 與 “gcc -O2” 的差別? 關於 TensorFlow 的格式,這裡有一篇文章寫得很詳細。這樣的格式要給別的 DL 框架使用,比如 PyTorch / Caffe2 或是 MXNet,就無法直接讀取了。

為了解決框架之間無法交換深度學習模型的問題,業界推動了一個通用的交換格式,稱為 ONNX(Open Neural Network Exchange)。ONNX 是一種針對機器學習的文件格式,它是一個開放的格式,任何人都可以使用它,用來儲存訓練好的深度學習模型。ONNX 讓使用不同前後端工具與框架 (TensorFlow / Pytorch / MXNet …) 所訓練出來的模型可以互通。ONNX 已開放 Source Code 的方式放在 Github 上面 (連結在此)。

深度學習模型與編譯器 (Compiler)

有了共通的交換格式之後,下一個問題就是如何將深度學習 ONNX 模型部署到不同架構的硬體上面呢?這個時候你就需要 DL 模型編譯器來幫忙了。DL 編譯器的兩個重要的性能是

  1. 連接圖型優化 — Graph Optimization
  2. 運算優化 — Operation Optimization

優化是將深度學習模型部署到不同的硬體上面是一個很重要的問題,這和 CPU 上跑的程序需要厲害的 Compiler 來優化是一樣的道理。而與程式語言編譯器一樣,DL 編譯器也有很多中間過程,以下做個名詞解釋。

  1. 高階中間表示 (High-level Intermediate Representation, High-level IR): 這裡是和硬體無關的中間形式,也被為 Graph IR,顧名思義,這個中間形式是要建立出資料流與操作資料中間的控制流程,編譯器就可以針對這個中間表現層在語意和流程層次上加以優化,這是 DL 模型優化的第一步。
  2. 低階中間表示 (Low-level Intermediate Representation, Low-level IR): 這裡定義的中間形式就和硬體的特性相關了,透過這個中間形式,就可以針對部署的硬體架構做優化,這個部分也會採取一些第三方的表現形式,例如 LLVM 的優化中間表示形式。
  3. 前端 (Frontend): 前端的工作在於讀取 DL 框架所產生的格式(例如,ONNX) 然後轉換成 High-level IR,也就是上面說的Graph IR。前端處理器要能了解並且專換 DL 框架的格式,然後分別進行節點層次 (node-level),區塊層次 (block-level) 和資料流層次 (dataflow-level) 的優化,至三種優化都需要一些篇幅來解釋,要再寫另外一篇文章討論。
  4. 後端 (Backend): 後端的工作在於接收前端與硬體無關的優化結果,並且根據要部署硬體的特性做低階優化,跟硬體相關的優化包含直接翻譯成機器代碼 (intrinsic mapping),記憶體配置與抓取 (memory allocation / fetching) 優化,優化排程的記憶體延遲隱藏 (memory latency hiding) 技術,計算平行化(parallelization),迴圈導向的優化( loop oriented optimizations)… 以及自動排程(auto-scheduling),自動調整 (auto-tuning) 技術…解說這些優化技術應該也是要再寫另外一篇文章討論。

編譯器就是由這些不同的優化策略所組成,他們之間的關係,就如同下圖。輸入就是 AI 工程師訓練出來的 DL 模型,轉換成 ONNX,進入編譯器編譯並且優化成指定硬體架構的目的碼,然後就能部署到指定的硬體上了。

DL 編譯器

提供 DL 能力的晶片將來將會變得很普遍,針對每個特殊的硬體設計,工程師們還是需要熟悉 DL 模型的編譯形式和特性,才能夠無痛部署 AI 算法到適當的硬體上從手機到家用 IoT 設備,AIoT 時代,編譯 DL 算法並了解硬體特性的能力,將會和現在編譯 C 語言並了解 CPU 處理能力一樣重要。

下一篇再來解釋 DL 編譯器內部的細節,各個 OpenSouce DL 編譯器的比較與實際操作 DL 編譯與部署。

--

--