華盛頓大學聯合AWS開源NNVM：面向AI框架的新型端到端編譯器

選自AWS

機器之心編譯

參與：劉曉坤

開發 AI 演算法有很多的人工智慧框架可以選擇，而訓練和開發 AI 模型也能用很多種硬體做到。框架和硬體的多樣性對於維持 AI 生態系統的健康發展是至關重要的。但這種多樣性也給 AI 開發者帶來了不小的挑戰。這篇文章將簡要描述這些挑戰，並介紹一種新的編譯器解決這些問題。

我們先回顧這些挑戰，介紹 UW 和 AWS 研究團隊，然後看看這種編譯器是如何工作的。

三個挑戰

首先，由於前端介面以及後端實現的不同，從一個 AI 框架轉換到另一個就很不容易。此外，演算法開發者也可能使用不止一種框架以作為開發和輸出管道的一部分。在 AWS，我們就有客戶要求將他們的 Caffe 模型部署到 MXNET 上，以在 Amazon EC2 上獲得性能加速。根據 Joaquin Candela 最近的博客文章（https://research.fb.com/facebook-and-microsoft-introduce-new-open-ecosystem-for-interchangeable-ai-frameworks/），用戶可以先用 PyTorch 快速開發模型然後再部署到 Caffe2 上。然而，人們對將模型從一個框架轉換到另一個後，對結果的調試過程的差異還是頗不滿意。

其次，框架開發者需要維護許多的後端以確保硬體（從智能手機晶元到數據中心 GPU）的性能。以 MXNet 為例，它有一個從頭構建的可移植 C++實現，它還搭配了目標相關的後端支持如 Nvidia GPU 的 CuDNN 和 Intel CPUs 的 MKLML。要保證這些不同的後端傳送一致的數值結果是很大的挑戰。

最後，晶元供應商需要使其開發的每一個新的晶元都能支持多個 AI 框架。每一個框架的工作負載都以獨特的方式進行表示和執行，因而即使一個單獨的計算比如卷積也可能需要用多種方式定義。支持多個框架需要巨大的工程工作。

UW 和 AWS 研究團隊

多樣的 AI 框架和硬體能給用戶帶來莫大的好處，而這對於 AI 開發者來說，傳送一致的結果給終端用戶很困難。幸運的是，我們並不是第一個遇到這種問題的人。計算機科學就有在不同的硬體上運行多種程序的漫長歷史。而解決問題的一種關鍵的技術就是編譯器。由於編譯器技術的啟發，包括 Tianqi Chen,，Thierry Moreau，Haichen Shen,，Luis Ceze,，Carlos Guestrin，Arvind Krishnamurthy 在內的來自華盛頓大學的團隊聯合來自 AWS AI team 的 Ziheng Jiang 共同提出了 TVM 堆棧以簡化這個問題。

今天，AWS 聯合 UW 公布了新的基於 TVM 堆棧的端對端編譯器，它能直接從多個深度學習前端將工作負載編譯為優化的機器代碼。我們先看看其架構。

架構

我們觀察到一個典型的 AI 框架可以大致劃分成三個部分：

前端向用戶展示易用的界面；

從前端接收的工作負載通常表示成計算圖，由數據變數（a，b，和 c）和運算元（*和+）組成；

為多種硬體實現和優化運算元（從基本的算術運算到神經網路層）。

新的編譯器稱作 NNVM 編譯器，建立在 TVM 堆棧的兩個組成的基礎上：用於計算圖的 NNVM，和用於張量計算的 TVM。

NNVM – 計算圖中間表示（intermediate representation，IR）堆棧

NNVM 的目標是將來自不同的框架的工作負載表示成標準的計算圖，然後將這些高級圖翻譯成執行圖。將計算圖表示成框架不可知（framework-agnostic）的形式的想法，受到了 Keras 的層定義和 numpy 的張量運算元的啟發。

按照 LLVM 的慣例，NNVM 還搭配了稱作 Pass 的常式，以操作這些計算圖。這些常式要麼往圖上添加新的屬性以執行它們，要麼調整圖以提高效率。

