面向航天異構平臺的深度學習編譯器加速技術優化

劉功晗 李 悅 王曉玲

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

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

0 引言

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

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

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

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

1 相關技術

1.1 DL編譯器

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

圖1 DL編譯器架構圖

1.2 Ansor

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

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

圖2 Ansor 在TVM中的環節

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

圖3 Ansor編譯流程

1) 任務調度：將完整的計算圖劃分為多個子圖，對有潛力的子圖進行重點優化。

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

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

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

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

2 成本模型設計

2.1 數據集

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

表1 Ansor提取的部分特征

2.2 模型設計

圖4 LightGBM成本模型的整體設計

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

表2 模型重要參數及優化結果

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

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

3 實驗評估

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

3.1 優化時間

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

表3 優化時間比較

圖6 RestNet-18網絡優化時間比較

3.2 推理時間

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

圖7 代碼性能比較

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

圖8 推理時間比較

4 結束語

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