自動生成硬體優化內核:陳天奇等人發布深度學習編譯器TVM

2021-01-13 機器之心Pro

本文由機器之心編輯,「機器之心」專注生產人工智慧專業性內容,適合開發者和從業者閱讀參考。點擊右上角即刻關注。

有了 TVM,業界與學界開發者們可以快速、輕鬆地在各個系統(包括手機、嵌入式設備與低功耗晶片)上部署深度學習應用程式,同時無須擔心資源與速度的限制。「TVM 作為神經網絡和硬體後端之間的共同層,消除了為每類設備或伺服器優化各自基礎架構的需要。」TVM 項目負責人陳天奇表示,「我們的框架允許開發人員快速、輕鬆地部署和優化大量硬體設備上的深度學習系統。」

項目連結:https://github.com/dmlc/tvm

深度學習如今無處不在且必不可少。這次創新部分得益於可擴展的深度學習系統,比如 TensorFlow、MXNet、Caffe 和 PyTorch。大多數現有系統針對窄範圍的伺服器級 GPU 進行了優化,並且需要在其他平臺(如手機、IoT 設備和專用加速器(FPGA、 ASIC))上部署大量工作。隨著深度學習框架和硬體後端數量不斷增加,我們提出了一個統一的中間表徵(IR)堆棧,可以彌補以生產力為中心的深度學習框架與性能或效率導向的硬體後端之間的差距。

我們很高興宣布 TVM 的發行將解決這一問題。TVM 是一個全新的框架,它可以:

優化 CPU、GPU 和其他專業化硬體的常規深度學習的計算量

自動轉換計算圖以最小化內存利用率,優化數據布局,融合計算模式提供從現有的前端框架到裸機硬體的端到端編譯,一直到瀏覽器可執行的 Javascript有了 TVM 的幫助,只需少量額外工作我們即可輕易地在手機端、嵌入式設備甚至瀏覽器上運行深度學習任務。TVM 還為多硬體平臺上的深度學習工作負載提供了統一的優化框架,其中包括依賴全新計算原語的專用加速器。

通過採用編譯器社區的共同理念,並提供兩個中間表徵層,我們有效地將高級深度學習算法降低到多硬體後端。在今天的版本中,我們開源了包含 x86、ARM、OpenCL、Metal、CUDA 和 Javascript 優化原語的 TVM 工具包。目前,我們正積極地為專用硬體加速和英偉達的 GEMM 優化的 Volta 架構提供支持。

技術細節

TVM 堆棧的目標在於提供一個可重複使用的工具鏈來將高級神經網絡描述從深度學習框架前端向下編譯為多硬體後端的低級機器代碼。拿 Apache MXNet 作為前端的例子,以下代碼段演示了如何使用 TVM 將深層學習模型的高級描述編譯為針對目標硬體定製的優化可執行模塊。

挑戰在於支持多個硬體後端,同時將計算、內存和能源足跡(energy footprint)保持在最低水平。我們借鑑了編譯器社區的共同理念,以彌合大量深度學習框架和硬體後端之間的差距:我們構建了一個由 NNVM 組成的兩級中間層,一個可用於任務調度、內存管理以及 TVM(用於優化計算內核的富有表現力的低級 IR)的高級中間表徵(IR)。

堆棧的第一級是基於計算圖的表徵。計算圖是一個有向無環圖,其表徵作為節點的計算和作為邊緣的數據流依賴性。這一表徵非常強大:它允許我們將操作屬性融入到計算圖中,並指定轉換規則以迭代優化計算圖。這是一個被大多數現有深度學習框架的採用的慣常方法,包括 TVM 堆棧中的 NNVM 圖表徵、TensorFlow XLA 和英特爾的 ngraph。

很多強大的優化可由圖優化框架支持。例如,我們提供了一個子線性內存優化功能,允許用戶在單塊 GPU 上訓練 1000 層的 ImageNet ResNet。

