【AI系統】Im2Col 算法

作為早期的 AI 框架，Caffe 中卷積的實現采用的是基于 Im2Col 的方法，至今仍是卷積重要的優化方法之一。

從上一篇文章的介紹中可以看到，在 CNN 中卷積直接計算的定義中，卷積核在輸入圖片上滑動，對應位置的元素相乘后相加求和，滑窗的大小由卷積核決定。由于滑動操作時的窗口的數據橫向是連續的，但是不同行在內存中是不連續的，在計算時有可能需要多次訪問內存。由于多次訪問內存直接增加了數據傳輸時間，從而進一步影響了卷積計算速度，這樣就造成了整體效率低等問題。

Im2Col 是計算機視覺領域中將圖片轉換成矩陣的矩陣列的計算過程。Img2col 的作用就是將卷積通過矩陣乘法來計算，從而能在計算過程中將需要計算的特征子矩陣存放在連續的內存中，有利于一次將所需要計算的數據直接按照需要的格式取出進行計算，這樣便減少了內存訪問的次數，從而減小了計算的整體時間。

在 AI 框架發展的早期，Caffe 使用 Im2Col 方法將三維張量轉換為二維矩陣，從而充分利用已經優化好的 GEMM 庫來為各個平臺加速卷積計算。最后，再將矩陣乘得到的二維矩陣結果使用 Col2Im 將轉換為三維矩陣輸出。其處理流程如下圖所示。

數學原理

Im2Col 算法的優化策略為用空間換時間，用連續的行向量的存儲空間作為代價優化潛在的內存訪問消耗的時間。其實施過程依據的是線性代數中最基本的矩陣運算的原理，根據上一篇文章的特征圖計算公式可得：

\[\begin{aligned} f_{out}(i,j)&=\sum_{C=0}^{k_{C}}\sum_{m=0}^{3}\sum_{n=0}^{3}I(i+m,j+n)K(m,n,C)\\ &=\sum_{C=0}^{k_{C}}\sum_{m=0}^{3}[I(i+m,j+0),I(i+m,j+1),I(i+m,j+2)]_{1×3}[K(m,0,C),K(m,1,C),K(m,2,C)]^{T}_{1×3}\\ &= \sum_{C=0}^{k_{C}}[I(i+0,j+0),…,I(i+2,j+2)]_{1×9}[K(0,0,C),…,K(2,2,C)]^{T}_{1×9}\\ &= \sum_{C=0}^{k_{C}}[K(0,0,C),…,K(2,2,C)]_{1×9}[I(i+0,j+0),…,I(i+2,j+2)]_{1×9}^{T}\\ &= [K(0,0,0),…,K(2,2,k_{C})]_{1×9×k_{C}}[I(i+0,j+0),…,I(i+2,j+2)]_{1×9×k_{C}}^{T}\\ \end{aligned} \]

根據上述式子可知，每個窗口下對應的卷積計算和兩個一維向量的點乘計算是等價的。由此可以推知，在未改變直接卷積計算的參數量和連接數的情況下，這種將每個窗口中的特征子矩陣展開成一維的行向量后再進行矩陣乘計算的思路，能夠通過減少計算過程中的訪存需求，優化整體的計算時間。

Im2Col 算法就是在這樣的基礎上進行設計的。其將卷積通過矩陣乘法來計算，最后調用高性能的 Matmul (矩陣乘法)進行計算。該方法適應性強，支持各種卷積參數的優化，在通道數稍大的卷積中性能基本與 Matmul 持平，并且可以與其他優化方法形成互補。

算法過程

除非特別說明，本文默認采用的內存布局形式為 NHWC。其他的內存布局和具體的轉換后的矩陣形狀或許略有差異，但不影響算法本身的描述。

Im2Col+Matmul 方法主要包括兩個步驟：

使用 Im2Col 將輸入矩陣展開一個大矩陣，矩陣每一列表示卷積核需要的一個輸入數據，按行向量方式存儲。

使用上面轉換的矩陣進行 Matmul 運算，得到的數據就是最終卷積計算的結果。

卷積過程

