卷積神經網絡中的Winograd快速卷積算法

目錄

• 寫在前面
• 問題定義
• 一個例子 F(2, 3)
• 1D winograd
• 1D to 2D，F(2, 3) to F(2x2, 3x3)
• 卷積神經網絡中的Winograd
• 總結
• 參考

博客：blog.shinelee.me | 博客園 | CSDN

寫在前面

隨便翻一翻流行的推理框架（加速器），如NCNN、NNPACK等，可以看到，對於卷積層，大家不約而同地采用了Winograd快速卷積算法，該算法出自CVPR 2016的一篇 paper：Fast Algorithms for Convolutional Neural Networks。

本文將嘗試揭開Winograd算法的神秘面紗。

問題定義

將一維卷積運算定義為\(F(m, r)\)，\(m\)為Output Size，\(r\)為Filter Size，則輸入信號的長度為\(m+r-1\)，卷積運算是對應位置相乘然后求和，輸入信號每個位置至少要參與1次乘法，所以乘法數量最少與輸入信號長度相同，記為

\[\mu(F(m, r))=m+r-1 \]

在行列上分別進行一維卷積運算，可得到二維卷積，記為\(F(m\times n, r\times s)\)，輸出為\(m\times n\)，卷積核為\(r\times s\)，則輸入信號為\((m+r-1)(n+s-1)\)，乘法數量至少為

\[\begin{aligned} \mu(F(m \times n, r \times s)) &=\mu(F(m, r)) \mu(F(n, s)) \\ &=(m+r-1)(n+s-1) \end{aligned} \]

若是直接按滑動窗口方式計算卷積，一維時需要\(m\times r\)次乘法，二維時需要\(m\times n \times r \times s\)次乘法，遠大於上面計算的最少乘法次數。

使用Winograd算法計算卷積快在哪里？一言以蔽之：快在減少了乘法的數量，將乘法數量減少至\(m+r-1\)或\((m+r-1)(n+s-1)\)。

怎么減少的？請看下面的例子。

一個例子 F(2, 3)

先以1維卷積為例，輸入信號為\(d=\left[ \begin{array}{llll}{d_{0}} & {d_{1}} & {d_{2}} & {d_{3}}\end{array}\right]^{T}\)，卷積核為\(g=\left[ \begin{array}{lll}{g_{0}} & {g_{1}} & {g_{2}}\end{array}\right]^{T}\)，則卷積可寫成如下矩陣乘法形式：

\[F(2, 3) = \left[ \begin{array}{lll}{d_{0}} & {d_{1}} & {d_{2}} \\ {d_{1}} & {d_{2}} & {d_{3}}\end{array}\right] \left[ \begin{array}{l}{g_{0}} \\ {g_{1}} \\ {g_{2}}\end{array}\right]=\left[ \begin{array}{c}{r_0} \\ {r_1}\end{array}\right] \]

如果是一般的矩陣乘法，則需要6次乘法和4次加法，如下：

\[\begin{array}{l}{r_{0}=\left(d_{0} \cdot g_{0}\right)+\left(d_{1} \cdot g_{1}\right)+\left(d_{2} \cdot g_{2}\right)} \\ {r_{1}=\left(d_{1} \cdot g_{0}\right)+\left(d_{2} \cdot g_{1}\right)+\left(d_{3} \cdot g_{2}\right)}\end{array} \]

但是，卷積運算中輸入信號轉換成的矩陣不是任意矩陣，其中有規律地分布着大量的重復元素，比如第1行和第2行的\(d_1\)和\(d_2\)，卷積轉換成的矩陣乘法比一般矩陣乘法的問題域更小，這就讓優化存在了可能。

Winograd是怎么做的呢？

\[F(2,3)=\left[ \begin{array}{lll}{d_{0}} & {d_{1}} & {d_{2}} \\ {d_{1}} & {d_{2}} & {d_{3}}\end{array}\right] \left[ \begin{array}{l}{g_{0}} \\ {g_{1}} \\ {g_{2}}\end{array}\right]=\left[ \begin{array}{c}{m_{1}+m_{2}+m_{3}} \\ {m_{2}-m_{3}-m_{4}}\end{array}\right] \]

其中，

\[\begin{array}{ll}{m_{1}=\left(d_{0}-d_{2}\right) g_{0}} & {m_{2}=\left(d_{1}+d_{2}\right) \frac{g_{0}+g_{1}+g_{2}}{2}} \\ {m_{4}=\left(d_{1}-d_{3}\right) g_{2}} & {m_{3}=\left(d_{2}-d_{1}\right) \frac{g_{0}-g_{1}+g_{2}}{2}}\end{array} \]

乍看上去，為了計算\(\begin{array}{l}{r_{0}=m_1 + m_2 + m_3 } \\ {r_{1}=m_2 - m_3 - m_4}\end{array}\)，需要的運算次數分別為：

• 輸入信號\(d\)上：4次加法（減法）
• 卷積核\(g\)上：3次加法（\(g_1+g_2\)中間結果可保留），2次乘法（除法）
• 輸出\(m\)上：4次乘法，4次加法