然而,我們發現只基於 IR 的計算圖並不足以應對支持不同硬體後端的挑戰。原因是像卷積或矩陣乘法一樣的圖運算符可以非常不同的方式映射和優化每個硬體後端。這些特定硬體的優化在內存布局、並行線程模式、緩存訪問模式和硬體原語選擇方面的變數可能很大。我們希望能夠以通用表徵的方式明確地表達這些優化旋鈕(optimization knob),以有效地導航優化空間。

我們構建了一個低級表徵以解決這一問題。該表徵基於索引公式,並額外支持遞推計算。

低級 IR 採用現有的圖像處理語言(比如 Halide 或 darkroom)的原理以制定一個富有表現力的深度學習 DSL。TVM 構建了由循環轉換工具比如 loopy、多面體分析啟發的圖優化。我們也從深度學習框架比如 MXNet、TensorFlow、Theano 使用的數據流描述語言汲取靈感,然後在調度階段對 TVM 中描述的算法進行處理,以應用針對目標硬體後端定製的轉換。

TVM 包括通常出現在 CPU 優化框架中的標準轉換原語。更重要的是,TVM 通過利用線程協作模式、數據布局轉換和強大的新計算原語,集成了針對 GPU 的全新優化原語。TVM 與 NNVM 的結合為優化軟體堆棧中深度學習工作負載提供了一個很好的機會,從而實現了圖級別與運算符級別優化的聯合計算。

多語言和平臺支持

TVM 的眾多優勢之一在於它可以支持多種語言和平臺。TVM 框架中有兩個組件:編譯器棧,其中包含完整的優化庫以產生優化的機器代碼;以及輕量級運行環境,提供在不同平臺上部署編譯模塊所需的可移植性。

TVM 目前支持嵌入式編譯器堆棧的 Python 和 C++ 接口。我們在設計中最大限度地重用框架,以便編譯器堆棧的改進可以在 Python 和 C++ 組件之間互換使用。我們還提供了一個輕量級運行環境,可以在包括 Android、iOS、樹莓派和 Web 瀏覽器在內的平臺上直接運行諸如 Javascript、Java、Python 和 C++ 等語言的 TVM 編譯代碼。

遠程部署和執行

TVM 支持使用 TVM PRC 進行交叉編譯,測試嵌入式設備,這是一種輕量級界面,用於在遠程嵌入式設備上部署和執行 TVM 交叉編譯模塊。我們 TVM 用戶提供了易用的高級 Python 界面,用於在各種低級嵌入式設備上進行遠程編譯、優化,並測試深度學習算法。

性能

TVM 目前仍然處在早期發展階段,我們會在未來看到更多改進,但目前它已經展現出了強大的實力。

TVM 為我們提供了靈活性,讓我們可以探索不同硬體平臺上,不同深度學習內核中豐富的優化空間。例如,TVM 允許我們裁剪數據布局,為內核與平臺融合模式需求。注意,構建基線庫可用於更通用的問題,而 TVM 的優化內核通過自動調整過程對人們評估的工作負載進行了大量調整。TVM 可以作為快速生成專屬內核的絕佳工具。

本節列出的結果仍有提升空間。

樹莓派

在結果的第一部分,我們比較了 TVM CPU 計劃與 nnpack 在樹莓派 3B 上執行 ResNet 任務的表現。由於時間有限,我們使用 TVM 實現直接卷積,而 nnpack 對 3×3 內核執行 winograd 轉換。

我們可以看到,使用 TVM 自動調整內核,我們可以獲得類似於 nnpack 中的手動優化內核的性能(實驗在樹莓派上進行)。

GPU 結果

PS. 該測試結果由 Leyuan Wang(AWS/UCDavis)、Yuwei Hu(圖森未來)與 Weitang Liu(AWS/UCDavis)提供。

