NVDLA中Winograd卷積的設計

時間 2019-11-09

標籤 nvdla winograd 設計简体版

原文原文鏈接

在AI芯片：高性能卷積計算中的數據複用曾提到，基於變換域的卷積計算——譬如Winograd卷積——並不能適應算法上對卷積計算多變的需求。但Winograd卷積依舊出如今剛剛公開的ARM Ethos-N57和Ethos-N37 NPUs的支持特性中，本文將利用Nvidia開源的NVIDIA Deep Learning Accelerator (NVDLA)爲例，分析在硬件中支持Winograd卷積的實現方式，代價和收益；以期對基於變換域卷積的優點和不足有更深的認識。html

1. Windgrad卷積的計算方式

卷積神經網絡中的三維卷積（後文簡稱爲卷積）計算過程能夠表示以下，將這種直接經過原始定義計算卷積的方式稱爲直接卷積（Direct Convolution）。python

for i = 1 : Ho
    for j = 1 : Wo
        for k = 1 : Co
            for l = 1 : R
                for m = 1 : S
                    for n = 1 : Ci
                        out[i,j,k] += In[i*s+l.j*s+m,n]*F[l,m,n];

其中各參數的含義以下表git

數據維度	描述
Ho/Wo	輸出feature map的高和寬
Co	輸出的channel數目
R/S	filter的高和寬
Ci	輸入的channel數目
s	卷積計算的stride

和通常的乘加運算不一樣，卷積計算中有滑窗的過程，充分利用這一點特性能夠節約計算過程當中的乘法次數。關於Winograd的原理和推導，能夠參考http://www.javashuo.com/article/p-wdkrqpuo-dx.html中的相關內容。此處直接給出3x3, stride=1卷積下Winograd卷積的形式(參見NVDLA Unit)。
\[S = A^T\left[\left(GgG^T\right) \odot \left( C^TdC \right) \right]A \]github

\[ g = \begin{bmatrix} wt_{0,0} & wt_{0,1} & wt_{0,2} \\ wt_{1,0} & wt_{1,1} & wt_{1,2} \\ wt_{2,0} & wt_{2,1} & wt_{2,2} \end{bmatrix} \]算法

\[d = \begin{bmatrix} x_{0,0} & x_{0,1} & x_{0,2} & x_{0,3}\\ x_{1,0} & x_{1,1} & x_{1,2} & x_{1,3}\\ x_{2,0} & x_{2,1} & x_{2,2} & x_{2,3}\\ x_{3,0} & x_{3,1} & x_{3,2} & x_{3,3} \end{bmatrix} \]bash

\[C = \begin{bmatrix} 1 & 0 & 0 & 0 \\ 0 & 1 & -1 & 1 \\ -1& 1 & 1 & 0 \\ 0 & 0 & 0 & -1 \end{bmatrix} \]網絡

\[G = \begin{bmatrix} 1 & 0 & 0 \\ 0.5 & 0.5 & 0.5 \\ 0.5 & -0.5& 0.5 \\ 0 & 0 & 1 \end{bmatrix} \]架構

\[A_T = \begin{bmatrix} 1 & 1 & 1 & 0 \\ 0 & 1 &-1 &-1 \end{bmatrix} \]dom

其中$g$是3x3的kernel，$d$是4x4的feature map，$\odot$表示矩陣對應位置元素相乘。$s$表示2x2的卷積結果。矩陣$C$, $G$, $A$爲常量，用於Wingrad卷積中的變換。因爲$C$, $G$, $A$中各元素取值爲$\pm1,\pm0.5$, 所以計算能夠經過加減和簡單移位獲得，認爲不須要進行乘法運算。ide

所以，採用Winograd卷積計算獲得4哥輸出結果須要16次乘法計算，而直接卷積須要36次乘法計算。可是因爲Winograd在變換中加入了加法計算，所以加法次數會有必定增長。注意上述討論中並無加入Channel方向，這是由於此處卷積在Channel上實際上依舊退化成了簡單的乘加運算，所以不管在變換先後進行Channel方向計算均沒有區別。

一段直接卷積和Winograd卷積對比的代碼以下所示

import numpy as np
g = np.random.randint(-128,127,(3,3))
d = np.random.randint(-128,127,(4,4))
direct_conv = np.zeros((2,2))
for i in range(2):
    for j in range(2):
        for r in range(3):
            for s in range(3):
                direct_conv[i,j] =  direct_conv[i,j] + d[i+r,j+s]*g[r,s]

