作者：王博文 | 曠視 MegEngine 架構師

一、背景

對於深度學習框架來說，網路的訓練/推理時間是使用者非常看中的。在實際生產條件下，使用者設計的 NN 網路是千差萬別，即使是同一類數學計算，引數也各不相同。如果沒有針對性的優化，框架就完全喪失競爭力。因此，在一類數學計算中，開發者們會開發多種高效的演算法，分別適用於不同的引數，以保證網路的效能。接下來開發者們需要解決一個新問題，當計算引數確定以後，如何讓最快的演算法執行該計算。

大部分框架靠先驗的經驗選擇演算法，MegEngine 亦總結有優秀的先驗經驗值，實現計算時自動選擇演算法。但是依靠經驗不能保證一定選擇了最快的演算法。很多實際場景中，使用者希望網路有最極致的效能。為此，MegEngine 設計了專門的流程，可以為每個計算自動選擇最快的演算法，從而保證整個網路的執行時間最短。並且同時能夠將計算的引數和其對應的演算法資訊以及裝置資訊記錄到記憶體或檔案，當用戶再次執行網路時，可以直接獲取效能最好的演算法。這一提升效能的流程被稱為 Fast Run，它能讓 MegEngine 的使用者執行不同的網路時都能收穫最好的效能。

二、Fast Run 簡述

目前，主流的框架幾乎都使用了運算元（Operator）的概念來抽象數學計算，如卷積運算元，矩陣乘運算元等。MegEngine 也使用了運算元 這一概念。此外，在底層，我們開發了名為 MegDNN 的計算庫，用以完成實際的數學計算。 MegDNN 僅提供數學計算能力。MegDNN 的頂層也是按照運算元的概念組織的，對不同的後端，分別封裝了 MegDNN 運算元。一個 MegDNN 運算元內部則可能有多個該運算元的演算法，MegEngine 將演算法抽象為 Algorithm，一個 Algorithm 物件可以完成該運算元的計算。

以卷積運算元為例，ARM 上，MegEngine 實現了非常通用的 Im2col 演算法，有特定條件下效能卓越的 Winograd 演算法，有在小尺寸卷積時高效能的 Direct 直接卷積演算法等。CUDA 上，有呼叫 cuDNN 庫函式的方法等等。從 MegEngine 運算元到 MegDNN 運算元再到演算法的關係如下圖所示：

一個 MegEngine 運算元可能持有一個或多個 MegDNN 運算元來完成計算，一個 MegDNN 運算元需要從多個演算法物件中選擇一個來執行計算。為了極致的計算效能，需要在開始網路計算之前，給 MegDNN 運算元選好最快的演算法。

Fast Run 的思路很直接，在網路計算開始之前，將每個 MegDNN 運算元中所有可行的演算法全部執行一次（Profiling），並將效能資料記錄下來，將最快的演算法設定給 MegDNN 運算元。Fast Run 成立的前提條件是演算法執行時間是穩定的，這樣比較每個演算法的 Profiling 資料才有意義。

最後是確定 Fast Run 執行的時間點。MegEngine 有統一的記憶體管理，各 MegEngine 運算元需要在計算開始前向記憶體規劃單元申請足夠的計算時記憶體，這一記憶體包括了其內部的 MegDNN 運算元計算時需要的記憶體，而 MegDNN 運算元計算時需要的記憶體完全由演算法決定。這就要求，MegDNN 此刻已經確定了將要使用的演算法。自然地，MegEngine 選擇在呼叫該介面之前執行 Fast Run 流程。這樣，當 Fast Run 流程完成時，各 MegDNN 運算元都設定了效能最好的演算法。

Fast Run 執行的代價是顯然的，它會顯著增加第一次網路執行的時間。Fast Run 的流程如下圖：

Fast Run 有下面兩種使用方式，區別在於上圖中寫入的 Cache 檔案不同：

  • 離線 Fast Run，離線 Fast Run 分兩步，分別在不同的程序中完成。第一步先將整個網路計算執行一遍，這一過程中，Fast Run 會將各個演算法的效能資料寫到一個專門的資料結構中，最後資料被統一寫入一個 Cache 檔案，隨後程序退出，這個過程稱之為“搜參”。第二步，載入同樣的網路，通過 MegEngine 的介面將 Cache 檔案讀入。可以看出，離線 Fast Run甚至可以在不同的裝置上進行。
  • 線上 Fast Run，線上 Fast Run 在同一個程序完成的。前半段與離線 Fast Run 的流程相同，Fast Run 後，各演算法的效能資料儲存在記憶體中的一個數據結構之中。此時，程序不會退出。後續可以給網路載入不同的輸入資料，此時各 MegDNN 運算元中已設定好效能最好的演算法。並且，也可以初始化另外的網路，亦可以像離線 Fast Run 的後半部分一樣，從當前的資料結構中讀取演算法。

