劉功晗 李 悅 王曉玲

1.宇航智能控制技術(shù)國家級重點實驗室，北京 100854

2.北京航天自動控制研究所，北京 100854

0 引言

目前，航天系統(tǒng)正向智能化方向邁進，各種智能算法的應(yīng)用在航天裝備中取得了優(yōu)異的表現(xiàn)。比如，在導(dǎo)彈突防控制中有學(xué)者提出應(yīng)用深度強化學(xué)習(xí)技術(shù)，設(shè)計了一種深度神經(jīng)網(wǎng)絡(luò)(Deep Neural Network，DNN)模型架構(gòu)，實現(xiàn)了對彈道中段突防最優(yōu)控制模型的逼近[1]。在目標(biāo)檢測中，通過Faster RCNN及其改進算法為遙感任務(wù)圖像中復(fù)雜地理環(huán)境背景下的目標(biāo)檢測方法提供新思路[2]，在故障檢測中，通過基于雙向長短記憶神經(jīng)網(wǎng)絡(luò)的故障診斷模型，提高了對于發(fā)動機故障診斷的準(zhǔn)確率[3]。然而，以DNN模型為例，其優(yōu)異表現(xiàn)是以高計算復(fù)雜性為代價的，尤其是對實時性有著高要求的航天領(lǐng)域，更需要通過調(diào)用底層加速庫來提高DNN模型的實現(xiàn)效率，如PyTorch和TensorFlow采用人工優(yōu)化內(nèi)核如Intel MKL-DNN或是Nvidia cuDNN作為后端。然而，隨著新興算子的不斷增加、網(wǎng)絡(luò)模型的不斷擴大，人工優(yōu)化的加速庫很難適配不斷增加的DNN模型，進而導(dǎo)致其可能無法滿足航天裝備對于實時性的要求，也極大阻礙了DNN模型在航天裝備異構(gòu)平臺上的部署，限制了航天智能技術(shù)的進一步發(fā)展。

為了解決這一難題，我們通過引入深度學(xué)習(xí)編譯器TVM[4]來保證航天裝備對于實時性的高要求，提高部署的靈活性。Ansor[5]是一個被集成到TVM中的自動搜索框架，為深層神經(jīng)網(wǎng)絡(luò)生成高性能的張量程序。編程人員只需要提供網(wǎng)絡(luò)模型，Ansor會自動搜索出一套調(diào)度策略，再通過TVM的代碼生成，就可以得到優(yōu)化后的代碼，實現(xiàn)與人工優(yōu)化庫相似甚至更好的性能。

雖然Ansor可以生成高度優(yōu)化的DNN模型代碼，但這個過程可能要花費很長的優(yōu)化時間。例如，在ResNet-18中，通過Ansor生成的代碼性能優(yōu)于PyTorch，但是高性能代碼的生成可能需要幾個小時甚至十幾個小時。隨著當(dāng)前深度學(xué)習(xí)的研究不斷深入，網(wǎng)絡(luò)模型的參數(shù)也到達了數(shù)百萬甚至數(shù)十億的規(guī)模，這使得減少TVM的優(yōu)化時間的問題變得十分突出。此外，隨著新的神經(jīng)網(wǎng)絡(luò)結(jié)構(gòu)[6-7]的快速出現(xiàn)，以及在不同的航天場景下硬件平臺的選擇，使得相關(guān)人員被迫更頻繁地優(yōu)化網(wǎng)絡(luò)。過長的優(yōu)化時間給優(yōu)化過程帶來了極大的不便，甚至使當(dāng)前基于編譯器的解決方案受到了質(zhì)疑。