一般圖像的三通道卷積，其輸入為 3 維張量 $(H, W, 3)$，其中 $H$,$W$ 為輸入圖像的高和寬，3 為圖像的通道數；卷積核為 4 維張量 $(N, C, KH, KW)$，其中 $N$ 為卷積核的個數,$KH$,$KW$ 為卷積核的高和寬，$C$ 為卷積核的通道數，卷積核的通道數應與輸入圖像的通道數一致；輸出為 3 維張量 $(N, H, W)$,其中 $H$,$W$ 為輸入圖像的高和寬，$N$ 為輸出圖像的通道數，輸出圖的通道數應與卷積核個數一致。此段中 $H,W$ 只是代指，并不表示數值通用，輸出圖像的寬高具體數值需要按照公式另行計算。其卷積的一般計算方式為：

在神經網絡中，卷積默認采用數據排布方式為 $NHWC$，意為（樣本數，高，寬，通道數）。輸入圖像/特征圖為 4 維張量 $(N,IH,IW,IC)$，其中 $N$ 為輸入圖像的個數,也可以理解為單次訓練的樣本數，$IH$,$IW$ 為輸入圖像的高和寬，$IC$ 為通道數；卷積核為 4 維張量 $(OC,KH,KW,IC)$，卷積核的通道數應與輸入圖像的通道數一致，所以均為 $IC$，這里的 $OC$ 應與下圖中卷積核個數 N 在數值上對應相等，表示卷積核的個數，其遵循一個卷積核計算得到一張特征圖這樣一一對應的規則。注意圖中的 N 與輸入圖像的個數 $N$ 并不相關，數值也不一致，圖中的 N 表示卷積核個數。輸出為 4 維張量 $(N,OH,OW,OC)$，這里的 $N$ 等價于前述輸入圖像的個數，$OC$ 表示輸出特征圖的個數，也就是每個樣本卷積后的輸出的通道數。卷積的一般計算方式為：

Im2Col 算法原理

Im2Col 算法的核心是改變了數據在內存中的排列存儲方式，將卷積操作轉換為矩陣相乘，對 Kernel 和 Input 進行重新排列。將輸入數據按照卷積窗進行展開并存儲在矩陣的列中，多個輸入通道的對應的窗展開之后將拼接成最終輸出 Matrix 的一列。其過程如下圖所示，卷積核被轉化為一個 $N×(KW*KH*C)$ 的二維矩陣，輸入被轉化為一個 $(KW*KH*C)×(OH*OW)$ 的矩陣，輸入矩陣的列數由卷積核在輸入圖像上的滑動次數所決定，具體數值可由特征圖的尺寸 $(OH*OW)$ 計算得到，這兩個數值根據上一篇文章特征圖的公式進行計算。

Input 重排

對 Input 進行重排，得到的矩陣見下圖右側，矩陣的行數對應輸出 $OH*OW$ 個數，也就是卷積核在 Input 上的滑動次數；每個行向量里，先排列計算一個輸出點所需要輸入上第一個通道的 $KH*KW$ 個數據，根據卷積窗的大小逐行拼接成一段行向量，排完當前通道的數據后，以同樣模式再按次序排列之后的通道的數據，直到第 $IC$ 個通道，最終構成前述完整的一個行向量。

權重數據重排

對權重數據進行重排，將 $N$ 個卷積核展開為權重矩陣的一行，因此共有 $N$ 行，每個行向量上先排列第一個輸入通道上 $KH*KW$ 數據，根據卷積窗的大小逐行拼接成一段行向量，排完當前通道的數據后，以同樣模式再依次排列后面的通道數據直到 $IC$。

通過數據重排，完成 Im2Col 的操作之后會得到一個輸入矩陣，卷積的 Weights 也可以轉換為一個矩陣，卷積的計算就可以轉換為兩個矩陣相乘的求解，得到最終的卷積計算結果。

推理引擎中的數據重排

首先歸納一下 Im2Col 算法計算卷積的過程。其具體過程如下（簡單起見忽略 Padding 的情況，即認為 $OH=IH,OW=IW$）：

將輸入由 $N×IH×IW×IC$ 根據卷積計算特性展開成 $(OH*OW)×(N*KH*KW*IC)$ 形狀二維矩陣。顯然，轉換后使用的內存空間相比原始輸入多約 $KH?KW?1$ 倍；
權重形狀一般為 $OC×KH×KW×IC$ 四維張量，可以將其直接作為形狀為 $(OC)×(KH*KW*IC)$ 的二維矩陣處理；
對于準備好的兩個二維矩陣，將 $(KH*KW*IC)$ 作為累加求和的維度，運行矩陣乘可以得到輸出矩陣 $(OH*OW)×(OC)$；
將輸出矩陣 $(OH*OW)×(OC)$ 在內存布局視角即為預期的輸出張量 $N×OH×OW×OC$，或者使用 Col2Im 算法變為下一個算子輸入 $N×OH×OW×OC$；

