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

華盛頓大學(xué)計(jì)算機(jī)系博士生陳天奇、以及上海交通大學(xué)和復(fù)旦大學(xué)的研究團(tuán)隊(duì)提出一個(gè)基于學(xué)習(xí)的框架，以優(yōu)化用于深度學(xué)習(xí)工作負(fù)載的張量程序。該研究使用基于機(jī)器學(xué)習(xí)的方法來(lái)自動(dòng)優(yōu)化張量運(yùn)算核心并編譯AI工作負(fù)載，從而可以將最優(yōu)的性能部署到所有硬件。實(shí)驗(yàn)結(jié)果表明，該框架能夠?yàn)榈凸?a href="http://www.wenjunhu.com/v/tag/132/" target="_blank">CPU，移動(dòng)GPU和服務(wù)器級(jí)GPU提供與最先進(jìn)手工調(diào)優(yōu)庫(kù)相媲美的性能。

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

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

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

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

摘要

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

實(shí)驗(yàn)結(jié)果表明，我們的框架能夠?yàn)榈凸腃PU，移動(dòng)GPU和服務(wù)器級(jí)GPU提供與最先進(jìn)手工調(diào)優(yōu)庫(kù)相媲美的性能。

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

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

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

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

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

我們?cè)谶@個(gè)框架中提供了詳細(xì)的組件設(shè)計(jì)選擇和實(shí)證分析。

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

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

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

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

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

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

討論和結(jié)論

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

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

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