通過研究我們發(fā)現(xiàn)，TVM使用成本模型來預(yù)測每個程序在真實硬件上的運行時間，再通過搜索算法選擇其中的最優(yōu)解。以Ansor為例，它采用XGBoost[8]作為成本模型進行預(yù)測，因此，XGBoost預(yù)測的速度，對整個優(yōu)化時間會產(chǎn)生重要影響。本文首先對相關(guān)的技術(shù)原理進行介紹，然后設(shè)計了以LightGBM (Light Gradient Boosting Machine)[9]為預(yù)測方法的成本模型，替代原始的XGBoost模型，這是一種旨在得到類似或更好的優(yōu)化質(zhì)量但是優(yōu)化時間更短的方法。最后通過實驗，證明我們所提出的方法是更優(yōu)的，它使我們能夠以最高快1.6倍的速度優(yōu)化DNN模型，而且得到了類似甚至更優(yōu)的推理時間，打破了TVM在航天裝備中的應(yīng)用瓶頸。

1 相關(guān)技術(shù)

1.1 DL編譯器

目前，像TVM這樣的DL編譯器因為可以自動優(yōu)化DL程序而變得流行。與傳統(tǒng)編譯器類似，典型的DL編譯器就是將DNN模型抽象看成各種編程語言，并對其進行優(yōu)化，最后生成后端可執(zhí)行的高性能代碼的過程。如圖1所示，DL編譯器的結(jié)構(gòu)主要包含2部分：編譯器前端和編譯器后端，中間表示(IR)橫貫前端和后端，IR是程序的抽象，用于程序的優(yōu)化。DNN模型在DL編譯器中轉(zhuǎn)換為多級IR(Intermediate Representation)，分為高級IR和低級IR?；诟呒塈R，編譯器前端負責(zé)獨立于硬件的轉(zhuǎn)化和優(yōu)化；基于低級IR，編譯器后端負責(zé)特定于硬件的優(yōu)化、代碼生成和編譯。

圖1 DL編譯器架構(gòu)圖

1.2 Ansor

TVM最初需要編程人員手動或半手動的指定調(diào)度策略，然后通過TVM的代碼Pipeline進行代碼生成，但是通過手寫或采用模板的優(yōu)化方式可能會陷入局部最優(yōu)，同時，想要寫出一個好的調(diào)度策略，編程人員需要對硬件有深入的理解。

為了解決上述問題，TVM引入了Ansor，這樣，編程人員只需要提供模型，Ansor會自動選擇出一套調(diào)度策略，再通過TVM進行代碼生成，就可以得到優(yōu)化后的代碼，如圖2所示。

圖2 Ansor 在TVM中的環(huán)節(jié)

Ansor的代碼生成流程可以分為以下幾個步驟，如圖3所示：

圖3 Ansor編譯流程

1) 任務(wù)調(diào)度：將完整的計算圖劃分為多個子圖，對有潛力的子圖進行重點優(yōu)化。

2) 草圖構(gòu)建：提取算子中的高層次特征，對算子進行粗粒度的優(yōu)化，確定代碼的基本結(jié)構(gòu)。

3) 隨機注釋：隨機初始化Tilling Size和一些for循環(huán)的策略，獲得計算圖的完整表示。

4) 搜索評估：訓(xùn)練成本模型，根據(jù)成本模型對代碼的性能進行評估，選取評估中高分數(shù)的一組實現(xiàn)，獲得其在真實硬件測量中的正確結(jié)果，選取實際性能最優(yōu)的實現(xiàn)作為Ansor的輸出。

總結(jié)來說，任務(wù)調(diào)度模塊劃分DNN模型的計算子圖，確定各個子圖的優(yōu)化次數(shù)；草圖構(gòu)建模塊搭建代碼的整體框架；隨機注釋模塊對細節(jié)進行填充；搜索評估模塊訓(xùn)練一個模型，對生成的代碼進行評價和更新，獲得最后的代碼生成結(jié)果。

2 成本模型設(shè)計

2.1 數(shù)據(jù)集