在權重數據的重排過程中，以上四個階段在推理引擎中的執行方式和執行模塊不一定在同一個階段進行。其中 1，3，4 可能會在 Kernel 層執行，但 2 可能會在預編譯階段或者離線轉換優化模塊去執行。

而在 Input 數據的重排則會在正式計算時感知到數據流后進行。總的來說，不同的推理引擎在應用 Im2Col 算法時，輸入數據的重排通常發生在以下階段：

模型轉換或圖優化階段：在這個階段，推理引擎可能會分析模型的卷積層，并決定是否應用 Im2Col 算法。如果使用，引擎會相應地調整模型的結構，以便在執行卷積時進行數據重排。例如，TensorRT 在構建優化圖時，會考慮是否將卷積層轉換為 Im2Col 形式。
推理初始化階段：在執行推理之前，推理引擎會進行初始化，這可能就包括為 Im2Col 操作分配必要的內存空間，并準備執行數據重排的代碼路徑。
預處理階段：在實際執行推理之前，輸入數據需要經過預處理，以滿足模型的輸入要求。這可能包括縮放、歸一化和其他數據轉換。在某些情況下，Im2Col 的重排也可以看作是預處理的一部分，雖然它與卷積操作更相關。
卷積層的前向傳播階段：在執行卷積層的前向傳播時，如果使用了 Im2Col 算法，推理引擎會在這一階段對輸入數據進行重排。即在卷積操作之前，輸入特征圖會被轉換為一個二維矩陣，其中每一行對應于卷積核在輸入特征圖上的一個位置。這個重排操作是 Im2Col 算法的核心部分。
后處理階段：在卷積操作完成后，如果需要，推理引擎可能會將數據從 Im2Col 格式轉換回原始格式。這通常是在卷積層的輸出被進一步處理或傳遞到下一個網絡層之前完成的。

在 AI 框架中，Im2Col 通常是為了優化卷積操作而設計的，它通過將多次卷積操作轉換為一次大矩陣乘法，從而可以利用現有的高性能線性代數庫來加速計算。隨著 AI 框架的發展，很多框架也實現了更加高效的卷積算法，比如 Winograd 算法或者直接使用 cuDNN 等專門的卷積計算庫，這些庫內部可能對 Im2Col 操作進行了進一步的優化。

Im2Col 算法總結

Im2Col 計算卷積使用 GEMM 庫的代價是額外的內存開銷。使用 Im2Col 將三維張量展開成二維矩陣時，原本可以復用的數據平坦地分布到矩陣中，將輸入數據復制了 $KH?KW?1$ 份。
轉化成矩陣后，可以在連續內存和緩存上操作，而且有很多庫提供了高效的實現方法（BLAS、MKL），Numpy 內部基于 MKL 實現運算的加速。
在實際實現時，離線轉換模塊實現的時候可以預先對權重數據執行 Im2Col 操作，而 Input 數據的 Im2Col 和 GEMM 的數據重排會同時進行，以節省運行時間。

空間組合優化算法

Im2Col 是一種比較樸素的卷積優化算法，在沒有精心處理的情況下會帶來較大的內存開銷。空間組合（Spatial pack）是一種類似矩陣乘中重組內存的優化算法。它是指在應用 Im2Col 算法時，通過對卷積操作中的空間數據進行特定的組合和重排，以減少計算和內存訪問的復雜性，從而提高推理效率。

空間組合優化算法是一種基于分治法（Divide and Conquer）的方法，基于空間特性，將卷積計算劃分為若干份，分別處理。

具體優化方法

常見的空間組合優化方法有：