TVM – 張量中間表示堆棧

TVM 起源於 Halide，用於在計算圖中實現運算元，並對其進行優化然後用於目標後端硬體。和 NNVM 不同，它提供了一個硬體無關的、領域特定的語言以在張量指標級上簡化運算元實現。TVM 還提供了調度原語比如多線程、 平鋪和緩存以優化計算，從而充分利用硬體資源。這些調度是硬體相關的，要麼用手工編寫，要麼也可以自動搜索優化模式。

支持的前端框架和後端硬體

支持的前端框架和後端硬體如下圖所示。

MXNet 通過轉換計算圖為 NNVM 圖，能直接被支持。對 Keras 的支持還在開發中，以和 MXNet 相似的方式。NNVM 編譯器也可以包括模型格式，比如 CoreML。因此任何能使用這些格式的框架都能使用這個編譯堆棧。

TVM 目前搭配了多種代碼生成器以支持多種後端硬體。例如，它能為 CPU（如 X86 和 ARM）生成 LLVM 中間表示。它能為多種 GPUs 輸出 CUDA、OpenCL 和 Metal kernels。

添加新的支持是很簡單的。對於一個新的前端，只需要將其同時定義了計算圖和操作者規範的工作負載轉換進 NNVM。如果要支持新的硬體，可以通過 TVM 再次使用運算元實現，並只需要明確規定如何有效的調度它們就行了。

性能結果

我們展示了以 MXNet 為前端，以及配置兩種典型硬體（樹莓派電腦的 ARM CPU 以及 AWS 的 Nvidia GPU）的 NNVM 編譯器的性能表現。儘管這兩種晶元有根本的區別，但也只在調度部分的代碼才出現了差異。

Nvidia GPUs

GPU 的調度大部分由 Leyuan Wang (AWS) 和 Yuwei Hu (TuSimple) 在他們實習期間編寫。我們比較了 NNVM 編譯器，和 MXNet 搭配 Nvidia K80 的 cuDNN 為後端的組合。在 cuDNN 不能得到有效支持的運算元比如深度（depthwise）卷積由手動優化 CUDA 的核實現。如上圖所示，在 ResNet18 和 MobileNet 的表現中，NNVM 編譯器都比 cuDNN 為後端要稍微好一些（快 1.2 倍）。

樹莓派電腦 3B

最適宜的調度由自動調整運動參數的模組（auto tuner）選出。特別是，我們通過基準問題測試在樹莓派電腦上的性能，為每一個給定形狀（shape）的運算元找到了最好的調度。

我們對比了 NNVM 編譯器和 MXNet。MXNet 使用 OpenBLAS 生成併兼容 NNPACK 庫，我們手工在 NNPACK 中啟用了 winograd 卷積演算法以獲得最佳性能。

如上圖所示，NNVM 編譯器在 Resnet18 中的表現要快 2.2 倍，並且在 MobileNet 上要快 11.5 倍。這主要都是由於深度卷積在 MXNet 中沒有得到優化（在 dnn 庫中缺少這樣的運算元），而 NNVM 編譯器受益於能直接為其生成有效的代碼的能力。

結論

本文中，我們介紹了 NNVM 編譯器，能將高級的計算圖編譯進優化的機器代碼中。該編譯器主要基於 TVM 堆棧的兩個主要部件：NNVM 提供計算圖、計算路徑優化的運算元的明細，通過使用 TVM 可為特定硬體部署、優化運算元。我們證明，在極小付出下，該編譯器就能在兩種完全不同的硬體上（ARM CPU 和英偉達 GPU）得到匹配甚至超越頂級表現。

我們希望 NNVM 編譯器能極大的簡化新型 AI 前端框架和後端硬體的設計，並能跨越前端和後端為用戶提供一致的結果。?

本文為機器之心編譯，轉載請聯繫本公眾號獲得授權。

?------------------------------------------------

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