如1.2節(jié)所述，隨機注釋模塊隨機生成代碼，但不考慮代碼的性能。因此，待優(yōu)化的計算可能十分復(fù)雜，遍歷所有的隨機注釋模塊來獲取最優(yōu)實現(xiàn)是不現(xiàn)實的。用成本模型預(yù)測隨機注釋模塊生成代碼的性能，需要通過提取完整程序的上下文中的特征為程序最里面的非循環(huán)語句構(gòu)建特征向量，提取的特征包括算術(shù)特征和內(nèi)存訪問特征，這些特征可以通過遍歷TVM的TIR獲取，使用獨熱碼對其進行編碼，包含一個語句的所有列出特征的特征向量的長度是164。將這些特征向量作為成本模型的輸入數(shù)據(jù)，提取的部分特征如表1所示。

表1 Ansor提取的部分特征

2.2 模型設(shè)計

圖4 LightGBM成本模型的整體設(shè)計

LightGBM算法使用帶深度限制的leaf-wise的葉子生長策略，與深度增長算法相比，葉向增長算法收斂速度更快。然而，如果模型參數(shù)設(shè)置不當(dāng)，葉向生長可能會出現(xiàn)過擬合現(xiàn)象，導(dǎo)致預(yù)測精度降低。本文采用Scikit-Learn的GridSearchCV函數(shù)實現(xiàn)參數(shù)搜索工作。在指定的參數(shù)范圍內(nèi)，按照步長依次調(diào)整參數(shù)，遍歷所有可能的參數(shù)組合，挑選驗證集誤差最小的超參數(shù)組合返回，最終結(jié)果如表2所示。

表2 模型重要參數(shù)及優(yōu)化結(jié)果

Ansor的搜索評估模塊使得優(yōu)化能夠更快地收斂，減少了遍歷時間，在采用LightGBM后的搜索評估模塊整體流程如圖5所示。其中，可以看到特征提取的輸入有3個來源：1) 隨機注釋模塊的樣本，保證了樣本的豐富性；2) 之前迭代中獲取的注釋模塊的優(yōu)秀樣本，有助于訓(xùn)練得到準(zhǔn)確度更高的LightGBM模型；3) 優(yōu)秀樣本的變種，根據(jù)進化算法，好的樣本“突變”往往會更容易得到好的樣本，因此采取對好樣本的參數(shù)進行微調(diào)的方式擴充樣本集。提取的特征用來訓(xùn)練LightGBM模型,得到程序p及其對應(yīng)的性能Perf(p)，同時該模型又能對注釋模塊產(chǎn)生的樣本進行一次初篩選，將評分高的樣本通過TVM 運行時模塊獲得相應(yīng)的正確結(jié)果Runtime (Perf(p))，根據(jù)正確值可以選擇出性能更優(yōu)的代碼，即整個Ansor的輸出，整個算法至此結(jié)束。該輸出可以無縫轉(zhuǎn)換成調(diào)度策略，省去了手寫調(diào)度策略的麻煩，并直接通過TVM 代碼生成模塊 (TVM CodeGen) 進行代碼的生成。

圖5 搜索評估模塊整體流程

3 實驗評估

在本節(jié)中，通過實驗評估LightGBM模型，尋找它能夠有效加速優(yōu)化時間的答案。用Python編寫LightGBM模型，利用scikit-learn實現(xiàn)模型的參數(shù)優(yōu)化。從優(yōu)化時間和推理時間2個方面評估。在優(yōu)化時間方面，分別比較了Ansor使用XGBoost和LightGBM兩種成本模型在搜索高性能代碼時花費的時間。在推理時間方面，從局部的代碼生成到整體的網(wǎng)絡(luò)推理進行比較。在下面的內(nèi)容中，把Ansor利用XGBoost算法的成本模型簡稱為xgb，利用LightGBM算法的成本模型簡稱為lgb。

3.1 優(yōu)化時間