C = np.array([[1,0,0,0],[0,1,-1,1],[-1,1,1,0],[0,0,0,-1]])
G = np.array([[1,0,0],[0.5,0.5,0.5],[0.5,-0.5,0.5],[0,0,1]])
AT = np.array([[1,1,1,0],[0,1,-1,-1]])
U = G.dot(g).dot(G.transpose())
V = C.transpose().dot(d).dot(C)
wg_conv = AT.dot(U*V).dot(AT.transpose())

print(direct_conv)
print(wg_conv)

由計算結果可知，二者結果徹底一致（若是採用浮點數時可能會有量化偏差，但都在合理範圍內）

>>> print(direct_conv)
[[-23640.    -51.]
 [-10740.   8740.]]
>>> print(wg_conv)
[[-23640.    -51.]
 [-10740.   8740.]]

2. NVDLA中的的直接卷積

在硬件設計過程當中不可能爲直接卷積和Winograd卷積分別設計徹底獨立的計算和控制邏輯，因爲直接卷積有計算靈活和適應性強的特色，各種神經網絡加速器都有支持。所以，Winograd必定是創建在直接卷積硬件結構基礎上的拓展功能。在探究NVDLA中的Winograd卷積設計以前，必須先明確NVDLA中的的直接卷積的計算方式。

Nvidia的相關文檔中十分詳細的NVDLA計算直接卷積的流程（NVDLA Unit)，其將卷積計算分紅了五級（下述描述中，以數值精度爲Int16爲例）

Atomic Operation （原子操做，完成16次64次乘法並將其加在一塊兒）
Stripe Operation （條帶操做，完成16次獨立的Atomic Operation）
Block Operation （塊操做，完成kernel的R/S方向的累加）
Channel Operation（通道操做，完成Channel方向計算的累加）
Group Operation （分組操做，完成一組kernel的所有計算）

NVDLA Unit中給出了可視化的圖像用於描述這個過程，這一過程實際上就是卷積的六層循環計算過程的拆解，能夠表示以下

for k = 1 : Co/16
    for i = 1 : Ho/4                                // Group Operation
        for j = 1 : Wo/4                            // Group Operation
            for n = 1 : Ci/64                       // Channel Operation
                for l = 1 : R                       // Block Operation
                    for m = 1 : S                   // Block Operation
                        for ii = 1:4                // Strip Operation
                            for ji = 1:4            // Strip Operation
                                for ki = 1：16      // Antomic Operation
                                    for ni = 1：64  // Antomic Block
                                        out[i*4+ii,j*4+jj,k*16+ki] += 
                                        In[(i*4+ii)*s+l.(j*4+jj)*s+m,n*64+ni]*F[l,m,n*64+ni];

其中，Atomic Operation決定了NVDLA乘法陣列的設計。根據設計能夠看出，NVDLA有16份徹底一致的乘法陣列用於計算16個不一樣Kernel的乘法；而每一個乘法陣列中有64個乘法和一棵64輸入的加法樹。

計算順序還必定程度肯定了NVDLA的Buffer設計和數據路徑設計。在計算直接卷積時，每週期須要128Byte的Feature/Pixel數據，實際上時規則的64Channel的數據；所以在存儲時只須要每一個Bank上存儲64Channel數據，使用時經過MUX選出指定Bank數據便可。在進行結果寫回時，每週期須要寫回16個Feature數據。因爲Winograd卷積使用的Weight能夠提早算好，對比直接卷積和Winograd卷積時能夠忽略Weight路徑。

3. NVDLA中的Winograd卷積

創建在直接卷積的硬件架構上，NVDLA針對Winograd卷積進行了一系列的修改。從計算方式上來講，再也不同時計算64個Channel的乘加；從硬件架構上來講，進行了計算修改和數據路徑修改。根據NVDLA的設計，Winograd卷積的計算$S = A^T\left[\left(GgG^T\right) \odot \left( C^TdC \right) \right]A$ 實際上分佈在不一樣的階段/模塊進行。

$U = GgG^T $是離線預先計算好的
$V = C^TdC $是在數據路徑上計算的
$S = A^T\left[ U\odot V\right]A$ 是在計算陣列中計算的

