用TVM在硬件平臺(tái)上部署深度學(xué)習(xí)工作負(fù)載的端到端 IR 堆棧

深度學(xué)習(xí)已變得無(wú)處不在，不可或缺。這場(chǎng)革命的一部分是由可擴(kuò)展的深度學(xué)習(xí)系統(tǒng)推動(dòng)的，如滕索弗洛、MXNet、咖啡和皮托奇。大多數(shù)現(xiàn)有系統(tǒng)針對(duì)范圍狹窄的服務(wù)器級(jí) GPU 進(jìn)行了優(yōu)化，需要在其它平臺(tái)，如移動(dòng)電話、物聯(lián)網(wǎng)設(shè)備和專用加速器（FPGA、ASIC）上部署大量精力。隨著深度學(xué)習(xí)框架和硬件后端數(shù)量的增加，建議建立一個(gè)統(tǒng)一的中間表示 （IR） 堆棧，以縮小以生產(chǎn)力為中心的深度學(xué)習(xí)框架與面向性能或效率的硬件后端之間的差距。

TVM 是一個(gè)新穎的框架，可以：
? 表示并優(yōu)化 CPU、GPU 和其它專業(yè)硬件的常見深度學(xué)習(xí)計(jì)算工作負(fù)載
? 自動(dòng)轉(zhuǎn)換計(jì)算圖以最大限度地減少內(nèi)存利用、優(yōu)化數(shù)據(jù)布局和融合計(jì)算模式
? 提供從現(xiàn)有前端框架到裸機(jī)硬件的端到端編譯，一直到瀏覽器可執(zhí)行的 javascripts。
TVM可以輕松地在移動(dòng)電話、嵌入式設(shè)備甚至瀏覽器上運(yùn)行深度學(xué)習(xí)工作負(fù)載，不需要額外的支持。TVM 還為眾多硬件平臺(tái)（包括依賴新型計(jì)算原始的專用加速器）的深度學(xué)習(xí)工作量提供了統(tǒng)一的優(yōu)化框架。

采用編譯器方案，提供兩個(gè)中間表示層，有效地將高級(jí)深度學(xué)習(xí)算法降低到多個(gè)硬件后端。
開源TVM包包含x86、ARM、OpenCL、Metal、CUDA和Javascript的原型優(yōu)化。致力于專業(yè)硬件加速和 Nvidia 的 GEMM 優(yōu)化Volta架構(gòu)提供支持。
技術(shù)細(xì)節(jié)
TVM 堆棧的目標(biāo)是提供可重復(fù)使用的工具鏈，以編譯從深度學(xué)習(xí)框架前端到低級(jí)機(jī)器代碼的高層神經(jīng)網(wǎng)絡(luò)描述，用于多個(gè)硬件后端。以 Apache MXNet 為例，以下代碼片段演示，如何使用 TVM 將深度學(xué)習(xí)模型的高level描述，編譯為針對(duì)目標(biāo)硬件定制的優(yōu)化可執(zhí)行模塊。

挑戰(zhàn)在于支持多個(gè)硬件后端，同時(shí)將計(jì)算、內(nèi)存和能量消耗保持在最低水平。借用AICompiler的智慧，以彌合眾多深度學(xué)習(xí)框架和硬件后端之間的差距：
構(gòu)建了由 NNVM 組成的兩級(jí)中間層、用于任務(wù)調(diào)度和內(nèi)存管理的高級(jí)中級(jí)代表 （IR） 和用于優(yōu)化計(jì)算內(nèi)核的具有表現(xiàn)力的低level IR 的 TVM。
堆棧的第一個(gè)級(jí)別是基于計(jì)算圖形的表示。計(jì)算圖是指示式循環(huán)圖，表示計(jì)算作為節(jié)點(diǎn)，數(shù)據(jù)流依賴性表示邊緣。此表示非常強(qiáng)大：
將操作屬性bake到計(jì)算圖中，指定轉(zhuǎn)換規(guī)則，迭代優(yōu)化計(jì)算圖形。這是大多數(shù)現(xiàn)有深度學(xué)習(xí)框架所采用的常見方法，包括 TVM 堆棧中的 NNVM 圖形表示、TensorFlow XLA 和英特爾的 ngraph。