表3展示了Ansor使用2種不同的成本模型在GPU上的優(yōu)化時間。分別比較了xgb和lgb在ResNet-18，VGG16，和SqueezeNet1_1上端到端的優(yōu)化時間。在表3中，“num_trials”是在調(diào)優(yōu)過程中的迭代次數(shù)，適當(dāng)增加這個值有利于搜索的充分收斂。對于ResNet-18，Ansor將其總共劃分成18個任務(wù)進行調(diào)優(yōu)，圖6展示了分別使用lgb和xgb調(diào)優(yōu)該網(wǎng)絡(luò)所花費的時間。由表3和圖6，總體而言，lgb實現(xiàn)了1.6倍的提速。這是因為通過使用LightGBM模型，減少了特征的數(shù)量，過濾掉了小梯度的樣本，并將遍歷方式從遍歷樣本變?yōu)楸闅v直方圖的形式，這些方法有效降低了時間復(fù)雜度，減少了大量不必要的計算。

表3 優(yōu)化時間比較

圖6 RestNet-18網(wǎng)絡(luò)優(yōu)化時間比較

3.2 推理時間

首先，分別對使用xgb和lgb生成代碼的性能進行比較，如圖7所示，其中y軸表示調(diào)優(yōu)后的代碼在真實硬件中測量得到的GFLPOS，x軸表示迭代次數(shù)。分別選取2個任務(wù)：圖7(a)是在CPU上對矩陣乘法進行調(diào)優(yōu)，圖7(b)是在GPU上，選取了ResNet-18中最后一層的卷積進行調(diào)優(yōu)，為使2個任務(wù)均能充分收斂，取num_trials=1000. 最終結(jié)果如圖7所示?？梢钥闯觯菏紫?，使用lgb調(diào)優(yōu)后的代碼性能與xgb類似，有時甚至可以超過xgb，證明了lgb在縮短優(yōu)化時間的同時保證了生成代碼的質(zhì)量。其次，通過lgb進行預(yù)測時，收斂速度更快。更快的收斂速度對于調(diào)優(yōu)DNN模型以獲得更好的性能至關(guān)重要。

圖7 代碼性能比較

圖8比較了在GPU上ResNet-18,VGG16和SqueezeNet1_1三種網(wǎng)絡(luò)的推理時間，用目前最流行的深度學(xué)習(xí)框架PyTorch作為測量的基線?？偟膩碚f，lgb的推理速度比xgb快了6.1%，比PyTorch快90.1%。這一方面是由于兩者決策樹的生成方式不同，XGBoost算法采用按層生長的策略(level-wise)，不加區(qū)分的對待同一層葉子，使得部分收益很低的葉子也進行了搜索和分裂。LightGBM使用帶有深度限制的葉子生長算法(leaf-wise)。因此在分裂次數(shù)相同的情況下，LightGBM可以降低更多的誤差，得到更好的精度。另一方面，lgb和xgb在SqueezeNet上的推理時間顯著降低，這或許是因為通過Ansor自動搜索的策略與PyTorch相應(yīng)的加速庫相比，更適配底層硬件。

圖8 推理時間比較

4 結(jié)束語

隨著航天智能技術(shù)的不斷發(fā)展，對于DNN模型的需求也不斷增加，TVM的出現(xiàn)滿足了航天裝備對實時性的要求，有效降低了DNN模型在航天異構(gòu)平臺上部署的難度。然而，雖然高度優(yōu)化的代碼可以通過TVM實現(xiàn)，但它們需要生成不同的張量程序，并在真實硬件上測量，才能得到最優(yōu)的結(jié)果，導(dǎo)致優(yōu)化時間過長的問題。在本文中，我們建議將LightGBM模型集成到Ansor中，該模型在保證精度的情況下，減少了優(yōu)化過程的時間成本。因此，對于在不同的DNN模型中搜索高性能的張量程序，使用帶有LightGBM的Ansor花費時間更少，優(yōu)于之前基于XGBoost的成本模型，是一種更加高效的方法。該方法使得在航天異構(gòu)平臺中，基于TVM解決實時性和靈活性的問題成為一種可能。