陳天奇團隊發布NNVM編譯器，性能優於MXNet，李沐撰文介紹

最新 10-07

夏乙若朴編譯整理

量子位出品 | 公眾號 QbitAI

亞馬遜和華盛頓大學今天合作發布了開源的端到端深度學習編譯器NNVM compiler。

先提醒一句，NNVM compiler ≠ NNVM。

NNVM是華盛頓大學博士陳天奇等人2016年發布的模塊化深度學習系統，今年8月中旬，他們又推出了將深度學習工作負載部署到硬體的端到端IR堆棧TVM，也就是把深度學習模型更簡單地放到各種硬體上。

當時，陳天奇把TVM+NNVM描述為「深度學習到各種硬體的完整優化工具鏈」，而這次推出的NNVM compiler，是一個基於TVM工具鏈的編譯器。

項目作者之一陳天奇在微博上這樣介紹這個編譯器：

我們今天發布了基於TVM工具鏈的深度學習編譯器NNVM compiler。支持將包括mxnet，pytorch，caffe2, coreml等在內的深度學習模型編譯部署到硬體上並提供多級別聯合優化。速度更快，部署更加輕量級。支持包括樹莓派，伺服器和各種移動式設備和cuda, opencl, metal, javascript以及其它各種後端。歡迎對於深度學習，編譯原理，高性能計算，硬體加速有興趣的同學一起加入dmlc推動領導開源項目社區。

NNVM compiler對CoreML的支持，讓開發者可以在非iOS設備上部署CoreML模型。

AI開發界的挑戰

AWS AI首席科學家李沐（MXNet作者）在亞馬遜博客撰文介紹稱，推出這個編譯器，是為了應對深度學習框架多樣化為AI開發界帶來的三個挑戰：

一、

對於演算法的開發者來說，由於各AI框架的前端交互和後端實現之間都存在很多區別，換框架很麻煩，而開發和交付過程中可能會用到的框架不止一個。

比如說有的亞馬遜AWS雲服務用戶，為了獲得EC2上的加速性能，會想要把Caffe模型部署到MXNet上。

為了應對這個問題，之前Facebook和微軟也聯合發布了模型間轉換工具ONNX。

二、

框架的開發者需要維護多個後端，來保證自己的框架能適用於從手機晶元到數據中心GPU的各種硬體。

比如說MXNet，要支持英偉達GPU的cuDNN，還要支持英特爾CPU的MKLML。

三、

從晶元供應商的角度來看，他們每新開發一款晶元都需要支持多個AI框架，每個框架表示和執行工作負載的方式都不一樣，所以，就連卷積這樣一個運算，都需要用不同的方式來定義。

支持多個框架，就代表要完成巨大的工作量。

通過將框架中的深度學習模型直接部署到硬體，NNVM compiler自然也就解決了這些問題。

結構

NNVM compiler可以將前端框架中的工作負載直接編譯到硬體後端，能在高層圖中間表示（IR）中表示和優化普通的深度學習工作負載，也能為不同的硬體後端轉換計算圖、最小化內存佔用、優化數據分布、融合計算模式。

編譯器的典型工作流如下圖所示：

這個編譯器基於此前發布的TVM堆棧中的兩個組件：NNVM用於計算圖，TVM用於張量運算。

其中，NNVM的目標是將不同框架的工作負載表示為標準化計算圖，然後將這些高級圖轉換為執行圖。

TVM提供了一種獨立於硬體的特定域語言，以簡化張量索引層次中的運算符實現。另外，TVM還支持多線程、平鋪、緩存等。

對框架和硬體的支持

編譯器中的NNVM模塊，支持下圖所示的深度學習框架：

具體來說，MXNet的計算圖能直接轉換成NNVM圖，對Keras計算圖的直接支持也正在開發中。

同時，NNVM compiler還支持其他模型格式，比如說微軟和Facebook前不久推出的ONNX，以及蘋果CoreML。

通過支持ONNX，NNVM compiler支持Caffe2、PyTorch和CNTK框架；通過支持CoreML，這個編譯器支持Caffe和Keras。

而編譯器中的TVM模塊，目前附帶多個編碼生成器，支持多種後端硬體，其中包括為X86和ARM架構的CPU生成LLVM IR，為各種GPU輸出CUDA、OpenCL和Metal kernel。

性能

NNVM compiler聯合使用圖級和張量級優化以獲得最佳性能。常規的深度學習框架會將圖優化與部署runtime進行打包，而NNVM編譯器將優化與實際部署運行時分離。

採用這種方法，編譯的模塊只需要依賴於最小的TVM runtime，當部署在Raspberry Pi或移動設備上時，只佔用大約300KB。

陳天奇團隊對NNVM compiler的性能進行了基準測試，並與MXNet進行了比較。這個測試基於兩種典型的硬體配置：樹莓派上的ARM CPU和AWS上的Nvidia GPU。

Nvidia GPU

GPU的基準和時間表由Leyuan Wang（AWS/UCDavis）和Yuwei Hu（圖森）提供。他們在Nvidia K80上對NNVM編譯器和MXNet進行了比較，以CUDA8和CuDNN7作為後端。這是一個非常強的基線，因為MXNet開啟了從CuDNN中選擇最佳內核的自動調整功能。另外，他們還使用了MXNet中優化深度內核來優化MobileNet工作負載。

如圖所見，NNVM編譯器生成的代碼在K80上優於MXNet。這些改進源於圖和內核級別的優化。值得注意的是，NNVM編譯器自己升恆所有的優化GPU內核，而不需要依賴諸如CuDNN這樣的外部庫。

樹莓派3b

樹莓派編譯堆棧由Ziheng Jiang（AWS/FDU）提供。他們使用OpenBLAS和NNPack對NNVM和MXNet進行了比較，嘗試不同的設置來獲得MXNet的最佳表現，例如為3×3卷積在NNPack中開啟Winograd卷積，啟動多線程，並禁用了額外的調度程序（所有的線程都被NNPack使用）。