圖形優(yōu)化框架可以支持許多強(qiáng)大的優(yōu)化。例如，提供了sublinear亞線性內(nèi)存優(yōu)化功能?？稍趩蝹€(gè) GPU 上訓(xùn)練 1000 層ImageNet ResNet。

發(fā)現(xiàn)僅基于計(jì)算圖的 IR 不足以解決支持不同硬件后端的挑戰(zhàn)。原因是，對(duì)于每個(gè)硬件的后端，可能會(huì)以非常不同的方式映射和優(yōu)化單個(gè)圖形算子，如卷積或矩陣乘法。這些針對(duì)硬件的優(yōu)化在內(nèi)存布局、并行線程模式、緩存訪問(wèn)模式和硬件原始選擇方面可能存在巨大差異。希望能夠以通用表示方式明確表達(dá)這些優(yōu)化模塊，有效引導(dǎo)優(yōu)化空間。
構(gòu)建一個(gè)低level級(jí)別的表示來(lái)解決這個(gè)問(wèn)題。此表示基于索引公式，為復(fù)發(fā)計(jì)算提供額外支持。

低級(jí) IR 采用現(xiàn)有圖像處理語(yǔ)言（如Halid或darkroom）的原則，定制具有表現(xiàn)力的深度學(xué)習(xí) DSL。TVM 構(gòu)建低level級(jí)別優(yōu)化的理論來(lái)自循環(huán)轉(zhuǎn)換工具，如循環(huán)和多面體分析。從 MXNet, TensorFlow, Theano.等深度學(xué)習(xí)框架中使用的數(shù)據(jù)流描述語(yǔ)言中汲取靈感。然后在調(diào)度階段處理 TVM 中描述的算法，應(yīng)用針對(duì)目標(biāo)硬件后端量身定制的轉(zhuǎn)換。

TVM 包括 CPU 優(yōu)化框架中常見的標(biāo)準(zhǔn)原型轉(zhuǎn)換。更重要的是，TVM 通過(guò)利用線程調(diào)度模式、數(shù)據(jù)布局轉(zhuǎn)換和強(qiáng)大的新原型計(jì)算，整合了針對(duì) GPU 的新型優(yōu)化。使用 TVM 與 NNVM 相結(jié)合，為優(yōu)化整個(gè)軟件堆棧中的深度學(xué)習(xí)工作負(fù)載提供了豐富的機(jī)會(huì)，實(shí)現(xiàn)了聯(lián)合計(jì)算圖形級(jí)別和算子級(jí)別的優(yōu)化。
多語(yǔ)言和平臺(tái)支持
TVM 的眾多優(yōu)勢(shì)之一，在于對(duì)多個(gè)平臺(tái)和語(yǔ)言的豐富支持。框架的兩個(gè)組件：包含完整優(yōu)化庫(kù)以生成優(yōu)化機(jī)器代碼的編譯器堆棧，以及輕量級(jí)lightweight的runtime，以及在不同平臺(tái)上部署編譯模塊所需的便攜性。

TVM 當(dāng)前支持嵌入式編譯器堆棧的python和C++接口。設(shè)計(jì)框架時(shí)，考慮到了最大的再利用，以便編譯器堆棧改進(jìn)，可以在 Python 和C++組件之間互換應(yīng)用。
提供輕量級(jí)的runtime，可以在包括 Android、iOS、raspberry pi和 Web 瀏覽器在內(nèi)的平臺(tái)上，以javascript, java, python, C++等語(yǔ)言直接運(yùn)行 TVM 編譯代碼。
遠(yuǎn)程部署和執(zhí)行

