「深度學習框架的未來」李沐：AWS開源端到端AI框架編譯器NNVM

2017 年 11 月 8 日，在北京國家會議中心舉辦的 AI WORLD 2017 世界人工智慧大會開放售票！早鳥票 5 折搶票最後1天！還記得去年一票難求的AI WORLD 2016盛況嗎？今年，我們邀請了冷撲大師」之父 Tuomas 親臨現場，且谷歌、微軟、亞馬遜、BAT、訊飛、京東和華為等企業重量級嘉賓均已確認出席。

有很多個人工智慧（AI）框架供你選擇來開發AI演算法。同樣你也可以選擇各種硬體來訓練和部署AI模型。框架和硬體的多樣性對於保持AI生態系統的健康至關重要。但是，這種多樣性也為AI開發人員帶來了幾個挑戰。本文將簡要介紹這些挑戰，並引入一個編譯器解決方案來輔助解決這些問題。

下面，我們首先回顧挑戰，介紹華盛頓大學和AWS研究團隊，然後介紹編譯器的工作原理。

三大挑戰

首先，由於前端介面和後端實現之間的差異，從一個AI框架切換到另一個AI框架是很麻煩的。此外，演算法開發人員在開發和交付流程中可能會使用多個框架。在AWS，我們有客戶希望在MXNet上部署其Caffe模型，從而享受Amazon EC2上的加速性能。根據Joaquin Candela最近的博客，用戶可能會使用PyTorch快速開發，然後在Caffe2上部署。但是，我們也聽到有人表示在將模型從一個框架轉換到另一個框架之後調試結果差異十分困難。

其次，框架開發人員需要維護多個後端，才能確保從智能手機晶元到數據中心GPU的各種硬體的性能。以MXNet為例，MXNet有一個從頭構建的攜帶型C++實現，還附帶了專門的後端支持，比如英偉達GPU的CuNNN和英特爾CPU的MKLML。確保這些不同的後端為用戶提供一致的數值結果十分具有挑戰性。

最後，晶元廠商製造的每款新晶元也需要支持多個AI框架。每個框架中的workload都以獨特的方式進行表徵和執行，因此即使像卷積一樣的運算也可能需要以不同的方式進行定義。因此，在晶元設計和製造方面，支持多個框架也需要大量的工程設計。

華盛頓大學和AWS研究團隊

不同的AI框架和硬體為用戶帶來巨大的收益，但AI開發人員為終端用戶提供一致性卻是非常具有挑戰性的事情。幸運的是，我們並不是第一個面對這種問題的人。計算機科學在不同硬體上運行各種編程語言的歷史由來已久。解決這個問題的一個關鍵技術是編譯器。由編譯技術驅動，來自華盛頓大學（UW）計算機科學與工程系保羅·艾倫學院的陳天奇，Thierry Moreau，Haichen Shen，Luis Ceze，Carlos Guestrin和Arvind Krishnamurth等一批研究人員，與AWS AI團隊的Ziheng Jiang，提出了TVM堆棧來簡化這個問題。

今天，AWS很高興與UW研究團隊一起宣布基於TVM堆棧的端到端編譯器，它將workload直接從各種深度學習前端編譯成優化的機器代碼。

架構

下面，我們先來看這個編譯器架構。

我們可以注意到，典型的AI框架大致分為三部分：

• 前端開放一個易於使用的界面給用戶；

• 從前端收到的workload通常被表示為由數據變數（a，b和c）和運算符（*和+）組成的計算圖；

• 從基本的算術運算到神經網路層的運算符都可以針對多個硬體進行實現和優化。

新的編譯器，NNVM編譯器，基於TVM堆棧中的兩個組件：用於計算圖的NNVM和用於張量運算符的TVM。

NNVM - 計算圖中間表示（IR）堆棧

NNVM的目標是將不同框架的workload表示為標準計算圖，然後將這些高級圖轉換為執行圖。這種計算圖的靈感來自Keras中的層定義和numpy的張量運算符。

NNVM還隨附常式，名叫Pass，遵循LLVM約定。這些常式可以向圖中添加新屬性來執行它們或修改圖以提高效率。

TVM - 張量IR堆疊

源自Halide的TVM實現了計算圖中使用的運算符，並針對目標後端硬體進行了優化。與NNVM不同，TVM提供了一種獨立於硬體的域特定語言，簡化張量索引級別中的操作符實現。TVM還提供調度原語（如多線程，平鋪和高速緩存）來優化計算以充分利用硬體資源。這些計劃與硬體有關，可以手動編碼，也可以自動搜索優化的模式。

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

支持的前端和後端如下圖所示。

MXNet直接將其計算圖轉換為NNVM圖得到支持，Keras也以類似的方式得到支持，不過正在開發中。NNVM編譯器也可以採用模型格式，如CoreML。因此，凡是能夠使用這些格式框架都可以使用此編譯堆棧。

TVM目前隨附多個代碼生成器，支持各種後端硬體。例如，TVM為CPU（如X86和ARM）生成LLVM IR，也可以為各種GPU輸出CUDA，OpenCL和Metal內核。

添加新的支持的操作也很簡單。對於新的前端，我們只需要將其workload轉換為定義計算圖和運算符規範的NNVM即可。要添加新的硬體，我們可以重複使用TVM的運算符實現，只需要指定有效schedule即可。

性能結果

我們使用MXNet展示NNVM編譯器性能，前端硬體配置是：ARM上的Raspberry PI和AWS上的英偉達GPU。雖然這兩種晶元間的架構差異巨大，但代碼上的唯一差異在於調度部分。

GPU的schedule主要由Leyuan Wang（AWS）和Yuwei Hu（TuSimple）在實習期間編寫。比較NNVM編譯器與MXNet，英偉達K80的cuDNN後端。像深度卷積的運算符這樣在cuDNN中沒有有效支持的，通過手動優化的CUDA內核來實現。可以看出，NNVM編譯器運行ResNet18和MobileNet比cuDNN後端略好（1.2倍快）。

在Raspberry PI的情況下，我們通過自動調諧器選擇最佳schedule。我們通過在Raspberry Pi上對運算符性能進行基準測試，給出了每個運算符在給定形狀的最佳schedule。

還是比較NNVM編譯器與MXNet。MXNet默認啟用OpenBLAS和NNPACK，我們還手動打開了NNPACK中的winograd卷積以獲得最佳性能。

可以看出，NNVM編譯器在Resnet18上快了2.2倍。而MobileNet上更有11.5倍的差異。這主要是由於在MXNet中沒有優化深度卷積（因為dnn庫中缺少這樣的運算符），而NNVM編譯器則受益於直接生成高效的代碼。

結論

我們介紹了NNVM編譯器，它將高級計算圖編譯成優化的機器代碼。NNVM編譯器基於TVM堆棧中的兩個組件：NNVM使用圖優化常式提供計算圖和運算符的specification，運算符通過使用TVM針對目標硬體實現和優化。我們證明，只需花費很少的精力，這種編譯器可以匹配甚至超過兩個完全不同的硬體（ARM CPU和Nvidia GPU）上的state-of-the-art性能。

我們希望NNVM編譯器可以大大簡化新的AI前端框架和後端硬體的設計，並幫助在各種前端和後端為用戶提供一致的結果。

原文鏈接：https://amazonaws-china.com/cn/blogs/ai/introducing-nnvm-compiler-a-new-open-end-to-end-compiler-for-ai-frameworks/

【掃一掃或點擊閱讀原文搶購五折「早鳥票」】

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