總的來說，Fast Run 提供搜參和記錄的功能。它的作用是給網路中的各個 MegDNN 運算元選擇當前引數下效能最好的演算法。由於 Fast Run 對每個 MegDNN 運算元執行同樣的操作，因此它在前向推理和反向傳播時都能使用。目前，MegEngine 支援 CUDA、CPU、ROCM 三個後端的 Fast Run ，MegEngine 的使用者們在訓練和部署時，均廣泛使用 Fast Run。

三、Fast Run 原理

Fast Run 中，Profiling 一個 MegDNN 運算元並設定演算法，會經歷 4 個步驟，其流程如下圖示：

這一流程中，需要注意一些細節：

1、遞迴搜參：MegDNN 中普遍存在運算元巢狀的情況。例如，Convolution 運算元中，Im2col 演算法會使用 MegDNN 的 MatMul 運算元執行矩陣乘計算。那麼，Convolution 的效能直接受到 MatMul 效能的影響。可以看到，在 Profiling 一個 Convolution 運算元之前，需要 MatMul 運算元執行的效能資料已知。為了解決這個問題，Fast Run 使用了遞迴的方式，來解決搜參時的運算元巢狀問題。如上圖中虛線框所示，一個 MegDNN 運算元，在獲取所有可用演算法之後，會呼叫每個演算法的介面，詢問該演算法是否依賴子運算元並儲存相關結果，若最終相關結果不為空，則會先對子運算元進行一次 Profiling，此後，再 Profiling 頂層的運算元時，其使用的子運算元會有最優的演算法儲存在 Cache 中。

2、Fast Run 效能資料儲存：Fast Run 效能資料存取離不開 Cache。MegEngine 提供了兩種 PersistentCache，兩種 Cache 區別於資料儲存的位置（記憶體或是檔案）。Cache 的結構如下圖所示：

MegEngine 中，PersistentCache 物件是單例的，兩種 Cache 都保證執行緒安全。Cache 維護一個從 category 資訊到一個集合的對映的集合，此處 category 是一個後端的記錄資訊。Category 是一個字串，由後端資訊和運算元型別拼接獲得，後端資訊 由裝置區分，例如 CUDA 的後端資訊由裝置名稱、NVIDIA 驅動版本和 CUDA 執行時庫版本資訊組成；CPU 作為後端時，則只記錄裝置名稱。MegEngine 中只有 CUDA、CPU、ROCM 三種類型有對應的 categoty 生成，這也是 MegEngine 目前僅支援在 CUDA、CPU、ROCM 三個後端支援 Fast Run 的原因。運算元型別 由運算元名稱、Cache 版本資訊兩部分組成。

一個 category 對映到一個集合，該集合維護單個 MegDNN 運算元的資訊到其所有可用演算法的 Profiling 結果的對映。該集合的 key值 由 MegDNN 運算元的所有輸入 Tensor 的尺寸和運算元的全部引數組成（這些引數能夠完全決定一個演算法是否可用）。value值 是一個數組，儲存每個 Profiling 過的演算法的時間、所需額外的空間等資訊，並排序。排序時，以執行時間進行升序排列，並且保證了序列中每個演算法使用的記憶體必須小於其前一個演算法使用的記憶體 – 這樣序列中不存在一個演算法既慢於另一個演算法，又使用更多的記憶體。一個 Cache 中可以存在不同後端的 Fast Run 結果，只要它們的 category 不同。

在一些常見的模型上，推理時關閉和開啟 Fast Run，效能表現如下：

從工程落地中 Fast Run 的使用情況來看，絕大部分場景下，能顯著降低網路執行時間。

四、Fast Run 使用

MegEngine 可配置的引數眾多，很多都是工程落地的解決方法，在工業上經過大量的實踐。其中一些引數與 Fast Run 的使用有密切的關係，這裡詳細闡述它們的使用。

4.1 開啟 Fast Run

原始碼級別使用 Fast Run 可以參照 MegEngine 自帶的可執行程式 load_and_run，如果僅關注利用 load_and_run 測試模型，有下面兩個引數需要使用：

  1. --full-run/--fast-run，搜參的兩種模式，需使用者選擇其中一種模式，兩者的區別在於 Profiling 時，生成的 MegDNN 運算元的可用演算法集大小不同。--full-run 時，會 Profiling MegDNN 運算元內所有的可用演算法，包括最樸素的演算法（MegDNN 運算元至少有一個演算法，保證任何引數下均可用，執行慢）。--fast-run 則會排除樸素演算法。如果想要減少 Profiling 的時間開銷，可以選擇使用 --fast-run 模式，此時需要注意的是，如果網路中有引數過於特殊的運算元，則該運算元可能面臨沒有可用演算法的情況（優化過的演算法不可用、樸素的演算法被排除），此時 MegEngine 會報出“沒有可用演算法”的錯誤並退出。
  2. --fast-run-algo-policy，指定 Cache 檔案的路徑，檔案中的效能資料會被讀入記憶體，被全域性唯一的 PersistentCache 物件持有。程序退出前，PersistentCache 中的效能資料會全部寫入該檔案。