TVM 支持使用 TVM RPC 對(duì)嵌入式設(shè)備進(jìn)行交叉編譯和測(cè)試，這是一個(gè)輕量級(jí)界面，可在遠(yuǎn)程嵌入式設(shè)備上部署和執(zhí)行 TVM 交叉編譯模塊。這為 TVM 用戶提供了熟悉的高級(jí) Python 界面，以便在各種低級(jí)嵌入式設(shè)備上遠(yuǎn)程編譯、優(yōu)化和測(cè)試深度學(xué)習(xí)算法。
性能
TVM能夠靈活地探索各種深度學(xué)習(xí)內(nèi)核的豐富優(yōu)化空間，用于多個(gè)硬件平臺(tái)。例如，TVM為最關(guān)心的內(nèi)核和平臺(tái)定制數(shù)據(jù)布局和融合模式要求。基線庫(kù)是針對(duì)更通用的問(wèn)題創(chuàng)建的，而 TVM 的優(yōu)化內(nèi)核則針對(duì)通過(guò)自動(dòng)調(diào)整過(guò)程評(píng)估的工作負(fù)載進(jìn)行了大量調(diào)整。TVM 充當(dāng)快速生成獲取此類專用內(nèi)核的橋梁。
Raspberry Pi樹莓派
在結(jié)果的第一部分，將TVM CPU調(diào)度與樹莓派3B執(zhí)行網(wǎng)格工作負(fù)載上的nnpack進(jìn)行了比較。使用TVM實(shí)現(xiàn)直接卷積，而nnpack用于執(zhí)行3x3內(nèi)核的winograd conv。

可以發(fā)現(xiàn)，通過(guò) TVM 的自動(dòng)調(diào)諧內(nèi)核，可以獲得類似于 nnpack 中hand-optimized kernels手動(dòng)優(yōu)化內(nèi)核的性能，用于aspberry pi實(shí)驗(yàn)。
GPU 結(jié)果
創(chuàng)建了一個(gè)端到端的編譯管道，可以將 MxNet 模型編譯到 TVM 執(zhí)行圖形。通過(guò)自動(dòng)將算子融合在一起，讓 TVM 生成融合內(nèi)核，在圖形節(jié)點(diǎn)內(nèi)和圖形節(jié)點(diǎn)之間應(yīng)用優(yōu)化。對(duì)mobilenet ImageNet的工作量進(jìn)行了基準(zhǔn)測(cè)試，并討論了以下結(jié)果：

TVM在速度方面可以優(yōu)于baseline基線方法。更有趣的是，內(nèi)核融合帶來(lái)了額外的加速。值得一提的是，TVM 能夠自行生成所有優(yōu)化的 GPU 內(nèi)核，無(wú)需依賴 CuDNN 等外部庫(kù)。
開源努力
TVM最初是華盛頓大學(xué)保羅·艾倫學(xué)校計(jì)算機(jī)科學(xué)與工程系的研究項(xiàng)目。TVM 堆棧旨在支持DLPack，這是多個(gè)主要深度學(xué)習(xí)框架對(duì)張力數(shù)據(jù)結(jié)構(gòu)的共識(shí)。收到了來(lái)自UW, AWS, Qiho 360, Facebook, HKUST, TuSimple, UCDavis, SJTU以及DMLC開源社區(qū)和DLPack倡議成員的早期貢獻(xiàn)。今后，該項(xiàng)目將遵循Apache open-source model開源模式，創(chuàng)建社區(qū)維護(hù)項(xiàng)目。
? Halide：TVM使用HalideIR作為數(shù)據(jù)結(jié)構(gòu)，用于arithematic簡(jiǎn)單化和低level lowering。HalideIR來(lái)自Halide。還從Halide那里學(xué)習(xí)，在TVM中實(shí)施lowing管道。
? Loopy：使用整數(shù)集分析及其循環(huán)轉(zhuǎn)換原始。
? Theano：symbolic scan算子的再現(xiàn)設(shè)計(jì)靈感。
源代碼Source code
? Github page頁(yè)面可以在這里找到： https://github.com/dmlc/tvm
? TVM 兼容DLPack，因此可以輕松支持采用該標(biāo)準(zhǔn)的框架，如 MXNet, PyTorch, Caffe2和微小 dnn。

總結(jié)

以上是生活随笔為你收集整理的用TVM在硬件平台上部署深度学习工作负载的端到端 IR 堆栈的全部?jī)?nèi)容，希望文章能夠幫你解決所遇到的問(wèn)題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。