優(yōu)化用于深度學(xué)習(xí)工作負(fù)載的張量程序

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

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

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

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

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

摘要

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

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

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

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

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

我們提出了一個基于機(jī)器學(xué)習(xí)的框架來解決這個新問題。

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

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

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

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

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

學(xué)習(xí)優(yōu)化張量程序算法

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

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

討論和結(jié)論

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

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

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