移動端 CPU 的深度學習模型推理效能優化——NCHW44 和 Record 原理方法詳解

MegEngine發表於2021-08-17
深度學習模型優化

使用者實踐系列，將收錄 MegEngine 使用者在框架實踐過程中的心得體會文章，希望能夠幫助有同樣使用場景的小夥伴，更好地瞭解和使用 MegEngine ~

作者：王雷 | 曠視科技 研發工程師

背景

隨著人工智慧技術的發展及應用領域的不斷擴大，算力較弱的移動裝置成為模型推理的重要運算載體，優化其推理效能因此成為重要的工程問題。一般認為，讓模型執行於 GPU 上會比執行於 CPU 上具有較大的優勢，取得可觀的效能提升。這通常是真實情況，但是，在工程實踐中我們也發現，對於某些模型維度較小的模型，在移動裝置上，GPU 執行並沒有帶來效能的提升，而且還額外引入了相容性的問題。所以，在某些應用場景下，我們需要以 CPU 為執行載體，嘗試各種方法，以提升模型推理效能。

我們在優化某關鍵點模型推理效能的工程實踐中，基於 MegEngine 推理引擎，發現有兩個優化方法比較有效，NCHW44 和 Record。本文將對它們的原理和使用方法做比較詳細的說明。

NCHW44 優化

原理

眾所周知，增加計算的並行程度是提升計算速度的重要手段。在 CPU 上，這就需要使用 SIMD 指令 ——Single Instruction, Multiple Data，單指令多資料，即執行單條指令，操作完成多個資料的運算。例如執行加法運算，如果使用非 SIMD 指令即一般的加法指令，每次只能操作一個數，而在模型推理中，這個數經常是 8 位、16 位，最大不過是 32 位的浮點數，這對於現代 64 位的暫存器來說，確實有些浪費。如果在暫存器中儲存多個數，一條指令完成運算，則能成倍的提升計算速度。在 x86 CPU 上，SIMD 的實現是 SSE、AVX 等指令集，而在 ARM CPU 上，則是 NEON 指令集。而且 CPU 還提供了 SIMD 指令的專用暫存器，在 x86 平臺上，暫存器位數為 128 位、256 位，甚至是 512 位，在 ARM 平臺上，暫存器位數是 128 位，這樣就可以一次完成 4 個 float32 資料的運算。因此，如果能想辦法在模型推理運算中儘可能多的使用 SIMD，就能提升推理的效能。

我們來看一下在模型推理中使用 SIMD 會遇到什麼問題。通常，張量在記憶體中的儲存方式為 NCHW（即每個通道的行列資料連續排布，再順序儲存各個通道），如在處理常見的卷積操作時，卷積核的尺寸可能多種多樣，比如 3x3，那麼每次需要取一行的 3 個連續的畫素資料與卷積核相應位置資料相乘（再處理其他列和通道），而對應 SIMD 指令，其使用的暫存器通常為 128 位，使用 float32 的話也需要一次處理 4 個資料才能充分發揮其優勢，而這四個資料必須在記憶體中處於相鄰位置，所以這種計算方式大大限制了 SIMD 指令。

作為改進，在 NCHW44（也稱 NC4HW4）佈局下，同一位置（HW）的 4 個通道的資料被連續排列到了一起，在卷積操作時它們一起參與計算，每次 SIMD 指令執行可以將它們一起載入暫存器，這樣就提升了計算效率。下圖示意了 NCHW44 的資料儲存排列方式。

實踐

MegEngine 支援兩種方式使用 NCHW44 優化：

1. 離線 dump（序列化）成 NCHW44 模型，推理時 MegEngine 會自動判斷出它的排列方式，執行相應的運算元實現。下面兩個用於 dump 的方法

megengine.jit.trace.dump

megengine.core.tensor.megbrain_graph.optimize_for_inference

支援關鍵字引數 enable_nchw44，將引數值設為 True，所輸出的就是 NCHW44 的模型。

對應的，如果想通過 load_and_run 預先測試效能，可以在使用 sdk/load-and-run/dump_with_testcase_mge.py 指令碼時新增引數 —enable-nchw44，生成的模型即是可被 load_and_run 載入執行的 nchw44 模型。