為了進行概念證明,我們構建了一種端到端編譯管道,可以將 MXNet 模型編譯成 TVM 執行圖。我們通過自動融合運算符並使 TVM 生成融合的內核,在圖形節點之間和之中應用優化。我們對 MobileNet 在 ImageNet 上的運行表現進行基準測試,結果如下:

我們可以看到,TVM 在速度上快過基準方法。更有意思的是,內核融合可以帶來更快的速度提升。值得一提的是,TVM 可以自行生成所有優化 GPU 內核,而不依賴於外部庫,如 CuDNN。

我們正在開展新的實驗,期待獲得更好的結果。

開源

TVM 始於華盛頓大學 Paul G. Allen 計算機科學與工程學院的研究項目。它被設計成支持 DLPack——一個被多種深度學習框架共享的張量數據結構。在開發過程中,我們得到了來自華盛頓大學、AWS、奇虎 360、Facebook、香港科技大學、圖森未來、UCDavis、上海交大以及 DMLC 開源社區和 DLPack 計劃成員們的支持。展望未來,該項目將繼續遵循 Apache 開源模式,力圖構建一個由社區維護的項目,歡迎大家的參與。

致謝

本項目的成功離不開以下所有人的通力合作:Yizhi Liu(奇虎 360)、Yuwei Hu(圖森未來)、Xingjian Shi(香港科技大學)、Leyuan Wang(AWS/UCDavis)、Nicolas Vasilache(Facebook)、Jian Weng(UCLA)、Weitang Liu(AWS/UCDavis)、Edward Z. Yang(Facebook)、Lianmin Zheng(上海交通大學)、Qiao Zhang(華盛頓大學)、William Moses(Facebook/MIT)與 Hu Shiwen。作者同時感謝 Xianyi Zhang 和他富有建設性的討論。

原始碼

本工具 GitHub 頁面:https://github.com/dmlc/tvm

下周將向 Apache MXNet 上遊提供基於 TVM 堆棧和 MXNet 的圖形編譯工具鏈示例TVM 與 DLPack 兼容,這意味著它可以輕鬆支持 MXNet、PyTorch、Caffe2 和 tiny-dnn 等採用該標準的框架。 原文連結:http://tvmlang.org/2017/08/17/tvm-release-announcement.html

