Facebook發布Tensor Comprehensions：自動編譯高性能機器學習核心的C＋庫

選自facebook research

機器之心編譯

機器之心編輯部

今天，Facebook 人工智慧實驗室宣布發布 Tensor Comprehensions，這是一個 C++庫和數學語言，它能幫助縮小使用數學運算的研究人員和專註在各種硬體後端運行大規模模型的工程師之間的距離。Tensor Comprehensions 的主要區別特徵是它使用一種獨特的準時化（Just-In-Time）編譯來自動、按需生成高性能的代碼，這正是機器學習社區所需要的。

生產力的數量級增長

創造全新高性能機器學習模型的典型工作流需要花費數天或者數周來完成兩個流程：

在 NumPy 層級的使用上，研究人員編寫一個全新層，並在 PyTorch 這樣的深度學習庫中鏈接已有運算，然後在小規模實驗上進行測試。代碼實現的表現需要通過數量級的加速來運行大規模試驗；

然後工程師使用該層，並為 GPU 和 CPU 編寫有效的代碼：

工程師需要成為高性能計算的專家，現在這方面的人才極為有限；

工程師需要明了這個環境，籌划出相關策略，編寫代碼以及做 debug；

將代碼移到與實際任務相關的後端（例如檢查冗長的評論），並添加樣板融合代碼。

結果導致，過去幾年深度學習社區一直依賴 CuBLAS、MKL、CuDNN 這樣的高性能庫來獲得在 GPU 和 CPU 上的高性能代碼。想要實驗不依賴於這些庫的新思路需要一定量級的工程量，這對研究人員來說可能是驚人的。

開源一種能將這一流程從數天或者數周縮減到數分鐘的工具包，我們覺得有非常大的實用價值。有了 Tensor Comprehensions，我們想像的是研究人員能夠以數學符號的方式編寫自己的想法，這種符號能被我們的系統自動編譯並調整，且結果是有很好表現的專業代碼。

在此發布中，我們能提供：

一種以簡單語法形式表達大量機器學習 idea 的數學符號；

一個基於 Halide IR 的 C++前端，面向此數學符號；

一個基於 Integer Set Library（ISL）的多面準時化（polyhedral Just-in-Time /JIT）編譯器；

一個基於進化搜索的多線程、多 GPU 的自動調節器。

早期研究

近期，高性能圖像處理領域中一種語言逐漸開始流行，即 Halide。Halide 使用類似的高級函數式語法描述圖像處理流程，接著在獨立的代碼塊中，將其明確調度到硬體上，詳細說明操作是如何被平鋪、矢量化、並行化和融合的。對於具有架構技能的人來說，Halide 是一個非常有生產力的語言，但是卻難以為絕大多數機器學習從業者所用。Halide 的自動調度是一個活躍的研究領域，但對於 GPU 上運行的 ML 代碼還沒有很好的解決方案。

Tensor Comprehension 將 Halide 編譯器作為所要調用的庫。我們構建了 Halide 的中間表徵（intermediate representation/IR）和分析工具，並將其與多面編譯（polyhedral compilation）技術配對，因此你可以使用相似的高階句法編寫層，而無需搞懂其運行原理。我們也發現了使語言更加簡潔的方法，而無需為縮減運算指定循環邊界。

細節

Tensor Comprehensions 使用 Halide 和多面編譯（Polyhedral Compilation）技術通過委託內存管理與協調自動合成 CUDA 內核。該編譯可對一般的運算符混合、快速局部內存、快速縮減和 JIT 專業化進行優化。由於我們並不嘗試獲取所有權或優化內存管理，因此我們的工作流可以簡單而高效地集成到任何能調用 C++函數的 ML 框架和語言中。

與經典編譯器技術和函數庫所採用的方法相反，多面編譯允許 Tensor Comprehensions 為每個網路按需調度單個張量元素的計算。在 CUDA 層級，它結合了仿射循環轉換、fusion/fission 和自動並行化，且能同時確保數據正確地流經內存層級。

