陳天奇團隊新研究：自動優化深度學習工作負載

新聞 05-24

新智元編譯

來源：arxiv.org

翻譯：肖琴

【新智元導讀】華盛頓大學計算機系博士生陳天奇、以及上海交通大學和復旦大學的研究團隊提出一個基於學習的框架，以優化用於深度學習工作負載的張量程序。該研究使用基於機器學習的方法來自動優化張量運算核心並編譯AI工作負載，從而可以將最優的性能部署到所有硬體。實驗結果表明，該框架能夠為低功耗CPU，移動GPU和伺服器級GPU提供與最先進手工調優庫相媲美的性能。

論文地址：https://arxiv.org/pdf/1805.08166.pdf

陳天奇團隊新研究：自動優化深度學習工作負載

深度學習在我們的日常生活中已經無處不在。深度學習模型現在可以識別圖像，理解自然語言，玩遊戲，以及自動化系統決策（例如設備放置和索引）。張量算符（tensor operators），如矩陣乘法和高維卷積，是深度學習模型的基本組成部分。

可擴展的學習系統依賴於手動優化的高性能張量操作庫，如cuDNN。這些庫針對較窄範圍的硬體進行了優化。為了優化張量算符，程序員需要從邏輯上等價的許多實現中進行選擇，但由於線程，內存重用， pipelining和其他硬體因素的不同，性能上的差別很大。

支持多種硬體後端需要巨大的工程努力。即使在當前支持的硬體上，深度學習框架和模型的開發也從根本上受到庫中優化操作符設置的限制，阻止了諸如操作符熔合（operator fusion）之類的優化，從而產生不受支持的操作符。

針對這個問題，華盛頓大學計算機系博士生陳天奇、以及上海交通大學和復旦大學的研究團隊提出一個基於學習的框架，以優化用於深度學習工作負載的張量程序（ tensor programs）。

陳天奇團隊新研究：自動優化深度學習工作負載

摘要

我們提出一個基於學習的框架，以優化用於深度學習工作負載的張量程序（ tensor programs）。矩陣乘法和高維卷積等張量算符（ tensor operators）的高效實現是有效的深度學習系統的關鍵。然而，現有的系統依賴於手工優化的庫，如cuDNN，這些庫只有很少的伺服器級GPU能很好地支持。對硬體有要求的操作庫的依賴限制了高級圖形優化的適用性，並且在部署到新的硬體目標時會產生巨大的工程成本。我們利用學習來消除這種工程負擔。我們學習了領域特定的統計成本模型，以指導在數十億可能的程序變體上搜索張量算符的實現。我們通過跨工作負載的有效模型遷移來進一步加快搜索速度。

實驗結果表明，我們的框架能夠為低功耗CPU，移動GPU和伺服器級GPU提供與最先進手工調優庫相媲美的性能。

學習優化張量程序問題的形式化方法

我們提出以下問題：我們是否可以通過學習來減輕這種工程負擔，並自動優化給定硬體平台的張量算符程序？本論文為這個問題提供了肯定的答案。我們建立了統計成本模型來預測給定的低級程序的程序運行時間。這些成本模型指導了對可能程序空間的探索。我們的成本模型使用可遷移的表示形式，可以在不同的工作負載之間進行泛化，以加速搜索。這一工作的貢獻如下：