相關焦點

  • 陳天奇等人提出TVM:深度學習自動優化代碼生成器
    TVM 是由華盛頓大學在讀博士陳天奇等人提出的深度學習自動代碼生成方法,去年 8 月機器之心曾對其進行過簡要介紹。該技術能自動為大多數計算硬體生成可部署優化代碼,其性能可與當前最優的供應商提供的優化計算庫相比,且可以適應新型專用加速器後端。
  • 陳天奇的tvm更新了:VTA,開源AI晶片棧
    陳天奇 發表於 tvm社區 量子位 搬運 | 公眾號 QbitAI華盛頓大學陳天奇博士近日在tvm社區介紹了VTA(Versatile Tensor Accelerator)項目,量子位全文搬運如下。然而,當我們的合作真正開始的時候,才發現問題比原來想的要困難的多:深度學習晶片的主要問題遠遠超過了硬體本身。深度學習硬體真正的目標,是希望大家都可以把自己的深度學習的模型跑在AI晶片上。而加速器帶來的特殊定製導致對於每個加速器,我們都需要重新設計一套驅動,以及上面的軟體棧—-包含編譯,優化和部署模塊。
  • 陳天奇:深度學習編譯技術的現狀和未來
    作者:陳天奇傳統的深度學習框架採用人工優化算子,然後建立運行時圖解釋器來解決內存分配調度等問題。深度學習編譯器技術路線一般指在優化過程中採用了自動或者半自動的代碼生成用以替代人工優化。深度學習編譯器無疑是最近非常熱門的話題。本文主要探討深度學習編譯技術的現狀和未來。
  • 2017中國AI英雄風雲榜技術創新人物候選人之陳天奇
    陳天奇,華盛頓大學計算機系博士生,深度學習編譯器TVM發明人。陳天奇,華盛頓大學計算機系博士生,此前畢業於上海交通大學ACM班,研究方向為大規模機器學習。
  • 陳天奇任 CTO,華盛頓大學TVM團隊創業,獲390萬美元種子輪融資
    作為機器學習領域知名的青年華人學者,陳天奇加入創業公司的消息引起了業內極大的關注。今日,OctoML 宣布完成 390 萬美元種子輪融資,主要面向缺少源或知識來構建優化深度學習模型的群體。 本輪融資由麥德羅納風險投資集團領投,Amplify Partners 參投。
  • MIT Taco項目:自動生成張量計算的優化代碼,深度學習加速效果提高...
    舉例來說,目前我們常用的深度學習框架,如TensorFlow、PyTorch等都會將一個深度學習的模型轉換為一個由基本計算符組成的數據流圖,再有下層的計算引擎一次調度執行這些節點對應的內核函數(對於數據圖的動態優化參見雷鋒網此前文章《緊跟未來深度學習框架需求,TensorFlow推出Eager Execution》)。
  • 陳天奇創業公司首個SaaS產品:快速構建部署ML應用,跨平臺優化
    去年,知名人工智慧青年學者、華盛頓大學博士陳天奇參與創建了 ML 創業公司 OctoML。該公司致力於打造一個可擴展、開放、中立的端到端棧,用於深度學習模型的優化和部署。這不,他們推出了首個 SaaS 產品 Octomizer,其能夠幫助開發者或數據科學家更方便、更省事省力地將 ML 模型部署到設備上。
  • 從ACM班、百度到亞馬遜,深度學習大牛李沐的開掛人生
    作為一經推出便大受追捧的 MXNet 深度學習框架的主要貢獻者之一,李沐功不可沒。值得注意的是,這個由 DMLC(Distributed Machine Learning Community)打造的深度學習框架,創建者都是中國人,以陳天奇、李沐、解浚源等為代表。現在,包括 AWS 在內的很多企業都將 MXNet 作為主流支持框架。
  • 阿里將 TVM 融入 TensorFlow,在 GPU 上實現全面提速
    有了這些預定義的形狀,我們可以提前生成高度優化的 CUDA 內核(固定形狀的計算可以帶來最佳優化潛能)。同時,還將生成一個適合大多數形狀的通用 batch 矩陣相乘內核,為沒有提前生成內核的形狀提供回退機制。我們將生成的針對特定形狀的高效內核和回退機制集成到 Tensorflow 中。
  • 內核融合:GPU深度學習的「加速神器」
    【按:在深度學習「紅透」半邊天的同時,當前很多深度學習框架卻面臨著共同的性能問題:被頻繁調用的代數運算符嚴重影響模型的執行效率。本文中,微軟亞洲研究院研究員薛繼龍將為大家介紹能夠輕鬆玩轉計算性能的「加速神器」——內核融合,探討內核融合在加速深度學習上的主要方法以及當前面臨的主要挑戰。】
  • 華為深度學習框架MindSpore正式開源:自動微分不止計算圖
    今年的華為開發者大會 HDC 2020 上,除了昇騰、鯤鵬等自研晶片硬體平臺之外,最令人期待的就是深度學習框架 MindSpore 的開源了。今天上午,華為 MindSpore 首席科學家陳雷在活動中宣布這款產品正式開源,我們終於可以在開放平臺上一睹它的真面目。
  • 華盛頓大學成立SAML實驗室:陳天奇參與,探索跨系統堆棧多層問題
    安妮 發自 凹非寺量子位 出品 | 公眾號 QbitAI△ 陳天奇昨天,TVM、XGBoost、cxxnet等機器學習工具的作者陳天奇宣布,自己所在的華盛頓大學新成立了一個實驗室陳天奇本人在微博上介紹說,實驗室將有整合系統、體系結構、機器學習和程序語言4個方向的團隊推進深度學習算法、系統的硬體和協同設計研究和TVM等開源項目。△ 陳天奇微博截圖四個方向華盛頓大學為什麼想建立這個實驗室?
  • 12倍端到端加速,陳天奇創業公司OctoML提克服二值網絡瓶頸新方法
    選自medium作者:Josh Fromm機器之心編譯機器之心編輯部去年,TVM 開發團隊陳天奇等人創建了 OctoML公司,旨在「讓機器學習可以部署在所有硬體上」。這些改進一部分是由新的算法和架構創新所推動的,但對於深度學習任務來說,不斷擴展的算力和內存也使其取得了重大進展。隨著 ML 準確率的提升,模型所需要的算力和內存也不斷增加。早在 2016 年,我們就可以看出模型大小和準確率之間的關係了。
  • 清華自研深度學習框架「計圖」開源!多項任務性能超過PyTorch
    與主流的深度學習框架TensorFlow、Pytorch不同,Jittor是一個完全基於動態編譯(Just-in-time)、使用元算子和統一計算圖的深度學習框架。研發團隊介紹稱,開發Jittor是為了將新技術、硬體和模型的能力,更好地釋放出來。
  • 深度學習框架簡史:TF和PyTorch雙頭壟斷,未來十年迎來黃金時期
    其中,谷歌開源了著名的 TensorFlow 框架,它至今仍是 ML 領域最流行的深度學習框架。Caffe 的發明者加入了 Facebook 並發布了 Caffe2;與此同時,Facebook AI 研究(FAIR)團隊也發布了另一個流行的框架 PyTorch,它基於 Torch 框架,但使用了更流行的 Python api。微軟研究院開發了 CNTK 框架。
  • 陳天奇:機器學習科研的十年
    MSRA 回來之後,在實驗室薛老師的建議下,我選擇了一個現在看來正確而又錯誤的方向 -- 深度學習。那是 AlexNet 出現之前兩年,深度學習的主流熱點是非監督學習和限制玻爾茲曼機。沒有導師的指導,沒有工具,當時我靠著實驗室的兩塊顯卡和自己寫的 CUDA 代碼開始了死磕深度學習的兩年半。實驗室的學長問我,你準備要幹啥,我說:「我要用卷積 RBM 去提升 ImageNet 的分類效率。」
  • 拜拜虛擬機,全面深度解密華為方舟編譯器
    Java:1995年5月發布。SUN公司為了在電視機頂盒等電子產品智能化過程中搶佔先機,解決跨平臺的問題,所以開發了Java語言。但Java需要藉助虛擬機機制來解釋原始碼並調度硬體資源。安卓系統使用的基礎語言就是Java。C#:2000年6月,由微軟發布。
  • 基於動態編譯(Just-in-Time)的全新深度學習框架
    Jittor 是一個採用元算子表達神經網絡計算單元、完全基於動態編譯(Just-in-Time)的深度學習框架。[1] 據介紹,Jittor 內部使用創新的元算子和統一計算圖的深度學習框架。和 Numpy 相比,元算子在保證易用性的同時,能夠實現更複雜、更高效的操作。
  • 深度作業系統20.1正式發布:Kernel 5.8內核 性能優化
    深度作業系統20.1(1010)正式發布,採用全新Kernel 5.8(Stable)內核,更新倉庫到Debian 10.6,提升系統穩定性和兼容性。 深度作業系統20.1針對啟動時間、安裝時間、資源佔用、啟動響應等方面優化,深度全家桶應用更新。
  • 要理解深度學習,必須突破常規視角去理解優化
    深度學習算法有一些重要的特性並不總是反映在目標值中。所以,要加深對深度學習的理解,還得超越常規視角。深度學習時代,機器學習通常歸結為首先為手頭的學習任務定義合適的目標/成本函數,然後使用梯度下降的某種變體(通過反向傳播實現)來優化這個函數。難怪每年有數以百計的機器學習論文貢獻給優化的不同領域。