卷積神經網絡性能優化
引言
卷積（Convolution）是神經網絡的核心計算之一，它在計算機視覺方面的突破性進展引領了深度學習的熱潮。卷積的變種豐富，計算復雜，神經網絡運行時大部分時間都耗費在計算卷積，網絡模型的發展在不斷增加網絡的深度，因此優化卷積計算就顯得尤為重要。
隨著技術的發展，研究人員提出了多種優化算法，包括 Im2col、Winograd 等等。本文首先定義卷積神經網絡的概念，繼而簡要介紹幾種常見的優化方法，并討論作者在該領域的一些經驗。
卷積神經網絡的概念
卷積神經網絡（Convolution Neural Networks, CNN）的概念拓展自信號處理領域的卷積。信號處理的卷積定義為
(1)
由于對稱性圖片 卷積計算在直覺上不易理解，其可視化后如圖一所示。圖中紅色滑塊在移動過程中與藍色方塊的積繪制成的三角圖案即為卷積結果 (?????)(??) 在各點上的取值。
圖一：信號處理中的卷積(https://jackwish.net/convolution-neural-networks-optimization.html)
公式1的離散形式為 ：
(2)
將該卷積拓展到二維空間即可得到神經網絡中的卷積，可簡寫為：
(3)
其中 ?? 為卷積輸出，?? 為卷積輸入，?? 為卷積核。該計算的動態可視化可以參考 conv_arithmetic(https://github.com/vdumoulin/conv_arithmetic)，這里不再介紹。
當應用到計算機視覺中處理圖片時，圖片的通道（Channel）可以對二維卷積簡單堆疊，即：
(4)
其中 ??c 是輸入的通道。這便是在三維張量中應用二維卷積的計算。
很多時候，公式描述顯得不是很直觀，圖二是堆疊的二維卷積的可視化。其中，與輸入、輸出、卷積核相關的標記帶有前綴 I、O、K。此外，本文圖例對輸出、輸入、卷積核三者的著色一般為：橙色、黃色、綠色。
圖二：卷積計算定義
當中張量的內存布局為 NHWC 時，卷積計算相應的偽代碼如下。其中外三層循環遍歷輸出 ??C 的每個數據點，對于每個輸出數據都需要經由內三層循環累加求和得到（點積）。
for (int oh = 0; oh < OH; oh++) {
for (int ow = 0; ow < OW; ow++) {
for (int oc = 0; oc < OC; oc++) {
C[oh][ow][oc] = 0;
for (int kh = 0; kh < KH, kh++){
for (int kw = 0; kw < KW, kw++){
for (int ic = 0; ic < IC, ic++){
C[oh][ow][oc] += A[oh+kh][ow+kw][ic] * B[kh][kw][ic];
}
}
}
}
}
}
和矩陣乘的優化方法類似，我們也可針對該計算進行向量化、并行化、循環展開的基本的優化操作。卷積的問題在于其 ???? 和 ???? 一般不超過 5 ，這不容易向量化，而計算特征又有輸入在空間維度存在數據復用。該計算的復雜性導致產生了幾種優化方法，下面我們介紹幾種。
Im2col 優化算法
作為早期的深度學習框架，Caffe 中卷積的實現采用的是基于 im2col 的方法，至今仍是卷積重要的優化方法之一。
Im2col 是計算機視覺領域中將圖片轉換成矩陣的矩陣列（column）的計算過程。從上一節的介紹中可以看到，二維卷積的計算比較復雜不易優化，因此在深度學習框架發展的早期，Caffe 使用 Im2col 方法將三維張量轉換為二維矩陣，從而充分利用已經優化好的 GEMM 庫來為各個平臺加速卷積計算。最后，再將矩陣乘得到的二維矩陣結果使用 Col2im 將轉換為三維矩陣輸出。
算法過程
除非特別說明，本文默認采用的內存布局形式為 NHWC 。其他的內存布局和具體的轉換后的矩陣形狀或許略有差異，但不影響算法本身的描述。
圖三：Im2col 算法計算卷積的過程
圖三是使用 Im2col 算法計算卷積的過程示例，具體的過程包括（簡單起見忽略 Padding 的情況，即認為 ????=????,????=????：
- 將輸入由 ????×????×???? 根據卷積的計算特性展開成 (????×????)×(????×????×????) 形狀的二維矩陣。顯然，轉換后使用的內存空間相比原始輸入多約 ??????????1 倍。
- 權重的形狀一般為 ????×????×????×???? 四維張量，可以將其直接作為形狀為 (????)×(????×????×????)) 的二維矩陣處理。
- 對于準備好的兩個二維矩陣，將 (????×????×????) 作為累加求和的維度，運行矩陣乘可以得到輸出矩陣 (????×????)×(????)。這一過程可以直接使用各種優化過的 BLAS（Basic Linear Algebra Subprograms）庫來完成。
- 輸出矩陣 (????×????)×(????) 在內存布局視角即為預期的輸出張量 ????×????×???? 。
Im2col 計算卷積使用 GEMM 的代價是額外的內存開銷。這是因為原始卷積計算中，卷積核在輸入上滑動以計算輸出時，相鄰的輸出計算在空間上復用了一定的輸入輸出。而用 Im2col 將三維張量展開成二維矩陣時，這些原本可以復用的數據平坦地分布到矩陣中，將輸入數據復制了 ??????????1 份。
當卷積核尺寸 ????×???? 是 1×1 時，上述步驟中的 ???? 和 ???? 可以消去，即輸入轉換后形狀為 (????×????)×(????)，卷積核形狀為 (????)×(????)，卷積計算退化為矩陣乘。注意觀察，對于這種情況，Im2col 過程實際上并沒有改變輸入的形狀，因此矩陣乘操作可以直接在原始輸入上運行，從而省去內存組織的過程（即上述步驟中的 1、2、4 步）。
內存布局與卷積性能
神經網絡中卷積的內存布局主要有 NCHW 和 NHWC 兩種，不同的內存布局會影響計算運行時訪問存儲器的模式，特別是在運行矩陣乘時。本小節分析采用 Im2col 優化算法時計算性能性能和內存布局的關系。
在完成 Im2col 轉換后，得到用于運行矩陣乘的輸入矩陣和卷積核矩陣。對計算過程施加矩陣計算中常用的數據劃分、向量化等優化方法（相關定義請參考通用矩陣乘（GEMM）優化算法：https://jackwish.net/gemm-optimization.html）。下面著重分析在這種場景下，不同內存布局對性能的影響。
首先考慮 NCHW 內存布局，將 NCHW 內存布局的卷積對應到矩陣乘 ??=???? 時，?? 是卷積核（filter），?? 是輸入（input），各個矩陣的維度如圖四所示。圖中的 ??,??,?? 用于標記矩陣乘，即圖片，同時標記出它們和卷積計算中各個維度的關系。
圖四：NCHW 內存布局卷積轉換成的矩陣乘
對該矩陣施行劃分后，我們詳細分析局部性的表現，并標記在圖四中。其中 Inside 表示 4×4 小塊矩陣乘內部的局部性，Outside 表示在削減維度方向小塊之間的局部性。
- 對輸出而言，小塊內訪存局部性較差，這是因為每次向量化加載會加載四個元素，每次加載都會發生緩存缺失（Cache miss）。外部表現取決于全局計算方向——行優先則局部性較好，列優先則較差。輸出的行為不是這里的討論終點，畢竟它沒有訪存復用。
- 對卷積核而言，小塊內訪存局部性較差，這和輸出類似。當小塊加載發生緩存缺失時，處理器會一次性加載一個緩存行（Cache line），這使得后續的若干個小塊訪存都能緩存命中（Cache hit），直到該緩存行全部命中后進入下一個緩存行的范圍。因此小塊外部局部性較好。
- 對輸入而言，小塊內訪存局部性表現同樣不好。然而不同于卷積核，小塊中因緩存缺失加載到緩存中的緩存行數據只會被使用一次，因為這種內存布局中下一個小塊的地址范圍一般超出了一個緩存行。因此輸入的幾乎每次內存加載都會發生高速緩存缺失——Cache 沒有起到加速的作用，每次訪存都需要到內存取數據。
因此，用 Im2col 處理卷積計算時，NCHW 布局對內存很不友好。
圖五是與之相對的 NHWC 內存布局的示例。值得注意的是，NHWC 和 NCHW 中 ??、?? 矩陣所代表的張量發生了調換——????????????=??????????×????????????（調換一下只是不想多畫一張圖）。具體的拆分方式仍然一樣，也正是上一小節中描述的步驟所構建的矩陣。
圖五：NHWC 內存布局卷積轉換成的矩陣乘
類似地，分析三個張量的訪存表現可知：
- 對輸出而言，NHWC 和 NCHW 表現一樣。
- 對輸入而言，小方塊的內部局部性表現不是很好，因為幾次向量加載都會發生緩存不命中；而外部局部性表現則較好，因為在削減維度滑動使用的內存是連續的。這種表現和 NCHW 中卷積核的表現一樣，整體來看都是對高速緩存比較友好的內存布局。
- 對卷積核而言，NHWC 的情況和 NCHW 中輸入的情況類似，小塊內和小塊外的局部性都較差。
兩種內存布局中的卷積核緩存表現并不是問題，因為卷積核在運行期間保持不變，可以在模型加載階段轉換卷積核的內存布局，使其在小塊外的內存對緩存友好（例如將 (????×????×????)×(????) 的布局轉換為 (????)×(????×????×???? ）。這里值得說明的是一般框架或引擎的運行都至少可分為兩個階段：準備階段和運行階段。一些模型的預處理工作可以放在準備階段完成，例如重新排布卷積核的內存布局這種在運行階段保持不變的數據。
因此，當使用 Im2col 方法計算時，整體的訪存表現取決于輸入的情況，即 NHWC 的內存布局要比 NCHW 內存布局更加友好。我們在實踐過程中的一個實驗表明，對于一個 1×1 卷積核的卷積，當采用類似的優化方法時，從 NCHW 轉換為 NHWC 可以將高速緩存缺失率從約 50% 降低到 2% 左右。這種程度的提高可以大幅改進軟件的運行性能（這里特指不使用特別設計過的矩陣乘優化方法）。
空間組合優化算法
Im2col 是一種比較樸素的卷積優化算法，在沒有精心處理的情況下會帶來較大的內存開銷。空間組合（Spatial pack）是一種類似矩陣乘中重組內存的優化算法。
圖六：空間組合優化算法對計算的劃分
空間組合優化算法是一種基于分治法（Divide and Conquer）的方法——它基于空間特性將卷積計算劃分為若干份，分別處理。圖六所示是在空間上將輸出、輸入劃分為四份。
劃分后，一個大的卷積計算被拆分為若干個小的卷積計算。雖然在劃分的過程中計算總量不變，但計算小矩陣時訪存局部性更好，可以借由計算機存儲層次結構獲得性能提升。這通過圖七中的步驟來完成。該步驟和上節中 Im2col 重組內存的過程類似：
- 在 H 和 W 維度劃分，將形狀為 ??×??×??×?? 的輸入張量拆分為 ???? 個（兩個方向分別拆分 ? 和 ?? 次）形狀為 ??×??/?×??/??×???? 的張量，分別將這些小的張量組織為連續內存；
- 將得到的 ???? 個輸入張量分別和卷積核做二維卷積操作，即可得到 ???? 個形狀為 ??×??/?×??/??×???? 的輸出張量；
- 將這些輸出張量重組內存布局得到最終形狀為 ??×??×??×???? 的輸出。
圖七：空間組合計算的步驟
值得注意的是，方便起見，上文的描述中忽略了 Padding 的問題。實際在步驟 1 中將輸入張量劃分為若干個小張量時，除了將劃分的小塊中原始數據拷貝外，還需要將相鄰的小張量的邊界數據拷貝。具體而言，如圖八所示，空間拆分得到的小張量的形狀實際上是：
??×(??/?+2(?????1))×(??/??+(?????1))×??.(5)
這里的 2(?????1) 和 2(?????1) 遵循 Padding 規則。規則為圖片時，它們可以忽略；規則為 SAME 時，位于源張量邊界的一邊 Padding 補圖片，不在源張量邊界的 Padding 則使用鄰居張量的值。只要考慮一下卷積的計算原理，這是顯而易見的。
圖八：空間組合算法的劃分細節
上面的三個示例圖都是拆分為 4 份的情況，實際應用中可以拆為很多份。例如可以拆成小張量邊長為 4 或者 8 ，從而方便編譯器向量化計算操作。隨著拆分出的張量越小，其局部性也越高，負面作用是消耗的額外內存也越多。這些額外內存是由于 Padding 引入的。當拆分為 ????h?w份時，拆分后 Padding 消耗的內存為：
可以看到，隨著拆分的粒度越小，額外消耗的內存越大。值得注意的是，當拆分到最細粒度時，即將在形狀為 ??×??×??×???? 的輸出的空間上拆分????? 份時，空間組合退化為 Im2col 方法。此時在一個元素上的卷積即為矩陣乘計算一個輸出元素時的點積。
只做空間劃分時，劃分與卷積核無關。而如果在輸出的通道維度劃分，卷積核也可做相應的拆分。通道維度的劃分相當于固定空間劃分后簡單的堆疊，不會對影響內存消耗，但會影響局部性。對于不同規模的卷積，尋找合適的劃分方法不是一件容易的事情。正如計算機領域的許多問題一樣，該問題也是可以自動化的，例如 AutoTVM （https://arxiv.org/abs/1805.08166）可以在這種情況下尋找較優的劃分方法。
Winograd 優化算法
前兩節介紹的兩種算法，Im2col 在將三維張量組織成矩陣后調用 GEMM 計算庫，這些計算庫很大程度上使用一些基于訪存局部性的優化；空間組合優化則本身就是利用局部性優化的方法。本小節介紹的 Winograd 優化算法則是矩陣乘優化方法中 Coppersmith–Winograd 算法的一種應用，是基于算法分析的方法。
這部分公式過多，排版不便，有興趣的話可以參考原文（
??https://jackwish.net/convolution-neural-networks-optimization.html ??
）或其他相關文獻。
間接卷積優化算法
Marat Dukhan 在 QNNPACK（Quantized Neural Network PACKage）中推出了間接卷積算法（The Indirect Convolution Algorithm），似乎到目前為止（2019 年中）依然是所有已公開方法中最快的。最近作者發表了相關的文章來介紹其中的主要思想。
雖然該算法在 QNNPACK 中的實現主要用于量化神經網絡（業界的其他量化優化方案都比較傳統 TensorFlow Lite 使用 Im2col 優化算法、騰訊出品的 NCNN使用 Winograd 優化算法；OpenAI 出品的 Tengine 使用 Im2col 優化算法），但其是一種同樣的優化算法設計。
本文寫作時設計文章尚未公開，而理解該算法設計很多細節內容，最好能結合代碼理解。我的 QNNPACK fork （https://github.com/jackwish/qnnpack）包含一個 explained 分支，其中對增加了對部分代碼的注釋，可作參考。
計算工作流
間接卷積算法的有效工作以來一個關鍵的前提——網絡連續運行時，輸入張量的內存地址保持不變。這一特性其實比較容易滿足，即使地址真的需要變化，也可以將其拷貝到固定的內存區域中。
圖九：間接卷積算法工作流
圖九是間接卷積算法工作流的詳細過程。最左側部分表示多個輸入使用相同的輸入緩沖區（Input Buffer）。間接卷積算法會在該輸入緩沖區基礎上構建如圖十的「間接緩沖區」（Indirect Buffer），而間接緩沖區是間接卷積算法的核心。如圖中右側，在網絡運行時，每次計算出 ??×?? 規模的輸出，其中 ?? 為將 ????×???? 視作一維后的向量化規模。一般 ??×?? 為 4×4、8×8 或 4×8 。在計算 ??×?? 規模大小的輸出時，經由間接緩沖區取出對應使用的輸入，并取出權重，計算出結果。計算過程等價于計算 ??×?? 和 ??×?? 矩陣乘。
在實現中，軟件的執行過程分為兩部分：
- 準備階段：加載模型，配置輸入緩沖區；重排權重，使其內存布局適用于后續計算；
- 運行階段：對于每個輸入，運行 ??????????/?????????/???次核心循環，每次使用 GEMM 方法計算出 ??×?? 大小的輸出。
圖十：間接緩沖區
如相關章節的討論，Im2col 優化算法存在兩個問題，第一是占用大量的額外內存，第二是需要對輸入進行額外的數據拷貝。這兩點如何才能解決呢？間接卷積算法給出的答案是間接緩沖區（Indirect Buffer），如圖十右半所示。
圖十是常規的 Im2col 優化算法和間接卷積優化算法的對比。正如相關小節介紹的那樣，Im2col 優化算法首先將輸入拷貝到一個矩陣中，如圖十中 Input 的相關箭頭，然后執行矩陣乘操作。間接卷積優化算法使用的間接緩沖區中存儲的其實是指向輸入的指針（這也是間接卷積優化算法要求輸入內存地址固定的原因），在運行時根據這些指針即可模擬出類似于 Im2col 的矩陣計算過程。
間接緩沖區布局
間接緩沖區可以理解為是一組卷積核大小的緩沖區，共有 ????×???? 個，每個緩沖區大小為 ????×????——每個緩沖區對應著某個輸出要使用的輸入的地址。每計算一個空間位置的輸出，使用一個間接緩沖區；空間位置相同而通道不同的輸出使用相同的間接緩沖區，緩沖區中的每個指針用于索引輸入中 ???? 個元素。在計算時，隨著輸出的移動，選用不同的間接緩沖區，即可得到相應的輸入地址。無需再根據輸出目標的坐標計算要使用的輸入的地址，這等同于預先計算地址。
圖十一繪制了當 ??×?? 為 4 、???? 和 ???? 均為 3 時，間接緩沖區的實際使用方法與計算過程。圖中命名為局部間接緩沖區意指目前考慮的是計算 ??×?? 時核心計算的過程。
當計算 ??×?? 大小的輸出時，使用的輸入為卷積核在對應輸入位置上滑動 ?? 步所覆蓋的趨于，即規模 (????)×(??+2(?????1))×???? 的輸入。在間接卷積算法中，這些輸入內存由 ??M個間接緩沖區中的指針索引，共有 ??×????×???? 個。圖十一中標識出了輸入空間位置左上角輸入和相應的輸入緩沖區的對應關系。可以看到，這里的 A、B、C、D 四個輸入緩沖區，相鄰的兩個緩沖區所指向的地址區域有 (?????????????????)/???? ，這里即為 2/3 ，各個緩沖區中指針的坐標也已標明。
圖十一：間接緩沖區詳解
圖中將平面緩沖區繪制成三維形式（增加 IC 維度），意在說明間接緩沖區的每個指針可索引 IC 個輸入元素，而每個間接緩沖區索引的內容即為與權重對應的輸入內存區域。
進一步地，左上角的輸入緩沖區排列方式并不是最終的排布方法，實際上這些指針會被處理成圖十一中部間接緩沖區的形式。將左上每個緩沖區中的指針打散，即可得到 ????×???? 指針，將 A、B、C、D 四個緩沖區的不同空間位置的指針收集到一起，即可得到圖中上部分的緩沖區排列方式 ????×????×??。可以看到， A、B、C、D 四個緩沖區內部相同空間位置的指針被組織到了一起。圖中中上部分是可視化的效果，中下部分則是間接緩沖區的真正組織方式。圖中褐色和深黃色的著色對應著相同的輸入內存或指針。值得注意的是，圖例中 Stride 為 1，當 Stride 不為 1 時，重新組織后 A、B、C、D 相同空間的坐標（對應于在輸入的坐標）不一定是連續的，相鄰的空間位置橫向坐標相差 ???????????? 大小。
使用間接緩沖區計算
我們已經知道了間接緩沖區的組織形式，以及其指針對應于輸入內存的地址趨于，現在來研究在計算過程中如何使用這些緩沖區。
和上一小節一樣，本小節的討論大體都在計算 ??×?? 規模的輸出，而這些輸出要使用 ?? 個 ????×???? 大小的輸入，其中有數據重用。現在回顧一下 Im2col 的算法（如圖十一中下左部分），當向量化運行計算時，對于 ??×?? 的矩陣乘計算，每次取出 ??×?? 規模的輸入和 ??×?? 規模的權重，對該塊運行矩陣乘即可得到相應的部分和結果。其中 ?? 是向量化計算在 ??K維度上的步進因子。
而卷積之所以可以使用 Im2col 優化算法，本質原因在于其拆解后忽略內存復用后的計算過程等價于矩陣乘。而間接緩沖區使得可以通過指針模擬出對輸入的訪存。在實際運行計算 ??×?? 輸出的計算核（micro kernel）時，會有 ?? 個指針掃描輸入。??個指針每次從圖十一中下部分的間接緩沖區結構中取出 ?? 個地址，即對應于輸入 ??×???? 的內存，如圖中右上角的布局。在這里的步進中，仍是以 ??×?? 形式運行，其中 ?? 在 ???? 維度上運動。當這部分輸入掃描完畢后，這 ?? 個指針從間接緩沖區中取出相應的指針，繼續下一次對 ??×???? 輸入內存的遍歷，每次計算出 1/(?????????) 的輸出部分和。當這個過程運行 ????×???? 次后，即得到了 ??×?? 的輸出。圖十一右下角的偽代碼對應了這一過程。由于間接緩沖區已經被組織成了特定的形狀，因此每次更新 ?? 個指針時，只需要從間接緩沖區指針（圖中偽代碼里的 p_indirect_buffer）中獲取。
這個過程的邏輯不易理解，這里再做一點補充說明。當上述的 ?? 個指針不斷運動掃描內存時，實際上是在掃描三維輸入 Im2col 之后的矩陣。而輸入緩沖區的特定是它將對二維矩陣的掃描轉化為了對三維張量的掃描。對輸入的掃描過程即是對圖十一中上部分可視化的輸入的掃描，聯系左上和左下對應的輸入關系，不難發現它每次掃描輸入中 ??×???? 塊內存。值得注意的是，這里的 ??×???? 由 ?? 個 1×???? 張量組成，它們之間 ?? 維度的間距為 ????????????。
這樣一來，只要運行該計算核心 ??????????/?????????/??? 次，即可得到全部輸出。
間接卷積優化算法解決了卷積計算的三個問題，第一是空間向量化問題，第二是地址計算復雜問題，第三是內存拷貝問題。一般計算卷積時都需要對輸入補零（對于 ????×???? 不是 1×1 的情況），這個過程傳統的方法都會發生內存拷貝。而間接卷積優化算法中的間接緩沖區可以通過間接指針巧妙地解決這一問題。在構造間接緩沖區時額外創建一塊 1×???? 的內存緩沖區，其中填入零值，對于空間中需要補零的位置，將相應的間接指針指向該緩沖區，那么后續計算時即相當于已經補零。例如圖十一中 A 的左上角對應于輸入空間位置 (0,0) 的，當需要補零時該位置一定要為零值，此時只要修改間接緩沖區的地址即可。
討論、總結與展望
至此，本文探討了一些已經發表或開源的卷積神經網絡的優化方法。這些優化方法或多或少地推動了深度學習技術在云端或移動端的應用，幫助了我們的日常生活。例如，最近上海人民乃至全中國人們頭疼的垃圾分類問題，也可以利用深度學習方法來幫助人們了解如何分類(https://news.mydrivers.com/1/633/633858.htm)。
從本文的集中優化方法中可以看到，卷積神經網絡的優化算法依然可以囊括在基于算法分析的方法和基于軟件優化的方法。實際上，在現代計算機系統中，基于軟件優化的方法，特別是基于計算機存儲層次結構的優化顯得更為重要，因為其方法更容易挖掘和應用。另一方面也是因為現代計算機新增的專用指令已經可以抹除不同計算的耗時差異。在這種場景下，基于算法分析的方法要和基于軟件優化的方法結合或許能取得更優的優化效果。
最后，本文討論的優化方法都是通用的方法，而隨著神經網絡處理器（如寒武紀 MLU、Google TPU）的發展，以及其他通用計算處理器的拓展（如點積相關的指令：Nvidia GPU DP4A、Intel AVX-512 VNNI、ARM SDOT/UDOT ），深度學習的優化或許還值得繼續投入資源。
