Introduction

由於機器學習的蓬勃發展,越來越多的機器學習框架為研究人員或開發者提供便利,例如 TensorFlow[1]、Keras[10] 和 PyTorch[15] 等。然而,沒有 AI 編譯器,各種機器學習框架無法實現端到端的應用。在這種情況下,TVM[3] 應運而生,以彌合機器學習框架、應用程式與後端之間的差距。由於推論需要在行動裝置等邊緣設備上運行,聯發科(MediaTek)推出了 NeuroPilot[4],這是一個跨平台框架,能夠將來自知名機器學習框架的 AI 模型部署至邊緣設備。

然而,NeuroPilot 和 TVM 各有其缺點。TVM 作為一個開源框架,提供通用的前端與後端功能,但 TVM 並不支援所有製造商提供的硬體,例如聯發科的行動裝置 AI 加速器。另一方面,聯發科開發的 NeuroPilot 雖然支援該公司的 AI 加速器,但支援的機器學習框架數量不如 TVM。因此,結合兩者的優勢成為一個極具價值的想法,最終可能實現雙贏的結果。

TVM 的 BYOC(Bring Your Own Codegen)流程為此場景提供了機會。研究人員和開發者可以透過 BYOC 流程,專注於自己熟悉的框架,並通過外部 CodeGen 和 runtime 來連接不同的運行時環境。這對於開發者和硬體後端供應商來說也是有益的,因為它能夠促進社群的廣泛參與。先前的研究已經利用 TVM BYOC 流程,透過 NNAPI 加速器[11] 在行動裝置上使用 NNAPI 進行 AI 推論。此外,另一項相關研究也基於類似概念,嘗試支援 NNAPI 的 NNEF 執行模型[2]。這些研究的共同目標是連接不同的機器學習框架,以促進更好的 AI 發展。

本研究利用 TVM BYOC 來彌合 TVM 和 NeuroPilot 之間的差距。首先,我們計算每個 TVM Relay 運算元(OP)的參數,然後將 Relay 抽象語法樹(AST)轉換為對應的 Neuron IR 類型。我們還將每個 TVM Relay 運算元對應到 Neuron IR,並使用 BYOC 流程為 TVM 和 NeuroPilot 生成運行時環境。透過這種方式,我們的解決方案流程能夠接受來自 TVM 前端的更多機器學習框架,並利用 NeuroPilot 提供的後端目標。

為了進一步展示我們的方法,我們構建了一個應用展示(圖 1)。該應用包含三個不同的機器學習模型,它們執行不同的任務並來自不同的機器學習框架。第一個模型是一個來自 TFLite 的 Mobilenet 量化模型,負責執行物件偵測,並在影片的每一幀中標記出每個物件。此外,我們使用了一個防偽模型[6] 來區分真假人臉。該模型基於卷積神經網路(CNN)框架,用於對抗演示攻擊,並具有深度像素級監督機制。最後,我們應用了一個情緒偵測模型,以提取被選取人臉的情緒。TVM BYOC 流程讓我們能夠整合來自不同機器學習框架的模型,使得我們的應用展示具有獨特性。

在我們的應用展示中,這三個模型之間存在依賴關係。也就是說,只有那些被物件偵測模型捕捉到並且經過人臉防偽模型判定為真實人臉的目標,才會進一步進行情緒偵測。由於聯發科 NeuroPilot 提供了多種後端選項,包括行動 CPU、GPU 及 AI 加速器,我們開發了一個早期的管線原型,以提升應用效能。

Background

Mediatek NeuroPilot

近年來,許多 AI 應用被廣泛應用於行動裝置和智慧車輛,關於 AI 在雲端與邊緣運算上的應用也受到廣泛討論。因此,許多公司紛紛開發自家的平台,以在邊緣設備上執行 AI 推論。例如,Intel 提供了一個端到端的邊緣運算 AI 框架 OpenVINO[7],而 MediaTek 也推出了 NeuroPilot,作為其專門用於行動裝置 AI 推論的平台。

儘管邊緣運算的重要性已廣為人知,但仍有許多挑戰需要克服。其中大多數挑戰來自於物理限制,例如功耗和散熱問題。因此,NeuroPilot 依賴於 MediaTek 的硬體來解決這些問題。然而,TVM 並不容易支援此類硬體,因此我們希望在兩者之間搭建橋樑,以同時發揮兩者的優勢。

NeuroPilot 的兩個核心概念與我們的研究直接相關:RuntimeCompiler。在編譯器部分,NeuroPilot 提供高層級的 IR(中介表示),能夠接受多種機器學習框架,並映射各種 AI 運算。此外,編譯器還提供 Execution Planner(執行計畫機制),用於將運算子分派至不同的後端目標。當編譯器完成工作後,Runtime 負責推理並產生最終的輸出二進位檔案。

TVM

TVM 是一個開源的端到端機器學習編譯框架,其運作方式類似於傳統的編譯器,例如 LLVM[12],並由三個主要階段組成:前端(Front-end)、中介表示(Intermediate Representation, IR)和後端(Back-end)。由於 TVM 具有易用性並能夠方便地執行端到端的 AI 解決方案,因此其受歡迎程度不斷提升。此外,一些研究已基於 TVM 開發工具,例如 NNBlocks[5],來輔助相關專案的開發。

TVM 與 NeuroPilot 之間的一個主要區別在於 TVM 的前端能夠接受多種機器學習框架,包括 TensorFlow[1]、Keras[10]、PyTorch[15],以及開放標準 ONNX[14]。在 TVM 中,有兩種中介表示(IR):Relay IR[17] 和 TIR

Relay IR 是一種高階中介表示(IR),用於描述整個機器學習模型,包括模型將使用的所有 AI 運算,但不會詳細說明每個運算的具體實作方式。而 TIR 則負責描述所有 AI 運算的實作細節。因此,Relay IR 是一個適合進行最佳化的層級,因為我們可以專注於數據流,而無需關心底層細節。在 TVM 的 Relay IR 上,我們可以執行節點融合(Node Fusion)、節點重排(Node Permutation),甚至對節點進行分區(Partition)。

在本研究中,我們使用 TVM BYOC(Bring Your Own Codegen)流程來整合 TVMNeuroPilot。在 TVM 方面,我們利用 TVM Relay IR 進行 圖分區(Graph Partition)與轉換(Transformation),以實現與 NeuroPilot 的互通。

NeuroPilot 支援 TVM

3.1 在 TVM 中使用 BYOC 流程

如前所述,TVM 是一個開源的 AI 編譯框架,支援主流的前端機器學習框架以及各種後端目標。然而,當硬體製造商希望推出自己的硬體目標(如機器學習加速器)時,該硬體供應商可能無法提供與現有機器學習框架(如 TensorFlow、PyTorch 或 MXNet)相容的程式介面,使其難以直接整合到現有的 AI 生態系統中。