动手写一个简单版的谷歌TPU-矩阵乘法和卷积
作者:互联网
谷歌TPU是一个设计良好的矩阵计算加速单元,可以很好的加速神经网络的计算。本系列文章将利用公开的TPU V1相关资料,对其进行一定的简化、推测和修改,来实际编写一个简单版本的谷歌TPU。计划实现到行为仿真为止,仅为更确切的了解TPU的优势和局限性,暂无在FPGA等硬件上进一步实现的计划。
系列目录
基本单元-归一化和池化(待发布)
TPU中的Instruction (待完成)
SimpleTPU实例: (计划中)
拓展
TPU的边界(规划中)
重新审视深度神经网络中的并行(规划中)
本文将对TPU中的矩阵计算单元进行分析,并给出了SimpleTPU中32×32的脉动阵列的实现方式和采用该阵列进行卷积计算的方法,并给出了一个卷积的设计实例,验证了其正确性。代码地址https://github.com/cea-wind/SimpleTPU/tree/master/lab1
1. 脉动阵列和矩阵计算
脉动阵列是一种复用输入数据的设计,对于TPU中的二维脉动阵列,很多文章中构造了脉动阵列的寄存器模型,导致阅读较为困难,而实际上TPU中的二维脉动阵列设计思路十分直接。譬如当使用4×4的脉动阵列计算4×4的矩阵乘法时,有
如上图所示,右侧是一个乘加单元的内部结构,其内部有一个寄存器,在TPU内对应存储Weight,此处存储矩阵B。左图是一个4×4的乘加阵列,假设矩阵B已经被加载到乘加阵列内部;显然,乘加阵列中每一列计算四个数的乘法并将其加在一起,即得到矩阵乘法的一个输出结果。依次输入矩阵A的四行,可以得到矩阵乘法的结果。
由于硬件上的限制,需要对传播路径上添加寄存器,而添加寄存器相对于在第i个时刻处理的内容变成了i+1时刻处理;这一过程可以进行计算结果上的等效。如下图所示,采用z-1代表添加一个延时为1的寄存器,如果在纵向的psum传递路径上添加寄存器,为了保证结果正确,需要在横向的输入端也添加一个寄存器(即原本在i进行乘加计算的两个数均在i+1时刻进行计算)。给纵向每个psum路径添加寄存器后,输入端处理如右图所示。(下图仅考虑第一列的处理)
当在横向的数据路径上添加寄存器时,只要每一列都添加相同延时,那么计算结果会是正确的,但是结果会在后一个周期输出,如下图所示
上述分析可以,一个4×4的乘加阵列可以计算一组4×4的乘加阵列完成计算,而对于其他维度的乘法,则可以通过多次调用的方式完成计算。譬如(4×4)×(4×8),可以将(4×8)的乘法拆分乘两个4×4的矩阵乘;而对于(4×8)×(8×4),两个矩阵计算完成后还需要将其结果累加起来,这也是为何TPU在乘加阵列后需要添加Accumulators的原因。最终脉动阵列设计如下所示(以4×4为例)
2. 脉动阵列的实现
如第一节所述,可通过HLS构建一个脉动阵列并进行仿真。关键代码包括,
1. Feature向右侧移动
for(int j=0;j<MXU_ROWNUM;j++){ for(int k=MXU_ROWNUM+MXU_COLNUM-2;k>=0;k--){ if(k>0) featreg[j][k] = featreg[j][k-1]; else if(i<mxuparam.ubuf_raddr_num) featreg[j][k] = ubuf[ubuf_raddr][j]; else featreg[j][k] = 0; } }
2. 乘法计算以及向下方移动的psum
for(int j=MXU_ROWNUM-1;j>=0;j--){ for(int k=0;k<MXU_COLNUM;k++){ ap_int<32> biasreg; biasreg(31,24)=weightreg[MXU_ROWNUM+0][k]; biasreg(23,16)=weightreg[MXU_ROWNUM+1][k]; biasreg(15, 8)=weightreg[MXU_ROWNUM+2][k]; biasreg( 7, 0)=weightreg[MXU_ROWNUM+3][k]; if(j==0) psumreg[j][k] = featreg[j][k+j]*weightreg[j][k] + biasreg; else psumreg[j][k] = featreg[j][k+j]*weightreg[j][k] + psumreg[j-1][k]; } }
完成代码编写后可进行行为级仿真,可以看出整个计算阵列的时延关系
1. 对于同一列而言,下一行的输入比上一行晚一个周期
2. 对于同一行而言,下一列的输入比上一列晚一个周期(注意同一行输入数据是一样的)
3. 下一列的输出结果比上一列晚一个周期
3. 从矩阵乘法到三维卷积
卷积神经网络计算过程中,利用kh×kw×C的卷积核和H×W×C的featuremap进行乘加计算。以3×3卷积为例,如下图所示,省略Channel方向,拆分kh和kw方向分别和featuremap进行卷积,可以得到9个输出结果,这9个输出结果按照一定规律加在一起,就可以得到追后的卷积计算结果。下图给出了3×3卷积,padding=2时的计算示意图。按F1-F9给9个矩阵乘法结果编号,输出featuremap中点(2,1)——指第二行第一个点——是F1(1,1),F2(1,2),F3(1,3),F4(2,1),F5(2,2),F6(2,3),F7(3,1),F8(3,2),F9(3,3)的和。
下面的MATLAB代码阐明了这种计算三维卷积的方式,9个结果错位相加的MATLAB代码如下所示
output = out1; output(2:end,2:end,:) = output(2:end,2:end,:) + out2(1:end-1,1:end-1,:); output(2:end,:,:) = output(2:end,:,:) + out3(1:end-1,:,:); output(2:end,1:end-1,:) = output(2:end,1:end-1,:) + out4(1:end-1,2:end,:); output(:,2:end,:) = output(:,2:end,:) + out5(:,1:end-1,:); output(:,1:end-1,:) = output(:,1:end-1,:) + out6(:,2:end,:); output(1:end-1,2:end,:) = output(1:end-1,2:end,:) + out7(2:end,1:end-1,:); output(1:end-1,:,:) = output(1:end-1,:,:) + out8(2:end,:,:); output(1:end-1,1:end-1,:) = output(1:end-1,1:end-1,:) + out9(2:end,2:end,:);
而在实际的HLS代码以及硬件实现上,部分未使用的值并未计算,因此实际计算的index和上述示意图并不相同,具体可参考testbench中的配置方法。
4. 其他
GPU的volta架构中引入了Tensor Core来计算4×4的矩阵乘法,由于4×4的阵列规模较小,其内部可能并没有寄存器,设计可能类似第一节图1所示。由于其平均一个周期就能完成4×4矩阵计算,猜测采用第一节中阵列进行堆叠,如下图所示。
一些FPGA加速库中利用脉动阵列实现了矩阵乘法,不过不同与TPU中将一个输入固定在MAC内部,还可以选择将psum固定在MAC内部,而两个输入都是时刻在变化的。这几种方式是类似的,就不再展开描述了。
标签:end,卷积,矩阵,TPU,计算,脉动阵列,output,乘法 来源: https://www.cnblogs.com/sea-wind/p/10995360.html