首先考慮計算陣列的設計。NVDLA計算3x3卷積，每次輸出2x2共計4個數，計算過程當中有4x4的矩陣點乘計算；結合直接卷積中64個乘法計算，Winograd卷積同時計算了4個Channel，共計4x4x4=64次乘法。乘法計算自己沒有區別，但在進行加法時，和直接卷積略有不一樣，用代碼可表示爲

//direct conv & winograd conv
for i = 1:16
    s1[i] = s0[i*4+0] + s0[i*4+1] + s0[i*4+2] + s0[i*4+3];

//direct conv
for i = 1:8
    s2[i] = s1[i*2+0] + s1[i*2+1];
for i = 1:4
    s3[i] = s2[i*2+0] + s2[i*2+1];
s4[i] = s3[0] + s3[1] + s3[2] + s3[3];

//winograd conv
for i=1:4
    s2_wg[0][i] = s1[i*4+0] + s1[i*4+1] + s1[i*4+2];
    s2_wg[0][i] = s1[i*4+1] - s1[i*4+2] + s1[i*4+3];
s3_wg[0][0] = s2_wg[0][0] + s2_wg[0][1] + s2_wg[0][2];
s3_wg[1][0] = s2_wg[1][0] + s2_wg[1][1] + s2_wg[1][2];
s3_wg[0][1] = s2_wg[0][1] - s2_wg[0][1] - s2_wg[0][2];
s3_wg[1][1] = s2_wg[1][1] - s2_wg[1][1] - s2_wg[1][2];

代碼中只有第一級的加法被direct conv和winograd conv徹底複用，其餘級的加法略有不一樣。在NVDLA中，加法是使用Wallace Tree完成的，以提升性能下降資源佔用。Direct Conv中和Winograd Conv中的後面幾級加法還進行了進一步複用。整體來講，從代碼上看（參見NV_NVDLA_CMAC_CORE_mac.v），爲了支持Winograd卷積

加法的第三級中增長了4棵4-2的Wallace Tree Compressor
加法的第四級中增長了2棵4-2的Wallace Tree Compressor
加法的第五級中增長了2棵6-2的Wallace Tree Compressor
增長了一些MUX以direct conv和winograd conv

其次考慮數據路徑，包括讀取的數據路徑和寫回的數據路徑。對於讀取而言，除了須要針對Winograd專門設計取址邏輯和數據選擇邏輯，還須要完成$V = C^TdC $的計算；根據文檔描述，這一計算過程是在PRA（Pre-addition）中完成的。從代碼上看（參見NV_NVDLA_CSC_dl.v）

針對Winograd的地址生成增長的控制邏輯能夠忽略
針對Winograd的數據選擇增長數千的寄存器
PRA採用MENTOR的HLS綜合工具實現，共實現了4份，和MAC陣列（1024乘加）對比，此處的計算資源較少

對於寫回路徑而言，爲了完成卷積計算，在乘加後增長了累加器和SRAM，其設計以下圖所示(ref. http://nvdla.org/_images/ias_image21_cacc.png)

和Direct Conv一次輸出16個結果相比，Winograd Conv輸出的結果爲64。這意味着爲了支持Winograd Conv，須要額外增長48組高位寬的累加器。同時，SRAM的大小也須要設置爲原先的四倍。

4. 相關討論

NVDLA爲了同時支持Direct Conv和Winograd Conv顯然付出了一些代價。定性的分析來看，包括

4組PRA，每組PRA中約有8次加法
16棵加法樹，每棵增長了約8次加法
48組高位寬加法
增長了約25KB的Accumulator SRAM

而做爲對比，一些典型數據包括

MAC陣列中有1024次乘法和約1024次加法
用於存放Feature/Pixel/Weight的Buffer大小爲512KB

顯然，爲了支持Winograd Conv增長的資源並不會太多。固然，雖然讀取路徑和計算陣列的設計受Winograd Conv的影響不大；可是對於寫回路徑而言，數據位寬發生了變化，必定程度影響了總體的架構設計。可能能夠優化的地方包括將Direct Conv的輸出也改爲2x2的大小，這樣寫回的數據路徑上Direct Conv和Winograd Conv就沒有差異了。

NVDLA是一個相對專用的加速器，從相關文檔中也能夠看出，NVDLA專門針對計算中的各類特性/數據排列進行了硬件上的處理。而現有的不少加速器，爲了兼顧不一樣網絡的計算效率，每每更爲靈活。在這種狀況下，Winograd Conv應該做爲設計的可選項，這是由於