2. 線上開啟轉換，dump 模型時不做 nchw44 配置，執行時通過 option 開啟轉換：

serialization::GraphLoader::LoadConfig load_config;

load_config.comp_graph = ComputingGraph::make();

auto &&graph_opt = ret.load_config.comp_graph->options();

graph_opt.graph_opt.enable_nchw44();

對應的，如果想通過 load_and_run 預先測試效能，可以在執行 load_and_run 時，新增命令列引數 —enable-nchw44。

兩種方式可以結合具體的使用情況來選擇：如果我們開發的 sdk 或 app 可能載入多個模型，有些使用 NCHW44 而有些不使用，比較適合選擇離線方式；如果因為某些原因，我們無法重新 dump 模型（比如原始的模型檔案丟失），則只能選擇線上方式。

效果

在我們的工程實踐中，某模型在當前主流 android 手機上的推理速度，大概有 20%-30% 左右的提升。

record 優化

原理

當 MegEngine 執行推理時，底層執行的是靜態圖，它的執行序列是確定的。對於圖中的每個運算元，執行都要分為兩個步驟：準備 kernel 和實際執行。在準備 kernel 階段，MegEngine 會依據 filter size、stride、shape 等資訊決定要執行的演算法，也就是選擇要執行的函式，即 kernel（對於卷積運算，可能會有多種不同的實現）。在執行階段，再實際呼叫這些函式。

如果選擇所需的依據不變（實際情況中主要是 shape 不要變），那麼這個準備 kernel 的過程就只需被執行一次，並把選擇的各個函式物件記錄到一個列表中，以後再執行的時候，直接順序地從列表中取出函式物件，執行即可。這樣就節省了後續各次執行時準備 kernel 的時間。這也就是 record 這個名字的含義所在。

目前 MegEngine 存在兩種級別的 record。record1 主要是為了加速執行，原理如上所述；record2 主要是為了節省記憶體，如果 shape 不變，MegEngine 可以析構圖上儲存的一些資訊（這些資訊可以在 shape 改變時用來做 shape 的推導）。對於我們希望提升計算效能的場景，一般 record1 比較合適。

注意 record 的一個最重要的限制條件是 shape 不能改變。對於某些檢測模型，可能需要依據輸入圖的尺寸，對模型進行 resize，這種情況就無法使用 record。對於輸入長寬和通道數不變的模型，仍需注意，batch 引數（即 NCHW 中的 N）也不能變，這是可能被忽略的。另外，模型載入後，在第一次執行之前，我們還是可以改變 shape 的，只要第一次執行之後不再改變 shape，就不影響 record 的使用。

除了 shape 不變這個條件之外，還有一些限制條件：

  1. 所有的運算元不能依賴動態記憶體分配，因為記錄的函式物件還包含輸入輸出的指標，動態記憶體情況下會發生變化；

  2. Host 端的輸入輸出指標不能變；

  3. 同步只能發生在網路執行的末尾，即不能在網路執行過程中，在某中間節點執行同步；

  4. 整個圖中不能存在多個 compnode。

這些條件對於一般的使用，基本可以滿足。

實踐

在 option 中開啟

serialization::GraphLoader::LoadConfig load_config;

load_config.comp_graph = ComputingGraph::make();

auto &&graph_opt = load_config.comp_graph->options();

graph_opt.comp_node_seq_record_level = 1; // 2

對應的，如果想通過 load_and_run 預先測試效能，可以在執行 load_and_run 時，新增命令列引數 --record-comp-seq 或 --record-comp-seq2。

效果

在我們的工程實踐中，某模型在當前主流 android 手機上的推理速度，大概有 10% 左右的提升。

總結

本文從原理和使用方面介紹了 MegEngine 的 NCHW44 和 record 兩個優化方法，它們只是我們在優化某關鍵點模型推理效能時嘗試發現比較有效的兩個方法。優化方法的有效性取決於模型的特點，因此對於具體的模型，可以嘗試 MegEngine 的其他優化選項，選擇比較合適的方法。當然，優化是多方面的，除了模型推理本身之外，優化預處理和後處理，減少資料複製，對於 Android 裝置合理的設定 CPU 親緣性等等，也是可以嘗試和考慮的方案。

相關文章