圖中的數字表示最初計算張量元素的順序，箭頭表示它們之間的依賴關係。在該案例中，圖像的旋轉對應著允許深層運算符混合的循環交換。

為了推動搜索過程，我們同樣提供了一個集成的多線程、多 GPU 自動調優的庫，它使用進化搜索來生成和評估數千種實現方案，並選擇性能最好的方案。只要在 Tensor Comprehension 上調用 tune 函數，就能實時地查看性能提升，並在滿意的時候終止進程。最好的策略是通過 protobuf 執行序列化，並立即或離線的情況下復用。

在性能方面，雖然我們在工作上還有很多地方需要提升，但 Tensor Comprehensions 已經能相匹配或勝過當前 ML 框架與手動調參庫所集成的性能。這主要通過將代碼生成策略適應於特定問題而解決。下面的條形圖說明了將 Tensor Comprehensions 自動生成的內核與 Caffe2 和 ATen 等相對比時得到的性能提升。更多性能提升的細節，請查看以下 Tensor Comprehensions 的研究論文。

隨著我們擴大對更多硬體終端的貢獻，Tensor Comprehension 將補足由英偉達和英特爾編寫的快速庫，並將與 CUDNN、MKL 或 NNPack 等庫聯合使用。

下一步計劃

該工具的出現讓研究者與程序員們可以使用符號編寫層，這種方式與論文中使用的，用以描述程序的簡潔數學表達方式相同。這意味著人們可以快速方便地將這種表達轉化為實現，此前需要數天的任務目前僅需數分鐘的時間。

FAIR 將在未來放出 Tensor Comprehensions 的 PyTorch 集成版本。

FAIR 致力於向科學領域作出貢獻，並一直積極與機器學習社區合作。Tensor Comprehensions 已經在 Facebook、India、ETH Zurich 和 MIT 開始了應用。目前，這項工作還處於開發的初始階段，FAIR 將在未來對其進行進一步改進。

論文：Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions

論文地址：https://arxiv.org/abs/1802.04730

摘要：卷積和循環模型的深度學習網路已經相當普及，可以分析大量的音頻、圖像、文本和圖表數據，並應用於自動翻譯、語音到文本轉換、場景理解、用戶偏好排序、廣告定位等。用於構建這些網路的深度學習框架，如 TensorFlow、Chainer、CNTK、Torch/PyTorch、Caffe1/2、MXNet 和 Theano，都在可用性和表達性、研究或產品導向以及硬體支持之間探索不同的權衡。它們在運算符的 DAG 上運行；封裝了高性能的庫，如 CUDNN（NVIDIA GPU）或 NNPACK（多種 CPU）；以及自動化的內存分配、同步化和分散式。當計算無法調用已有的高性能庫時，需要使用自定製的運算符，這通常需要很高的工程代價。當研究者發明了新的運算符時，這很常見：而這種運算符會使得運行性能大幅降低，導致其創新意義受限。此外，即使存在這些框架可以使用的運行時調用庫，對於用戶的特定網路和數據集也無法提供最優的性能，因為它沒有優化運算符之間的關係，也不存在對數據的大小和形狀進行優化的工具。我們的貢獻包括了（1）構建了一種和深度學習的數學過程很接近的語言，即 Tensor Comprehensions，提供了命令式和聲明式的語言；（2）構建了一個多面即時（JIT）編譯器，以將表徵深度學習的有向無環圖數學描述轉換成一個具備內存管理和同步化功能的 CUDA 核，並且提供了運算符融合和指定特定數據大小等優化功能；（3）構建了一個由模組填充的高速編譯緩存。特別是，我們證明了多面框架可以有效地針對 GPU 上的當前最佳深度學習模型構建領域特定的優化器。我們的工作流相比將英偉達庫應用在機器學習社區常用的的核，以及應用在 Facebook 的真實產品級模型，要快了 4 倍。它集成了主流的框架 Caffe2（產品導向）、PyTorch（研究導向），以及 ATen 非共時張量庫。

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

------------------------------------------------

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