優(yōu)化用于深度學習工作負載的張量程序

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

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

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

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

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

摘要

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

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

學習優(yōu)化張量程序問題的形式化方法

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

我們提供了學習優(yōu)化張量程序問題的一種形式化方法，并總結了其關鍵特征。

我們提出了一個基于機器學習的框架來解決這個新問題。

我們使用遷移學習將優(yōu)化速度進一步提高2倍至10倍。

我們在這個框架中提供了詳細的組件設計選擇和實證分析。

在實際的深度學習工作負載的實驗結果表明，我們的框架提供的端到端性能改進比現(xiàn)有框架好1.2倍至3.8倍。

圖1：該問題的一個例子。 對于給定的張量算符規(guī)范 ，有多種可能的低級別程序?qū)崿F(xiàn)，每種實現(xiàn)都有不同的loop順序， tiling 大小以及其他選項。每個選項都創(chuàng)建一個具有不同性能的邏輯等效程序。我們的問題是探索程序空間并找到一個優(yōu)化的程序。

圖2：學習優(yōu)化張量程序框架的概覽

學習優(yōu)化張量程序算法

圖3：編碼低級別循環(huán)AST的可能方法的示例

表1：單batch的ResNet-18推理中所有conv2d操作符的配置。H，W表示高度和寬度，IC表示輸入通道，OC表示輸出通道，K表示 kernel大小，以及S表示stride大小。

討論和結論

我們提出了一種基于機器學習的框架來自動優(yōu)化深度學習系統(tǒng)中張量算符的實現(xiàn)。我們的統(tǒng)計成本模型允許在工作負載之間進行有效的模型共享，并通過模型遷移加速優(yōu)化過程。這個新方法的優(yōu)秀實驗結果顯示了對深度學習部署的好處。

在我們的解決方案框架之外，這個新問題的具體特征使它成為相關領域創(chuàng)新的一個理想測試平臺，如神經(jīng)程序建模、貝葉斯優(yōu)化、遷移學習和強化學習。

在系統(tǒng)方面，學習優(yōu)化張量程序可以使更多的融合操作符、數(shù)據(jù)布局和數(shù)據(jù)類型跨不同的硬件后端。這些改進對于改進深度學習系統(tǒng)至關重要。我們將開放我們的實驗框架，以鼓勵在這些方向進行更多的研究。