兩個引數可以單獨使用，也可以一起使用：

  1. 單獨使用 --full-run/--fast-run，Profiling 資料儲存在記憶體中。
  2. 兩者一起使用，檔案中的效能資料首先會被讀入記憶體。如果檔案為空，所有 MegDNN 運算元完成搜參後，效能資料寫回檔案。如果檔案不為空，且某個 MegDNN 運算元能從 Cache 中查詢到效能資料，則不會進行搜參，餘下不能查到效能資料的，則會搜參。這樣實現了斷點搜參的功能，MegEngine 稱之為“續搜“。如果 Fast Run 時程式因為某些原因異常退出，”續搜“能使 Fast Run 在下一次能夠連上。“續搜”也能讓多個模型的效能資料可以合併在一個 Cache 檔案中。如果所有 MegDNN 運算元都能從 Cache 中查到效能資料，則搜參不會發生，網路具有最好的效能。
  3. 單獨使用 --fast-run-algo-policy，檔案中的效能資料首先會被讀入記憶體，如果 Cache 中沒有記錄，不“續搜”，以經驗值設定 MegDNN 運算元的演算法，效能可能不是最優。

在使用 Fast Run 時，可以配合 --verbose 一起使用，程式將詳細列印 Fast Run 時的除錯資訊，包括 MegDNN 運算元的名稱，輸入輸出的尺寸資訊，設定的演算法名稱等。如果發現效能不符合預期，比如當載入的模型和 Cache 檔案不匹配時，通常會發生“續搜”，造成網路執行時間很長的假象。因此，我們強烈推薦在此時使用 --verbose 引數來觀察程式工作是否符合預期。

4.2 演算法屬性

MegDNN 中某些演算法具有獨特的屬性，會影響向 MegDNN 運算元設定演算法，當前使用的 屬性 有：

  1. REPRODUCIBLE：具有 REPRODUCIBLE 屬性的演算法，可保證計算結果位元對齊。Fast Run 中，在從 Cache 中讀演算法資訊時提供了對 REPRODUCIBLE 屬性的支援。設定 --reproducible，Fast Run 會從 Cache 中選擇效能最好的且具有 REPRODUCIBLE 屬性的演算法。在 Profiling 階段，並不區分演算法是否 REPRODUCIBLE，這樣 Cache 中的演算法既有 REPRODUCIBLE 屬性的，也有非 REPRODUCIBLE 屬性的，具備一定的泛用性。
  2. NAIVE：只有 MegDNN 中最樸素的演算法具有 NAIVE 屬性。--full-run 和 --fast-run 的區別就在於 --fast-run 通過該屬性篩除了執行最慢的樸素演算法。

4.3 weight 前處理

有些演算法，在計算時需要對資料進行輔助轉換。其中，對權重 weight 的轉換可以是一次性的，這樣可以節省執行時間。例如 Winograd 演算法，其權重可以在進行卷積計算之前轉好，節約相當一部分執行時的效能開銷。MegEngine 在 GraphCommonOptimizeOptions 中提供了 weight_preprocess 選項來支援部署時權重的提前轉換功能。一旦設定 weight_preprocess，對於那些 weight 能夠提前轉換的演算法，其效能資料將不會包含權重轉換的時間。簡單的說，在搜參階段設定 weight_preprocess，會影響演算法的效能資料，從而 Cache 中演算法的效能資料排序可能不同。如果 Cache 是在開啟 weight 前處理的情況下搜參得到，部署時務必要開啟 weight 前處理以獲得更好的效能，否則有效能下降的風險。Fast Run 與 weight 前處理不是必需的關係，兩者可以分開使用。不過通常情況下，兩者結合使用可以獲得更好的效能.

4.4 Fast Run 版本

Fast Run 的版本資訊以字串的形式表示在 Cache 的 category 中。Cache 具有相容性，可以允許不同的版本的 MegEngine 下的搜參結果集合在同一個 Cache 中，Cache 中看到的是不同的 category。但是使用者在使用過程，依然需要注意 Fast Run 的版本。一般地，如果 MegDNN 的演算法發生了刪除或者是屬性的變動，Fast Run 的版本資訊會發生變化。Fast Run 版本資訊變化後，需要重新搜參。