AWS李沐：一種開源、端到端的新型人工智慧框架編譯器——NNVM編譯器

你可以選擇多種人工智慧(AI)框架來開發人工智慧演算法，還可以選擇多種硬體來訓練和部署人工智慧模型。框架和硬體的多樣性對於維護人工智慧生態系統的健康至關重要。然而，這種多樣性也給人工智慧開發人員帶來了一些挑戰。本文簡要介紹了這些挑戰，並介紹了一個可以解決這些問題的編譯器解決方案。

讓我們先來看看這些挑戰，並向您介紹來自華盛頓大學和AWS研究團隊，然後讓您了解編譯器是如何工作的。

三個挑戰

首先，由於前端介面和後端實現之間的差異，從一個人工智慧框架切換到另一個AI框架是非常重要的。此外，演算法開發人員可能會使用多個框架作為開發和交付流程的一部分。在AWS上，我們有客戶希望在MXNet上部署他們的Caffe模型，以享受Amazon EC2上的高速性能。Joaquin Candela最近在博客中寫道，用戶可能會使用PyTorch進行快速開發，然後部署在Caffe2上。然而，在將模型從一個框架轉換到另一個框架後，有人抱怨二者的調試結果存在差異。

其次，框架開發人員需要維護多個後端，以保證智能手機晶元，GPU晶元等硬體的性能。以MXNet為例，它有一個可移植的，可以從頭開始構建的C++實現。它還附帶了目標相關的後端支持，比如cuDNN用於英偉達的 GPU晶元，MKLML用於英特爾的CPU晶元。保證這些不同的後端向用戶提供一致的數字結果是很有挑戰性的。

最後，晶元供應商需要為他們所構建的每一個新晶元支持多個人工智慧框架。每個框架中的工作負載都以獨特的方式表示和執行，因此，即使是一個簡單的卷積（Convolution）操作，也可能需要以不同的方式定義。而支持多個框架需要花費工程師們大量的工作。

不同的人工智慧框架和硬體給用戶帶來了巨大的好處，但是對於人工智慧開發者來說，給最終用戶帶來一致的結果是非常困難的。幸運的是，我們並不是第一個面對這種問題的人。計算機科學在不同的硬體上運行不同的編程語言有著悠久的歷史。解決這個問題的一個關鍵技術是編譯器。受傳統編譯器技術的啟發，Tianqi Chen，Thierry Moreau ，Haichen Shen，Luis Ceze，Carlos Guestrin以及Arvind Krishnamurthy這組來自華盛頓大學計算機科學工程學院的研究人員，聯合了來自AWS AI團隊的Ziheng Jiang，介紹了TVM堆棧技術來簡化這個問題。

今天，AWS和來自華盛頓大學的研究團隊共同發布了一種基於TVM堆棧技術的端到端編譯器，叫做NNVM。這種編譯器能直接將來自不同的深度學習前端的工作負載編譯成優化過的二進位機器碼。讓我們來看一下傳統的AI框架和NNVM編譯器的結構對比：

我們觀察到一個典型的人工智慧框架可以大致劃分為三個部分：

1、前端向用戶展示了一個易於使用的界面。

2、來自前端的工作命令通常被表示成計算圖（computation graph），它由數據變數和操作運算符組成。

3、不管是基本算術的操作運算符，還是神經網路層的操作運算符都被優化執行成對應不同硬體的操作運算符

而NNVM編譯器基於TVM堆棧中的兩個組件:處理計算圖的NNVM組件和處理張量操作運算符的TVM組件。

NNVM組件——計算圖媒介表示堆棧 (intermediate representation (IR) stack)

NNVM組件的目標是將來自不同框架的工作命令表示為標準化的計算圖，然後將這些高級計算圖轉換為執行圖（execution graph）。將計算圖以一種與框架無關的形式表達出來的想法，是受到來自Keras（一種深度學習框架）的層定義，以及來自Numpy（Python的一種工具）的張量操作運算符的啟發。

NNVM組件還附帶了一些稱為Pass的常式，它們通過LLVM約定來操作這些計算圖。這些常式可以在計算圖中添加新的屬性來執行計算圖，或者修改計算圖以提高效率。

TVM組件——張量媒介表示堆棧

TVM起源於Halide語言，它的執行對象是計算圖中的操作運算符，它把操作運算符優化成對應目標後端硬體的操作運算符。它與NNVM組件不同，它提供了一種與硬體無關的、對應特定領域的語言，以簡化在張量索引級別中的操作符執行。TVM還提供了一些進程調度原語，例如多線程和緩存，從而優化計算以充分利用硬體資源。這些進程調度是硬體相關的，並且可以手動編碼，也可以自動搜索優化過的模式。

NNVM編譯器支持的前端框架和後端硬體

下圖說明了NNVM編譯器如何處理前後端之間的關係：

MXNet是通過將其計算圖直接轉換為NNVM組件中的計算圖來支持的，Keras仍在發展中，但也以類似的方式支持。NNVM編譯器還可以使用諸如CoreML之類的模型格式。

因此，任何能夠使用這些格式的框架都可以使用這個編譯堆棧。

TVM目前使用多個代碼生成器來支持各種後端硬體。例如，它為X86和ARM等cpu架構生成LLVM IR。它也能夠為各種GPU架構輸出CUDA、OpenCL和Metal kernels。

NNVM編譯器要添加新的支持很簡單。對於一個新的前端框架，我們只需要將它的工作命令轉換為定義了計算圖和操作符規範的NNVM組件的格式。如果要添加新的後端硬體，我們可以重複使用TVM的操作符執行，然後只需指定如何有效地調度它們。

性能測試結果

我們用MXNet分別在GPU與CPU硬體後端上，以及ResNet18與MobileNet模型上，來與NNVM編譯器進行性能對比測試。

這是在GPU上的性能結果測試對比圖，可以看出NNVM編譯器的耗時更少。

這是在CPU上的性能結果測試對比圖，可以看出NNVM編譯器的耗時更少。

結論

NNVM編譯器能將一個高級計算圖編譯成優化過的二進位機器碼，這個編譯器基於TVM堆棧中的兩個組件：NNVM組件提供了一個帶有計算圖優化常式的計算圖和操作運算符的規範。TVM組件把操作運算符執行優化成對應目標硬體的操作運算符。而且，我們在兩個完全不同的硬體上——ARM CPU和Nvidia gpu進行了性能測試對比，證明了這個編譯器可以匹配甚至超越了目前最先進的計算性能。我們希望NNVM編譯器能夠極大地簡化AI的前端框架和後端硬體的設計，並幫助用戶在不同的前端和後端提供一致的結果。

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