在神經網絡的推理階段，卷積核上的元素是固定的，因此\(g\)上的運算可以提前算好，預測階段只需計算一次，可以忽略，所以一共所需的運算次數為\(d\)與\(m\)上的運算次數之和，即4次乘法和8次加法。

與直接運算的6次乘法和4次加法相比，乘法次數減少，加法次數增加。在計算機中，乘法一般比加法慢，通過減少減法次數，增加少量加法，可以實現加速。

1D winograd

上一節中的計算過程寫成矩陣形式如下：

\[Y=A^{T}\left[(G g) \odot\left(B^{T} d\right)\right] \]

其中，\(\odot\)為element-wise multiplication（Hadamard product）對應位置相乘，

\[B^{T}=\left[ \begin{array}{cccc}{1} & {0} & {-1} & {0} \\ {0} & {1} & {1} & {0} \\ {0} & {-1} & {1} & {0} \\ {0} & {1} & {0} & {-1}\end{array}\right] \]

\[G=\left[ \begin{array}{ccc}{1} & {0} & {0} \\ {\frac{1}{2}} & {\frac{1}{2}} & {\frac{1}{2}} \\ {\frac{1}{2}} & {-\frac{1}{2}} & {\frac{1}{2}} \\ {0} & {0} & {1}\end{array}\right] \]

\[A^{T}=\left[ \begin{array}{llll}{1} & {1} & {1} & {0} \\ {0} & {1} & {-1} & {-1}\end{array}\right] \]

\[g=\left[ \begin{array}{lll}{g_{0}} & {g_{1}} & {g_{2}}\end{array}\right]^{T} \]

\[d=\left[ \begin{array}{llll}{d_{0}} & {d_{1}} & {d_{2}} & {d_{3}}\end{array}\right]^{T} \]

• \(g\)：卷積核
• \(d\)：輸入信號
• \(G\)：Filter transform矩陣，尺寸\((m+r-1)\times r\)
• \(B^T\)：Input transform矩陣，尺寸\((m+r-1)\times (m+r-1)\)
• \(A^T\)：Output transform矩陣，尺寸\(m \times (m+r-1)\)

整個計算過程在邏輯上可以分為4步：

• Input transform
• Filter transform
• Hadamar product
• Output transform

注意，這里寫成矩陣形式，並不意味着實現時要調用矩陣運算的接口，一般直接手寫計算過程速度會更快，寫成矩陣只是為了數學形式。

1D to 2D，F(2, 3) to F(2x2, 3x3)

上面只是看了1D的一個例子，2D怎么做呢？

論文中一句話帶過：

A minimal 1D algorithm F(m, r) is nested with itself to obtain a minimal 2D algorithm,F(m×m, r×r).

\[Y=A^{T}\left[\left[G g G^{T}\right] \odot\left[B^{T} d B\right]\right] A \]

其中，\(g\)為\(r \times r\) Filter，\(d\)為\((m+r-1)\times (m+r-1)\)的image tile。

問題是：怎么nested with itself？

這里繼續上面的例子\(F(2, 3)\)，擴展到2D，\(F(2\times 2, 3 \times 3)\)，先寫成矩陣乘法，見下圖，圖片來自SlideShare，注意數學符號的變化，

將卷積核的元素拉成一列，將輸入信號每個滑動窗口中的元素拉成一行。注意圖中紅線划分成的分塊矩陣，每個子矩陣中重復元素的位置與一維時相同，同時重復的子矩陣也和一維時相同，如下所示

令\(D_0 = [k_0, k_1, k_2, k_3]^T\)，即窗口中的第0行元素，\(D_1 \ D_2 \ D_3\)表示第1、2、3行；\(W_0=[w_0, w_1, w_2]^T\)，

\[\begin{aligned} \left[ \begin{array}{c}{r_0} \\ {r_1} \\ {r_2} \\ {r_3}\end{array}\right] &= \left[ \begin{array}{c}{R_0} \\ {R_1}\end{array}\right] = \left[ \begin{array}{c}{K_0 W_0 + K_1 W_1 + K_2 W_2} \\ {K_1 W_0 + K_2 W_1 + K_3 W_2} \end{array} \right] \\ &= \left[ \begin{array}{c} {A^{T}\left[(G W_0) \odot\left(B^{T} D_0 \right)\right] + A^{T}\left[(G W_1) \odot\left(B^{T} D_1 \right)\right] + A^{T}\left[(G W_2) \odot\left(B^{T} D_2 \right)\right]} \\ {A^{T}\left[(G W_0) \odot\left(B^{T} D_1 \right)\right] + A^{T}\left[(G W_1) \odot\left(B^{T} D_2 \right)\right] + A^{T}\left[(G W_2) \odot\left(B^{T} D_3 \right)\right]} \end{array} \right] \\ \\ &=A^{T}\left[\left[G [W_0 \ W_1 \ W_2 ] G^{T}\right] \odot\left[B^{T} [d_0 \ d_1 \ d_2 \ d_3] B\right]\right]A \\ \\ &=A^{T}\left[\left[G g G^{T}\right] \odot\left[B^{T} d B\right]\right] A \end{aligned} \]