分塊卷積（Blocked Convolution）：將大卷積核分解為多個小卷積核，每個小卷積核單獨計算，這樣可以減少內存訪問和提高緩存利用率。這種方法在處理大型圖像或特征圖時特別有效。
重疊數據塊（Overlapping Blocks）：在進行 Im2Col 操作時，可以通過重疊數據塊來減少邊緣效應和填充（padding）的需求。這種方法可以減少計算量，但可能會增加內存訪問的復雜性。
稀疏卷積（Sparse Convolution）：對于稀疏數據，可以通過只對非零值進行卷積來減少計算量。這種方法在處理稀疏特征圖時特別有效。
權重共享（Weight Sharing）：在卷積操作中，同一個卷積核會在不同的位置重復使用。通過在 Im2Col 操作中利用這一點，可以減少權重數據的重復加載和存儲。
張量化（Tensorization）：利用特定硬件（如 GPU）的特定指令來優化張量操作，例如使用 SIMD（單指令多數據）指令集。這種方法可以加速 Im2Col 操作中的矩陣乘法。
循環展開和軟件流水線：這些是編譯器優化技術，可以用于優化循環結構，減少循環開銷和提高指令級并行性。在 Im2Col 操作中，循環展開可以減少循環迭代的開銷，而軟件流水線可以優化數據流和處理流程。
內存復用：通過復用中間計算結果所占用的內存，可以減少總的內存使用。在 Im2Col 操作中，可以復用輸入和輸出數據的內存空間，以減少內存分配和釋放的次數。

這些空間組合優化方法可以單獨使用，也可以組合使用，以提高 Im2Col 操作的效率。在實際應用中，選擇哪種優化策略取決于具體的模型結構、硬件平臺和性能目標。隨著 AI 框架和硬件的發展，許多框架已經采用了更加高效的卷積實現，如直接卷積（Direct Convolution）、Winograd 算法或利用專用硬件加速器，這些實現可能不再需要顯式的 Im2Col 和 Col2Im 操作。

空間組合例子

以分塊卷積為例，如圖所示在空間上將輸出、輸入劃分為四份：

劃分后，大卷積計算被拆分為若干個小卷積進行計算，小卷積塊的大小必須與卷積核的大小相匹配。劃分過程中計算總量不變，但計算小矩陣時訪存局部性更好，可以借由計算機存儲層次結構獲得性能提升。

在傳統的卷積操作中，輸入特征和卷積核的每個通道都需要參與計算，這導致內存訪問模式復雜，難以優化。而分組卷積將通道分成多個組，每個組內的通道數減少，使得內存訪問更加規則和局部化，有利于提高緩存利用率。此外，這種基于分治法的分解策略，有助于提高并行處理的效率。在分組卷積中，每個組內的卷積核是獨立的，這意味著不同組之間的權重和激活可以共享內存空間。這種共享可以減少內存占用，尤其是在使用大量卷積核的網絡結構中。

算法注意與問題點

空間組合優化注意點

值得注意的是，上文的描述中忽略了 Padding 的問題。實際將輸入張量劃分為若干個小張量時，除了將劃分的小塊中原始數據拷貝外，還需要將相鄰的小張量的邊界數據拷貝：

\[N\times\left(\frac{H}{h}+2 (KH-1)\right)\times\left(\frac{W}{w}+2 (KW-1)\right)\times C \]

這里的 $2(KH?1)$ 和 $ 2(KW?1)$ 遵循 Padding 規則。規則為 VALID 時，可以忽略；規則為 SAME 時，位于 Input Tensor 邊界一邊 Padding 補 0，不在 Input Tensor 邊界 Padding 使用鄰居張量值。也就是說，在真正使用這種方法的時候，可以通過重疊數據塊來減少邊緣效應和填充（padding）的需求。這種方法可以減少計算量，但可能會增加內存訪問的復雜性。

空間組合優化問題點

實際應用中可以拆為很多份。例如可以拆成小張量邊長為 4 或者 8 ，從而方便編譯器向量化計算操作。隨著拆分出的張量越小，Padding 引起的額外內存消耗越大，其局部性也越高，負面作用是消耗的額外內存也越多。
對于不同規模的卷積，尋找合適的劃分方法不是一件容易的事情。正如計算機領域的許多問題一樣，該問題也是可以自動化的，例如通過使用 AI 編譯器可以在這種情況下尋找較優的劃分方法。

如果您想了解更多AI知識，與AI專業人士交流，請立即訪問昇騰社區官方網站https://www.hiascend.com/或者深入研讀《AI系統：原理與架構》一書，這里匯聚了海量的AI學習資源和實踐課程，為您的AI技術成長提供強勁動力。不僅如此，您還有機會投身于全國昇騰AI創新大賽和昇騰AI開發者創享日等盛事，發現AI世界的無限奧秘~

posted @ 2024-11-22 11:37 ZOMI醬醬閱讀(39) 評論(0) 收藏舉報

刷新頁面返回頂部