FAIR開源Tensor Comprehensions,讓機器學習與數學運算高性能銜接

..

雷鋒網AI科技評論消息,Facebook AI 研究院於近日開源了 C++ 庫及數學語言 Tensor Comprehensions,它能有效填補研究人員於數學運算領域的溝通鴻溝,並基於各種硬件後端上大規模運行工程模型。它的最大亮點在於,它採用了 Just-In-Time 的編譯自動生成機器學習社區所需的高性能代碼。

也就是說,通過 Tensor Comprehensions,研究人員能夠以數學符號的方式進行編寫,系統能夠根據需求進行編譯調整,並輸出專業的代碼。

數量級增長

傳統意義上,如果要從零創造一個具有高性能表現的機器學習層需要兩個大步驟,這可能需要耗費工程師數天甚至數周的努力。

  • 1. 在 NumPy 層級,研究人員需要寫一個全新層,並在以 PyTorch 為代表的深度學習庫中鏈接已有運算,然後進行小規模測試。如果要運行大規模實驗,需要通過數量級加速檢測其代碼實現。

  • 2. 工程師隨後採用這個層,並為 GPU 和 CPU 撰寫有效代碼(但這需要滿足多個條件):

    1)工程師需要是一位在高性能計算頗有了解的專家,但這一人才目前非常稀缺;

    2)文本、策略、寫代碼,debug,工程師需要樣樣精通;

    3)將代碼與實際任務相關的後端相連接,如冗長的參數檢查和添加樣板集成代碼

這也直接導致近年來深度學習社區一直依賴以 CuBLAS, MKL, 和 CuDNN 為代表的高性能庫而構建運行於 GPU 和 CPU 的代碼。對於研究員而言,如何尋找一條新的實現思路成為了巨大挑戰。

而在 Tensor Comprehensions 中,研究員們得以將這一流程從幾天甚至幾周減少到數分鐘。這一開源包含了:

• 用簡單語法表達一系列機器學習概念的數學符號

• 基於 Halide IR 數學符號的 C ++前端

• 基於整數集庫(ISL)的 Just-in-Time 編譯器,

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

早期工作

Halide 是一種最近在高性能圖像處理領域頗受歡迎的語言,它採用類似的高級函數語法來描述一個圖像處理的 pipeline,隨後在單獨代碼塊中調度到硬件上,並且詳細到如何平鋪、矢量化、并行化和融合。對於具有專業知識的人而言,這是一種非常高效的語言;但對於機器學習從業者來說,這一難度並不小。Halide 的自動調度在研究上非常活躍,但對於 GPU 上運行的機器學習代碼,目前還沒有很好的解決方案。

Tensor Comprehension 將 Halide 編譯器作為調用庫。FAIR 研究員構建了 Halide 的中間表徵(IR)和分析工具,並與多面編譯進行技術配對,因此,用戶可以在無需理解運行原理的條件下使用相似的高級語法編寫層。此外,FAIR 研究員也找到了簡化語言的方法,不需要為縮減運算制定循環邊界。

細節實現

使用 Halide 和多面編譯技術,Tensor Comprehensions 能通過委託內存管理和同步功能以自動合成 CUDA 內核。這一編譯能夠針對特定尺寸對一般操作符進行融合,對快速本地內存、快速縮減和 Just-in-Time 專業化都能優化。因不嘗試涉足內存管理,因此這一流程能夠輕鬆集成到機器學習的任意框架,以及任何允許調用 C++ 函數的語言中。

與經典編譯器技術和庫方法相反的是,多面編譯能讓張量理解為每個新網絡按需調度單個張量元素的計算。

在 CUDA 層面,它結合了仿射循環轉換、融合/分裂和自動并行處理,同時確保數據在存儲器層次結構中正確移動。

圖中數字錶示最初計算張量元素的順序,箭頭表示它們之間的依賴關係。在這個例子中,數字旋轉對應於循環交換,能夠實現深度操作器的融合。

研究員提供了一個集成的多線程、多 GPU 自動調節庫,以推進搜索過程,它使用 Evolutionary Search 來生成和評估數千種實現方案,並選擇性能最佳的方案。只需調用 Tensor Comprehension 上的 tune 功能,用戶便能目睹性能的改善。最好的策略是通過 protobuf 序列化,在當下或離線情況下可重用。

在性能方面,Tensor Comprehensions 可能仍然有很大提升空間,但在條件滿足的情況下,它已經可以匹配或打破當前機器學習框架與手動調整庫集成的性能,並通過將代碼生成策略適應特定問題大小的能力來實現。下圖展現了 Tensor Comprehensions 與 Caffe2 和 ATen 的對比。

更多信息請參考 arXiv 論文(鏈接見底部)。

未來,Tensor Comprehensions 將補充硬件製造商(如 NVIDIA 和 Intel)編寫的快速庫,並將與 CUDNN,MKL 或 NNPack 等庫一起使用。

下一步

雷鋒網AI科技評論了解到,該版本的開源庫能讓研究人員和程序員使用與他們在論文中使用的數學相似標記來編寫層,並簡明地傳達程序意圖。他們也能夠在幾分鐘內將這種表示方法進行編譯。未來隨着工具鏈的不斷增長,FAIR 認為它的預計可用性和性能將會提升,並受益整個社區。

晚些時間,FAIR 將會發佈 PyTorch 集成的 Tensor Comprehensions。

GitHub 鏈接:https://facebookresearch.github.io/TensorComprehensions/

arXiv 論文:https://arxiv.org/abs/1802.04730

Slack 頁面:https://tensorcomprehensions.herokuapp.com/

Email: tensorcomp@fb.com

雷鋒網AI科技評論編譯


想在手機閱讀更多程式設計資訊?下載【香港矽谷】Android應用
分享到Facebook
技術平台: Nasthon Systems