卷積運算為對應位置相乘再相加，上式中，\(A^{T}\left[(G W_0) \odot\left(B^{T} D_0 \right)\right]\)為列向量\(W_0\)與\(D_0\)的卷積，結果為長度為2的列向量，而\(A^{T}\left[(G W_0) \odot\left(B^{T} D_0 \right)+ (G W_1) \odot\left(B^{T} D_1 \right) + (G W_2) \odot\left(B^{T} D_2 \right)\right]\)方括號內對應位置相乘再相加，相當於在構成的行向量上卷積，據此，上面的推導就不難看出了。

卷積運算為對應位置相乘再相加，上式中，\(A^{T}\left[(G W_0) \odot\left(B^{T} D_0 \right)\right]\)表示長度為4的\(D_0\)與長度為3的\(W_0\)卷積結果，結果為長度為2的列向量，其中，\((G W_0)\)和\((B^{T} D_0)\)均為長度為4的列向量，進一步地，\(\left[(G W_0) \odot\left(B^{T} D_0 \right)+ (G W_1) \odot\left(B^{T} D_1 \right) + (G W_2) \odot\left(B^{T} D_2 \right)\right]\)可以看成3對長度為4的列向量兩兩對應位置相乘再相加，結果為長度為4的列向量，也可以看成是4組長度為3的行向量的點積運算，同樣，\(\left[(G W_0) \odot\left(B^{T} D_1 \right)+ (G W_1) \odot\left(B^{T} D_2 \right) + (G W_2) \odot\left(B^{T} D_3 \right)\right]\)也是4組長度為3的行向量的內積運算，考慮兩者的重疊部分\((B^T D_1)\)和\((B^T D_2)\)，恰好相當於\(G [W_0 \ W_1 \ W_2 ]\)的每一行在\(B^{T} [d_0 \ d_1 \ d_2 \ d_3]\)的對應行上進行1維卷積，上面我們已經進行了列向量卷積的Winograd推導，行向量的卷積只需將所有左乘的變換矩陣轉置后變成右乘就可以了，至此，上面的推導結果就不難得出了。

所謂的nested with itself如下圖所示，

此時，Winograd算法的乘法次數為16（上圖\(4\times 4\)），而直接卷積的乘法次數為36，降低了2.25倍的乘法計算復雜度。

卷積神經網絡中的Winograd

要將Winograd應用在卷積神經網絡中，還需要回答下面兩個問題：

• 上面我們僅僅是針對一個小的image tile，但是在卷積神經網絡中，feature map的尺寸可能很大，難道我們要實現\(F(224, 3)\)嗎？
• 在卷積神經網絡中，feature map是3維的，卷積核也是3維的，3D的winograd該怎么做？

第一個問題，在實踐中，會將input feature map切分成一個個等大小有重疊的tile，在每個tile上面進行winograd卷積。

第二個問題，3維卷積，相當於逐層做2維卷積，然后將每層對應位置的結果相加，下面我們會看到多個卷積核時更巧妙的做法。

這里直接貼上論文中的算法流程：

整體仍可分為4步，

• Input transform
• Filter transform
• Batched-GEMM（批量矩陣乘法）
• Output transform

算法流程可視化如下，圖片出自論文Sparse Winograd Convolutional neural networks on small-scale systolic arrays，與算法對應着仔細推敲還是挺直觀的。

注意圖中的Matrix Multiplication，對應3維卷積中逐channel卷積后的對應位置求和，相當於\((m+r-1)^2\)個矩陣乘積，參與乘積的矩陣尺寸分別為\(\lceil H / m\rceil\lceil W / m\rceil \times C\)和\(C \times K\)，把Channel那一維消掉。

總結

• Winograd算法通過減少乘法次數來實現提速，但是加法的數量會相應增加，同時需要額外的transform計算以及存儲transform矩陣，隨着卷積核和tile的尺寸增大，就需要考慮加法、transform和存儲的代價，而且tile越大，transform矩陣越大，計算精度的損失會進一步增加，所以一般Winograd只適用於較小的卷積核和tile（對大尺寸的卷積核，可使用FFT加速），在目前流行的網絡中，小尺寸卷積核是主流，典型實現如\(F(6\times 6, 3\times 3)\)、\(F(4\times 4, 3\times 3)\)、\(F(2\times 2, 3\times 3)\)等，可參見NCNN、FeatherCNN、ARM-ComputeLibrary等源碼實現。
• 就卷積而言，Winograd算法和FFT類似，都是先通過線性變換將input和filter映射到新的空間，在那個空間里簡單運算后，再映射回原空間。
• 與im2col+GEMM+col2im相比，winograd在划分時使用了更大的tile，就划分方式而言，\(F(1\times 1, r\times r)\)與im2col相同。

參考

• arxiv: Fast Algorithms for Convolutional Neural Networks
• video: Fast Algorithms for Convolutional Neural Networks by Andrew Lavin and Scott Gray
• video: Even Faster CNNs Exploring the New Class of Winograd Algorithms
• arxiv: Sparse Winograd Convolutional neural networks on small-scale systolic arrays
• ARM-software/ComputeLibrary