AI移動端優化之Im2Col+Pack+Sgemm
轉自:https://blog.csdn.net/just_sort/article/details/108412760
這篇文章是基于NCNN的Sgemm卷積為大家介紹Im2Col+Pack+Sgemm的原理以及算法實現,希望對算法優化感興趣或者做深度學習模型部署的讀者帶來幫助。
1. 前言
最近空閑時候在給MsnhNet貢獻Arm端的代碼,地址詳見:https://github.com/msnh2012/Msnhnet 。對于卷積運算一種最常見的實現方法就是Im2Col加上Sgemm,這樣可以把復雜的卷積運算轉換成兩個矩陣直接相乘,在OpenBlas中就提供了矩陣相乘的多種算法如Sgemm。所以這篇文章的目的就是講清楚如何對卷積層進行Im2Col將其轉換為矩陣乘法,以及為了進一步加速矩陣乘法的計算過程使用了NCNN中的Pack策略,最后再使用Neon對主要的計算過程進行優化。
2. Im2Col+Sgemm計算卷積原理
相信大家對于卷積的概念都非常熟悉了,這里就不再重復提起了,我大概說一下我了解的卷積的計算方式有哪些吧。首先是暴力計算,這是最直觀也是最簡單的,但是這種方式訪存很差,效率很低。其次是我在基于NCNN的3x3可分離卷積再思考盒子濾波介紹過的手工展開某些特定的卷積核并且一次處理多行數據,這樣做有個好處就是我們規避掉了一些在行方向進行頻繁切換導致的Cache Miss增加,并且在列方向可以利用Neon進行加速。再然后就是比較通用的做法Im2Col+Sgemm,這種方式可以使得訪存更好。其它常見的方法還有Winograd,FFT,Strassen等等。
本文的重點是嘗試講清楚Im2Col+Sgemm的算法原理及其實現思路,以及在Im2Col的基礎上進行數據Pack進一步改善矩陣乘法的計算中由于Cache Miss帶來的性能下降。
首先,讓我們來考慮一個單通道的長寬均為 444 的輸入特征圖,并且假設卷積核的大小是 3×33\times 33×3,那么它經過Im2Col之后會變成什么樣呢?如下圖所示:
其實變換的方式非常簡單,例如變化的第一步如下圖所示:
直接在卷積核滑動的范圍內,把所有元素順序排列成一列即可,其它位置處同理即可獲得Im2Col的結果。多通道和單通道的原理一樣,如下圖所示:
其次,讓我們來考慮一下卷積核經過Im2Col操作后應該是什么樣子的,對于多通道如下圖所示(單通道就是其中一種顏色):
最后,輸入特征圖和卷積核都進行了Im2Col之后其實都變成了一個二維的矩陣,我們就可以使用矩陣乘法來直接計算結果了。下圖展示了一個 4×4×34\times 4\times 34×4×3 的特征圖和一個 3×3×33\times 3\times 33×3×3 的卷積核經過Im2Col之后之后使用Sgemm進行計算的過程,其中每一種顏色都代表一個輸出通道。
相信看了上面幾個圖之后,大家對Im2Col和Sgemm結合用來做卷積的過程有所了解,為了更加嚴謹,下面還是以公式化的形式再來總結一下:
假設輸入的特征圖維度是 (1,D,H,W)(1,D,H,W)(1,D,H,W) ,表示Batch為1,通道數為 DDD,高為 HHH ,寬為 WWW,卷積核的維度是 (Dout,D,K,K)(D_{out},D,K,K)(Dout?,D,K,K),表示輸出通道數為 DoutD_{out}Dout? ,卷積核大小是 K×KK\times KK×K,則輸入特征圖的Im2Col過程如下所示,窗口從左到右上到下的順序在每個輸入通道同步滑動,每個窗口內容按行展開成一列,然后再按通道順序接上填到 im2colim2colim2col buffer對應的列,并且 im2colim2colim2col buffer按照從左到右的順序填寫。
由于這個地方沒有padding并且步長是1,那么卷積的輸出特征圖的高為 Hout=(H+2?p?K)/Stride+1=H?K+1H_{out}=(H+2*p-K)/Stride+1=H-K+1Hout?=(H+2?p?K)/Stride+1=H?K+1,同理寬為 W?K+1W-K+1W?K+1。所以最后 im2colim2colim2col buffer的寬維度為 (H?K+1)?(W?K+1)(H - K +1) * (W - K + 1)(H?K+1)?(W?K+1),高維度則是 K?K?DK * K * DK?K?D ,最后再把權值Im2Col成 (Dout,D?K?K)(D_{out},D*K*K)(Dout?,D?K?K),這樣卷積就可以直接變成兩個二維矩陣的乘法了。
最后再獲得了乘積矩陣之后只需要按照通道進行順序排列就可以獲得輸出特征圖了,這個過程就是Im2Col的逆過程,學名叫作Col2Im,也比較簡單就不再贅述了。
3. 如何實現Im2Col
這里我們先做一些約定,變量 src
表示輸入特征圖的數據,輸入特征圖的維度是 [1,inChannel,inHeight,inWidth][1,inChannel,inHeight,inWidth][1,inChannel,inHeight,inWidth],輸入卷積核的維度是$ [1,outChannel,KernelH,kernelW]$,最后卷積操作執行完之后的輸出特征圖維度就是 [1,outChannel,outHeight,outWidth][1,outChannel,outHeight,outWidth][1,outChannel,outHeight,outWidth]。
3.1 對輸入特征圖進行Im2Col
首先我們對輸入特征圖進行Im2Col操作,按照之前的介紹可知,我們要得到的結果矩陣的維度是 [inChannel?kernelH?kernelW,outHeight?outWidth][inChannel*kernelH*kernelW,outHeight*outWidth][inChannel?kernelH?kernelW,outHeight?outWidth],這個過程只需要窗口從左到右上到下的順序在每個輸入通道同步滑動,每個窗口內容按行展開成一列,然后再按通道順序接上填到 im2colim2colim2col buffer對應的列,并且 im2colim2colim2col buffer按照從左到右的順序填寫。不難寫出這部分的代碼實現,這里的 src_im2col
等價于 im2colim2colim2col buffer:
// 1. im2colfloat *src_im2col = new float[outWidth * outHeight * kernelH * kernelW * inChannel];const int Stride = kernelW * kernelH * outHeight * outWidth;//const int inSize = inHeight * inWidth;const int outSize = outHeight * outWidth; const int kernelSize = kernelH * kernelW;// inCahnnel x kW x kH// outWidth x outHeightfor(int cc = 0; cc < inChannel; cc++){const float *src0 = src + cc * kernelH * kernelW * inChannel;int dst_idx = Stride * cc;for(int i = 0; i < kernelH; i++){for(int j = 0; j < kernelW; j++){for(int x = 0; x < outHeight; x++){for(int y = 0; y < outWidth; y++){int row = x * StrideH + i;int col = y * StrideW + j;int ori_idx = row * inWidth + col;src_im2col[dst_idx] = src0[ori_idx];dst_idx++; }}}}}
3.2 對卷積核進行Im2Col
按照我們第二節的介紹,我們同樣可以輕易寫出對卷積核進行Im2Col的操作,代碼如下:
int Stride = 0;for(int cc = 0; cc < outChannel; cc++){int c = cc;const float* k0 = kernel + c * inChannel * kernelSize;Stride = kernelSize * inChannel;float* destptr = dest + c * Stride;for(int i = 0; i < inChannel * kernelSize; i++){destptr[0] = k0[0];destptr += 1;k0 += 1;}}
3.3 優化點之Pack策略
按照第二節介紹的思路,在獲取了輸入特征圖和卷積核的Im2Col變換矩陣之后其實就可以利用Sgemm計算出卷積的結果了。
但是如果直接使用矩陣乘法計算,在卷積核尺寸比較大并且輸出特征圖通道數也比較大的時候,我們會發現這個時候Im2Col獲得矩陣是一個行非常多列非常少的矩陣,在做矩陣乘法的時候訪存會變得比較差,從而計算效率邊地。這是因為當代CPU架構的某些原因,導致程序在行方向上的處理始終比列方向慢一個檔次,所以我們如果能想個辦法使得行方向的元素變少,列方向的元素變多這樣就可以有效改善緩存達到加速的目的,另外列方向上元素增多更容易讓我們發揮Neon指令集的作用。所以,這就是本文將要提到的第一個優化技巧,數據打包(Pack)。
具體來說,對于卷積核我們進行 4×44\times 44×4 的Pack(所謂 4×44 \times 44×4 的Pack就是在Im2Col獲得的二維矩陣的高維度進行壓縮,在寬維度進行膨脹,不知道我這種說法是否合適,參考一下下面的圖應該很好理解),如下圖所示:
這是一個 3×33\times 33×3 的卷積核并且輸出通道數為 444,它經過Im2Col之后首先變成上圖的上半部分,然后經過Pack 4×44\times 44×4 之后變成了上圖的下半部分,即變成了每個卷積核的元素按列方向交織排列。
這部分的代碼也比較容易實現,另外一個技巧是,每次在執行卷積計算時,對于Image的Im2col和Pack每次都會執行,但對于卷積核,Im2col和Pack在任意次只用做一次,所以我們可以在模型初始化的時候提前把卷積核給Pack好,這樣就可以節省卷積核Im2Col和Pack耗費的時間,代碼實現如下:
void ConvolutionLayerSgemm::convolutionTransformKernel(float *const &kernel, const int &kernelW, const int &kernelH, float* &dest, const int &inChannel,const int &outChannel){int kernelSize = kernelH * kernelW;int ccOutChannel = 0;int ccRemainOutChannel = 0;int Stride = 0;ccOutChannel = outChannel >> 2;ccRemainOutChannel = ccOutChannel << 2;for(int cc = 0; cc < ccOutChannel; cc ++){int c = cc << 2;const float* k0 = kernel + c * inChannel * kernelSize;const float* k1 = kernel + (c + 1) * inChannel * kernelSize;const float* k2 = kernel + (c + 2) * inChannel * kernelSize;const float* k3 = kernel + (c + 3) * inChannel * kernelSize;Stride = 4 * kernelSize * inChannel;float* destptr = dest + (c / 4) * Stride;for(int i = 0; i < inChannel * kernelSize; i++){destptr[0] = k0[0];destptr[1] = k1[0];destptr[2] = k2[0];destptr[3] = k3[0];destptr += 4;k0 += 1;k1 += 1;k2 += 1;k3 += 1;}}for(int cc = ccRemainOutChannel; cc < outChannel; cc++){int c = cc;const float* k0 = kernel + c * inChannel * kernelSize;Stride = 4 * kernelSize * inChannel;float* destptr = dest + (c / 4 + c % 4) * Stride;for(int i = 0; i < inChannel * kernelSize; i++){destptr[0] = k0[0];destptr += 1;k0 += 1;}}}
注意這個地方如果Pack的時候有拖尾部分,也就是說outChannel不能被4整除時,那么直接按照原始順序排列即可,后面在Sgemm方法計算的時候需要注意一下。下圖展示了一個 5×4×35\times 4\times 35×4×3 的卷積核經過Im2Col再經過Pack 4×44\times 44×4 之后獲得的結果,可以看到拖尾那一行是直接復制的,不夠的長度直接置0即可。
接下來我們繼續講一下對于輸入數據的Pack,我這里以對輸入數據進行8x8的Pack為例子來講解輸入數據Pack之后應該是什么樣子,以及如何和剛才已經4x4Pack好的卷積核執行矩陣乘法以完成整個卷積過程。至于為什么這里要使用Pack8x8而不是Pack4x4,因為考慮到MsnhNet后面的版本中需要支持Armv8的實現并且無論如何Pack也基本不影響計算的邏輯。下面我們舉個例子來加以說明。
下面展示了一個輸入維度為 [1,1,7,4][1,1,7,4][1,1,7,4] 的特征圖使用Im2Col進行展開并Pack8x8之后獲得結果(卷積核維度為 [1,1,3,3][1,1,3,3][1,1,3,3]),其中左邊代表的是Im2Col后的結果,右邊是將其進一步Pack8x8后獲得的結果。這個地方由于 outHeight×outWidthoutHeight\times outWidthoutHeight×outWidth 不能被8整除,所以存在拖尾部分無法進行Pack,即結果圖中的第二和第三列,直接順序放就可以了。
這部分的代碼實現也比較簡單,如下所示:
// pack 8x8// preapareconst int packChannel = outSize / 8 + outSize % 8;const int packHeight = inChannel; const int packWidth = 8 * kernelSize;int kernelPackChannel = outChannel / 4 + outChannel % 4;const int kernelPackHeight = inChannel;const int kernelPackWidth = 4 * kernelSize;float *src_im2col_pack = new float[packHeight * packWidth * packChannel];// pack startint colCount = outSize >> 3;int remainColCount = colCount << 3;#if USE_OMP#pragma omp parallel for num_threads(OMP_THREAD)
#endiffor(int i = 0; i < colCount; i++){int newi = i << 3;const float *src0 = src_im2col;src0 += newi;float *packptr = src_im2col_pack + i * packHeight * packWidth;for(int j = 0; j < inChannel * kernelSize; j ++){
#if USE_NEON
#if __aarch64__throw Exception(1, "Error: armv8 temporarily not supported!", __FILE__, __LINE__, __FUNCTION__);
#elseasm volatile("pld [%0, #256] \n""vld1.f32 {d0-d3}, [%0] \n""vst1.f32 {d0-d3}, [%1] \n": "=r"(src0), // %0"=r"(packptr) // %1: "0"(src0),"1"(packptr): "memory", "q0", "q1");
#endif#elsepackptr[0] = src0[0];packptr[1] = src0[1];packptr[2] = src0[2];packptr[3] = src0[3];packptr[4] = src0[4];packptr[5] = src0[5];packptr[6] = src0[6];packptr[7] = src0[7];
#endifpackptr += 8;src0 += outSize;}}// pack tail#if USE_OMP#pragma omp parallel for num_threads(OMP_THREAD)
#endiffor(int i = remainColCount; i < outSize; i++){const float *src0 = src_im2col;src0 += i;float *packptr = src_im2col_pack + (i / 8 + i % 8) * packHeight * packWidth;for(int j = 0; j < inChannel * kernelSize; j++){packptr[0] = src0[0];packptr += 1;src0 += outSize;}}
3.4 優化點之帶Pack的矩陣乘法
現在我們已經獲得了輸入特征圖和卷積核的Im2Col+Pack操作后的矩陣,那么我們就可以使用矩陣乘法來計算出最后的結果了,即Sgemm算法,但這里因為Pack的原因不能使用OpenBlas標準庫。這里介紹一下如何手寫一個Sgemm算法。首先Sgemm的接口大概是:
sgemm (int M, int N, int K, float *A, float *B, float *C)
其中輸入矩陣A的特征維度是 M×KM\times KM×K ,輸入矩陣B的特征維度是 K×NK\times NK×N,輸出矩陣 CCC 的特征維度是 M×NM\times NM×N ,不考慮任何優化的情況下復雜度是 O(N×K×N)O(N\times K \times N)O(N×K×N) 。因為我們這里Pack了數據,所以訪存相比于原始版本會變好一些,但計算量實際上還是沒變的。除此之外,由于行方向的數據增多我們可以更好的在行方向上進行利用Neon優化使得整個計算過程的效率更好。
所以就目前的情況來說,矩陣A就是我們的卷積核經過Im2Col+Pack4x4獲得的輸出矩陣,矩陣B就是我們的輸入特征圖經過Im2Col+Pack8x8之后獲得的輸出矩陣,然后矩陣C就是經過矩陣乘法之后獲得的輸出特征圖了。在實現矩陣乘法的時候,以矩陣A為基準一次處理4行,并且在列方向分別處理4或者8個元素,這個可以結合上面的輸入特征圖的Im2Col+Pack8x8的示意圖來想。最后,對于矩陣A不夠4的拖尾部分,進行暴力處理即可。這個地方考慮的是當代典型的CNN架構一般通道數都是 4 的倍數,暫時沒有實現拖尾部分的Neon優化,下面先看一下這個算法實現的整體部分,復刻NCNN。
這部分的代碼實現如下:
// sgemm (int M, int N, int K, float *A, float *B, float *C)
// A (M x K)
// B (K x N)
// C (M x N)//int M = outChannel;int N = outHeight * outWidth;int K = kernelSize * inChannel;int ccOutChannel = outChannel >> 2;int ccRemainOutChannel = ccOutChannel << 2;#if USE_OMP#pragma omp parallel for num_threads(OMP_THREAD)
#endiffor(int cc = 0; cc < ccOutChannel; cc++){int c = cc << 2;float *destptr0 = dest + c * outSize;float *destptr1 = dest + (c + 1) * outSize;float *destptr2 = dest + (c + 2) * outSize;float *destptr3 = dest + (c + 3) * outSize;int i = 0;// N = outHeight*outWidthfor(; i + 7 < N; i = i+8){const float *ptrB = src_im2col_pack + (i / 8) * packHeight * packWidth;
#if __aarch64__throw Exception(1, "Error: armv8 temporarily not supported!", __FILE__, __LINE__, __FUNCTION__);
#elseconst float *ptrA = kernel_im2col_pack + (c / 4) * kernelPackHeight * kernelPackWidth;
#endif#if USE_NEON#if __aarch64__throw Exception(1, "Error: armv8 temporarily not supported!", __FILE__, __LINE__, __FUNCTION__);
#elseasm volatile();
#endif#elsefloat sum0[8] = {0};float sum1[8] = {0};float sum2[8] = {0};float sum3[8] = {0};int j = 0;// K = kernelSize * inChannel// 同時計算4行,同時在每一列計算8個輸出for(; j + 7 < K; j = j + 8){for(int n = 0; n < 8; n++){sum0[n] += ptrA[0] * ptrB[n];sum1[n] += ptrA[1] * ptrB[n];sum2[n] += ptrA[2] * ptrB[n];sum3[n] += ptrA[3] * ptrB[n];ptrA += 4;sum0[n] += ptrA[0] * ptrB[n + 8];sum1[n] += ptrA[1] * ptrB[n + 8];sum2[n] += ptrA[2] * ptrB[n + 8];sum3[n] += ptrA[3] * ptrB[n + 8];ptrA += 4;sum0[n] += ptrA[0] * ptrB[n + 16];sum1[n] += ptrA[1] * ptrB[n + 16];sum2[n] += ptrA[2] * ptrB[n + 16];sum3[n] += ptrA[3] * ptrB[n + 16];ptrA += 4;sum0[n] += ptrA[0] * ptrB[n + 24];sum1[n] += ptrA[1] * ptrB[n + 24];sum2[n] += ptrA[2] * ptrB[n + 24];sum3[n] += ptrA[3] * ptrB[n + 24];ptrA += 4;sum0[n] += ptrA[0] * ptrB[n + 32];sum1[n] += ptrA[1] * ptrB[n + 32];sum2[n] += ptrA[2] * ptrB[n + 32];sum3[n] += ptrA[3] * ptrB[n + 32];ptrA += 4;sum0[n] += ptrA[0] * ptrB[n + 40];sum1[n] += ptrA[1] * ptrB[n + 40];sum2[n] += ptrA[2] * ptrB[n + 40];sum3[n] += ptrA[3] * ptrB[n + 40];ptrA += 4;sum0[n] += ptrA[0] * ptrB[n + 48];sum1[n] += ptrA[1] * ptrB[n + 48];sum2[n] += ptrA[2] * ptrB[n + 48];sum3[n] += ptrA[3] * ptrB[n + 48];ptrA += 4;sum0[n] += ptrA[0] * ptrB[n + 56];sum1[n] += ptrA[1] * ptrB[n + 56];sum2[n] += ptrA[2] * ptrB[n + 56];sum3[n] += ptrA[3] * ptrB[n + 56];ptrA -= 28;}ptrA += 32;ptrB += 64;}// K = kernelSize * inChannel * 4// 如果是pack4x4那么末尾一定是4的倍數for(; j < K; j++){for(int n = 0; n < 8; n++){sum0[n] += ptrA[0] * ptrB[n];sum1[n] += ptrA[1] * ptrB[n];sum2[n] += ptrA[2] * ptrB[n];sum3[n] += ptrA[3] * ptrB[n];}ptrA += 4;ptrB += 8;}for(int n = 0; n < 8; n++){destptr0[n] = sum0[n];destptr1[n] = sum1[n];destptr2[n] = sum2[n];destptr3[n] = sum3[n];}#endifdestptr0 += 8;destptr1 += 8;destptr2 += 8;destptr3 += 8;}
3.5 優化點之Neon的應用
上面的列方向一次處理多個元素的累乘和累加過程可以進一步使用Neon優化,將上面的一次處理4行,每列一次計算8個輸出的代碼翻譯為Neon匯編代碼即可,代碼實現以及簡要的注釋如下:
asm volatile("veor q1, q0, q0 \n""vdup.f32 q8, d2[0] \n""vdup.f32 q9, d2[0] \n""vdup.f32 q10, d2[0] \n""vdup.f32 q11, d2[0] \n""vdup.f32 q12, d2[0] \n""vdup.f32 q13, d2[0] \n""vdup.f32 q14, d2[0] \n""vdup.f32 q15, d2[0] \n"// r4 = K >> 2"lsr r4, %12, #2 \n"// 如果nn等于0,使用beq進行循環跳轉,即跳轉到循環1 "cmp r4, #0 \n""beq loop1 \n"// for(; nn != 0; nn--) && nn = K >> 2"loop0: \n" // kernel q0-q3"pld [%5, #512] \n""vldm %5!, {d0-d7} \n" // input q4-q7"pld [%4, #512] \n""vldm %4!, {d8-d15} \n"//calc// sum0[n] += ptrA[0] * ptrB[n];"vmla.f32 q8, q4, d0[0] \n"// sum1[n] += ptrA[1] * ptrB[n];"vmla.f32 q9, q5, d0[0] \n"// sum0[n] += ptrA[0] * ptrB[n + 8];"vmla.f32 q10, q4, d0[1] \n"// sum1[n] += ptrA[1] * ptrB[n + 8];"vmla.f32 q11, q5, d0[1] \n"// sum0[n] += ptrA[0] * ptrB[n + 16];"vmla.f32 q12, q4, d1[0] \n" // sum1[n] += ptrA[1] * ptrB[n + 16];"vmla.f32 q13, q5, d1[0] \n"// sum0[n] += ptrA[0] * ptrB[n + 24];"vmla.f32 q14, q4, d1[1] \n" // sum1[n] += ptrA[1] * ptrB[n + 24];"vmla.f32 q15, q5, d1[1] \n"// sum2[n] += ptrA[2] * ptrB[n];"vmla.f32 q8, q6, d2[0] \n" // sum3[n] += ptrA[3] * ptrB[n];"vmla.f32 q9, q7, d2[0] \n"// sum2[n] += ptrA[2] * ptrB[n + 8];"vmla.f32 q10, q6, d2[1] \n" // sum3[n] += ptrA[3] * ptrB[n + 8];"vmla.f32 q11, q7, d2[1] \n"// sum2[n] += ptrA[2] * ptrB[n + 16];"vmla.f32 q12, q6, d3[0] \n" // sum3[n] += ptrA[3] * ptrB[n + 16];"vmla.f32 q13, q7, d3[0] \n"// sum2[n] += ptrA[2] * ptrB[n + 24];"vmla.f32 q14, q6, d3[1] \n" // sum3[n] += ptrA[3] * ptrB[n + 24];"vmla.f32 q15, q7, d3[1] \n"// ptrA += 4x4"pld [%4, #512] \n""vldm %4!, {d8-d15} \n"// sum0[n] += ptrA[0] * ptrB[n + 32];"vmla.f32 q8, q4, d4[0] \n" // sum1[n] += ptrA[1] * ptrB[n + 32];"vmla.f32 q9, q5, d4[0] \n"// sum0[n] += ptrA[0] * ptrB[n + 40];"vmla.f32 q10, q4, d4[1] \n" // sum1[n] += ptrA[1] * ptrB[n + 40];"vmla.f32 q11, q5, d4[1] \n"// sum0[n] += ptrA[0] * ptrB[n + 48];"vmla.f32 q12, q4, d5[0] \n" // sum1[n] += ptrA[1] * ptrB[n + 48];"vmla.f32 q13, q5, d5[0] \n"// sum0[n] += ptrA[0] * ptrB[n + 56];"vmla.f32 q14, q4, d5[1] \n" // sum1[n] += ptrA[1] * ptrB[n + 56];"vmla.f32 q15, q5, d5[1] \n"// sum2[n] += ptrA[2] * ptrB[n + 32];"vmla.f32 q8, q6, d6[0] \n" // sum3[n] += ptrA[3] * ptrB[n + 32];"vmla.f32 q9, q7, d6[0] \n"// sum2[n] += ptrA[2] * ptrB[n + 40];"vmla.f32 q10, q6, d6[1] \n" // sum3[n] += ptrA[3] * ptrB[n + 40];"vmla.f32 q11, q7, d6[1] \n"// sum2[n] += ptrA[2] * ptrB[n + 48];"vmla.f32 q12, q6, d7[0] \n"// sum3[n] += ptrA[3] * ptrB[n + 48];"vmla.f32 q13, q7, d7[0] \n"// sum2[n] += ptrA[2] * ptrB[n + 56];"vmla.f32 q14, q6, d7[1] \n" // sum3[n] += ptrA[3] * ptrB[n + 56];"vmla.f32 q15, q7, d7[1] \n""subs r4, r4, #1 \n"// 第一個for循環的結束,nn>0"bne loop0 \n" // 開始寫第二個for循環"loop1: \n"// K = kernelSize * inChannel * 4// K >> 2 == inChannel>>2 = inChannel & 3// 計算完之后進行第三個for循環進行最后的賦值"and r4, %12, #3 \n""cmp r4, #0 \n""beq loop3 \n""loop2: \n" // kernel q0 && ptrA += 4// q0 = [d0, d1] = [ptrA[0], ptrA[1], ptrA[2], ptrA[3]]"pld [%5, #128] \n""vld1.f32 {d0-d1}, [%5]! \n"// input q4, q5 && ptrB += 8// q4, q5 = [d8, d9, d10, d11] = [ptrB[0], ..., ptrB[7]]"pld [%4, #256] \n""vld1.f32 {d8-d11}, [%4]! \n"// for(int n = 0; n < 8; n++){// sum0[n] += ptrA[0] * ptrB[n];// sum1[n] += ptrA[1] * ptrB[n];// sum2[n] += ptrA[2] * ptrB[n];// sum3[n] += ptrA[3] * ptrB[n];// }"vmla.f32 q8, q4, d0[0] \n" "vmla.f32 q9, q5, d0[0] \n""vmla.f32 q10, q4, d0[1] \n""vmla.f32 q11, q5, d0[1] \n""vmla.f32 q12, q4, d1[0] \n" "vmla.f32 q13, q5, d1[0] \n""vmla.f32 q14, q4, d1[1] \n" "vmla.f32 q15, q5, d1[1] \n""subs r4, r4, #1 \n""bne loop2 \n"// 完成賦值"loop3: \n" "vst1.f32 {d16-d19}, [%0] \n""vst1.f32 {d20-d23}, [%1] \n""vst1.f32 {d24-d27}, [%2] \n""vst1.f32 {d28-d31}, [%3] \n": "=r"(destptr0), // %0"=r"(destptr1), // %1"=r"(destptr2), // %2"=r"(destptr3), // %3"=r"(ptrB), // %4"=r"(ptrA) // %5: "0"(destptr0),"1"(destptr1),"2"(destptr2),"3"(destptr3),"4"(ptrB),"5"(ptrA),"r"(K) // %12: "cc", "memory", "r4", "q0", "q1", "q2", "q3", "q4", "q5", "q6", "q7", "q8", "q9", "q10", "q11", "q12", "q13", "q14", "q15");
4. 優化效果
下面給出一些測試數據來驗證一下此算法的有效性,注意算法的正確性筆者已經預先做過一些測試,基本可以保證,不久后此代碼也將合并到MsnhNet的下一個發行版本中。
-
下面的測試結果都默認開啟了Neon優化,單核A53,Armv7a架構。
-
并且所有的優化方法代碼實現都在:https://github.com/msnh2012/Msnhnet 的Arm部分。
輸入特征圖大小 | 特征圖輸入通道數 | 卷積核通道數 | 卷積核大小 | 步長 | 優化方法 | 速度 |
---|---|---|---|---|---|---|
14x14 | 512 | 1024 | 3x3 | 1 | 3x3s1手工優化 | 430ms |
14x14 | 512 | 1024 | 3x3 | 1 | Im2col+Pack+Sgemm | 315ms |
14x14 | 512 | 1024 | 3x3 | 2 | 3x3s2手工優化 | 120.49ms |
14x14 | 512 | 1024 | 3x3 | 2 | Im2col+Pack+Sgemm | 93ms |
可以看到這里Im2Col+Pack+Sgemm的速度明顯快于之前的手動展開版本,在stride為1時優化了100+ms。
下面我們再測一組更大的特征圖看一下此算法的表現,一般特征圖比較大那么通道數就比較小,我們對應修改一下:
輸入特征圖大小 | 特征圖輸入通道數 | 卷積核通道數 | 卷積核大小 | 步長 | 優化方法 | 速度 |
---|---|---|---|---|---|---|
112x112 | 64 | 128 | 3x3 | 1 | 3x3s1手工優化 | 542ms |
112x112 | 64 | 128 | 3x3 | 1 | Im2col+Pack+Sgemm | 588.75ms |
112x112 | 64 | 128 | 3x3 | 2 | 3x3s2手工優化 | 97.43ms |
112x112 | 64 | 128 | 3x3 | 2 | Im2col+Pack+Sgemm | 138.90ms |
什么鬼?竟然降速了,這是為什么呢?
有一個直觀的猜測就是當輸入通道數和卷積核尺寸的乘積也就是卷積核Im2Col獲得的矩陣的寬度比較大時,本文介紹的Im2Col+Pack+Sgemm能帶來加速效果。在輸入特征圖通道數較小并且特征圖較大時甚至會比手工優化更慢,因為這個時候對訪存的優化效果很小,因為這個時候Im2Col的矩陣行數和列數差距并不大,這個時候手工優化帶來的Cache Miss并不會對卷積的計算帶啦很大的影響。因此,此算法并非所有情況適用,需要靈活判斷。
本實驗進行得比較有限,關于何時切換Sgemm策略進行卷積可以關注MsnhNet的后續實現。
5. 后記
最近花了比較久的時間陸陸續續寫完這篇文章并在MsnhNet上實現了這個代碼,希望能給想了解Im2Col以及NCNN的Pack優化策略以及如何手動實現一個類Sgemm算法并對其進行優化的讀者一些幫助。完整代碼可以見:https://github.com/msnh2012/Msnhnet 的Arm端實現部分。另外提供了一個單個OP測試工程,地址為:https://github.com/BBuf/ArmNeonOptimization,目前支持Conv3x3s1,Conv3x3s2,BatchNorm,ConvSgemm 的測試。
6. 致謝
-
https://github.com/Tencent/ncnn
-
https://github.com/msnh2012/Msnhnet
7. 廣告時間
MsnhNet是一款基于純c++的輕量級推理框架,本框架受到darknet啟發。
項目地址:https://github.com/msnh2012/Msnhnet ,歡迎一鍵三連。
本框架目前已經支持了X86、Cuda、Arm端的推理(支持的OP有限,正努力開發中),并且可以直接將Pytorch模型(后面也會嘗試接入更多框架)轉為本框架的模型進行部署,歡迎對前向推理框架感興趣的同學試用或者加入我們一起維護這個輪子。