在AI芯片:高性能卷积计算中的数据复用曾提到,基于变换域的卷积计算——譬如Winograd卷积——并不能适应算法上对卷积计算多变的需求。但Winograd卷积依旧出现在刚刚公开的ARM Ethos-N57和Ethos-N37 NPUs的支持特性中,本文将利用Nvidia开源的NVIDIA Deep Learning Accelerator (NVDLA)为例,分析在硬件中支持Winograd卷积的实现方式,代价和收益;以期对基于变换域卷积的优势和不足有更深的认识。
1. Windgrad卷积的计算方式
卷积神经网络中的三维卷积(后文简称为卷积)计算过程可以表示如下,将这种直接通过原始定义计算卷积的方式称为直接卷积(Direct Convolution)。
代码语言:javascript复制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];
其中各参数的含义如下表
数据维度 | 描述 |
---|---|
Ho/Wo | 输出feature map的高和宽 |
Co | 输出的channel数目 |
R/S | filter的高和宽 |
Ci | 输入的channel数目 |
s | 卷积计算的stride |
和一般的乘加运算不同,卷积计算中有滑窗的过程,充分利用这一点特性可以节约计算过程中的乘法次数。关于Winograd的原理和推导,。此处直接给出3x3, stride=1卷积下Winograd卷积的形式。 [S = A^Tleft[left(GgG^Tright) odot left( C^TdC right) right]A ]
[ 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} ]
[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} ]
其中(g)是3x3的kernel,(d)是4x4的feature map,(odot)表示矩阵对应位置元素相乘。(s)表示2x2的卷积结果。矩阵(C), (G), (A)为常量,用于Wingrad卷积中的变换。由于(C), (G), (A)中各元素取值为(pm1,pm0.5), 因此计算可以通过加减和简单移位得到,认为不需要进行乘法运算。
因此,采用Winograd卷积计算得到4哥输出结果需要16次乘法计算,而直接卷积需要36次乘法计算。但是由于Winograd在变换中加入了加法计算,因此加法次数会有一定增加。注意上述讨论中并没有加入Channel方向,这是因为此处卷积在Channel上实际上依旧退化成了简单的乘加运算,因此无论在变换前后进行Channel方向计算均没有区别。
一段直接卷积和Winograd卷积对比的代码如下所示
代码语言:javascript复制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)
由计算结果可知,两者结果完全一致(如果采用浮点数时可能会有量化误差,但都在合理范围内)
代码语言:javascript复制>>> 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中给出了可视化的图像用于描述这个过程,这一过程实际上就是卷积的六层循环计算过程的拆解,可以表示如下
代码语言:javascript复制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^Tleft[left(GgG^Tright) odot left( C^TdC right) right]A) 实际上分布在不同的阶段/模块进行。
- $U = GgG^T $是离线预先计算好的
- $V = C^TdC $是在数据路径上计算的
- (S = A^Tleft[ Uodot Vright]A) 是在计算阵列中计算的
首先考虑计算阵列的设计。NVDLA计算3x3卷积,每次输出2x2共计4个数,计算过程中有4x4的矩阵点乘计算;结合直接卷积中64个乘法计算,Winograd卷积同时计算了4个Channel,共计4x4x4=64次乘法。乘法计算本身没有区别,但在进行加法时,和直接卷积略有不同,用代码可表示为
代码语言:javascript复制//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,其设计如下图所示
和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应该作为设计的可选项,这是因为
- 计算3x3卷积有2.25x的理论提升
- Winograd Conv的乘法依旧是矩阵计算
- Winograd Conv的数据路径和直接卷积没有必然的冲突
- Winograd Conv的加法可以直接在数据路径上完成,甚至不影响其他设计
- 如果加速器设计粒度足够细,甚至可以从软件调度上直接支持Winograd Conv
完全不考虑Winograd Conv的理由只可能是未来算法发展趋势下,3x3的普通卷积计算量占比会大大下降。