作者:安平博,Xilinx高级工程师;来源:AI加速微信公众号
前言
AI芯片(这里只谈FPGA芯片用于神经网络加速)的优化主要有三个方面:算法优化,编译器优化以及硬件优化。算法优化减少的是神经网络的算力,它确定了神经网络部署实现效率的上限。编译器优化和硬件优化在确定了算力的基础上,尽量最大化硬件的计算和带宽性能。经历了一年多的理论学习,开始第一次神经网络算法优化的尝试。之所以从一个FPGA开发者转向算法的学习,有几个原因:
第一是神经网络在AI芯片上的部署离不开算法的优化。一个浮点数的计算(加法或者乘法)和定点数的计算消耗的资源差距很大,对于FPGA这样逻辑资源有限的芯片而言,定点计算更加友好,而且能够提升几倍于浮点计算的性能。
第二是神经网络量化压缩需要密切的结合FPGA硬件的特点,需要考虑到FPGA的存储资源,计算符号是否能够被FPGA友好的实现等。在AI加速器项目中,算法和FPGA都有各自的开发者,FPGA会对算法组提出要求,比如激活函数量化,normalization如何做等,然后算法组在这些特定要求下去进行算法优化。如果一个人对FPGA和算法都比较熟悉的话,那么就会更容易发现算法优化的点。
第三是FPGA开发方式的趋势是多样化。使用RTL语言仍然是主要的开发方法,需要一个人有一定的数字电路基础。这种开发方式最底层,所以最灵活,可以更好的去调优。但是同时,FPGA一直渴望去突破固有的开发方式,让一个不懂得硬件的软件开发人员也可以很容易的上手,同时能够缩短开发周期,比如HLS。我相信,随着HLS的发展和FPGA芯片的演进,使用这种方式的开发者会越来越多。在那些算法复杂,更新较快的项目中,HLS更有优势,而在一些对资源,时序,功耗要求更高的项目中,RTL更有优势。当硬件平台逐渐软件化后,必然会对FPGA开发者的算法能力提出更高的要求。
Transformer网络结构
Google在《Attention is all your need》的文章中,提出了使用全attention结构替代LSTM的transformer模型,在翻译任务上取得了更好的成绩。这个网络结构计算量大,计算符号相对简单,有一定的应用,所以适合用于网络加速的展示。结构整体模型如下:
1 embedding
包含了input和output的embedding层,完成词汇到网络输入向量的转化,embedding的矩阵大小取决于词汇量的多少,对于翻译来讲,通常都是巨大的,所以其不适合放在FPGA上进行加速,没有量化的必要。Input和output以及softmax前的linear层都共享相同的参数,这样做的目的,是因为共享input和output权重能够降低word level perplexity,当然也降低了参数存储量。最后的linear使用embedding的权重是为了将网络向量转化为词语出现的logits。
2 positional encoding
Transformer是没有循环网络,为了获取词汇位置关系信息,对词汇进行位置编码。其实就是给每个词汇加上位置偏移,位置偏移函数选择了sin和cos函数:
Pos是词汇位置,i是词汇向量的维度位置。
3 encoder
由多层的multi-head attention和linear组成,multi-headattention和linear之间由norm和add,add是一个residual连接。
Multi-head attention结构如下:
Q,K,V分别是query,key和value,这是attention机制中抽象出来的三个重要变量,通过计算q和k的相似度,得到每个k对应的v的权重系数,然后对value进行加权求和就得到了attention值。这个是attention机制的本质思想。Transformer中使用softmax函数来描述相似度,当然还有很多其它方法来描述。
这里添加了一个scale1/squart(dk),这其实是一个参数的调节,防止矩阵乘法得到结果太大而导致softmax函数的梯度太小。
这里还要注意transformer网络没有对Q,K,V直接进行单一的attention计算,而是对这三个变量进行了拆分,平行计算拆分后的变量,得到的attention值最后在拼接在一起。
4 decoder
Decoder和encoder也有类似的结构,不同的是,在decoder中由三层:mask-multi-head attention,multi-head attention以及FC构成。带mask的multi-head是为了屏蔽target句子词之后的词,因为对句子的翻译应该是由前向后进行的,后边的词语不应该出现在前边词语的预测之中。
量化方法
量化实际是一个仿射变换:
其中s是scale,q是量化后的数据,z是偏移,如果采用对称变换,令z为0,那么就有:
去除中心z,可以消除矩阵计算中的交叉项。接下来就是如何获得q和s。q和s通过如下方式获得:
Clip操作是在最小值n和最大值p之间获得x/s的向下整数值,如果x/s向下整数值超过n或者p就取n和p。
S的值通过训练获得,为了保证能够很好的在FPGA上计算,s的值最好可以取得2的幂次。
由于s和x都是需要训练的参数,所以我们需要求得他们的梯度值,梯度值比较简单,对q(x, s)的x和x进行求导,有:
对x的梯度使用的是hinton提出的strait-through estimator,这样做是因为可以消除量化引入的噪声,更快的训练。
实践
transformer中有dense,matmul等操作,需要量化的数据有dense中的权重,matmul中的Q,V,K变量。第一次没有什么经验,还是一点点来。首先选择其中一个dense进行量化。从github上下载了一个transformer的实现源码https://github.com/Kyubyong/transformer,这个代码写的很简洁,容易看懂。官方的实现代码比较复杂,需要安装的库较多,曾经也尝试过,因为某些库无法安装成功,所以放弃了。在使用Kyubyong的transformer的时候,也遇到了一个问题,训练可以完成,但是在eval的时候,报了维度的错误,后来找到是在positional encoding的embedding中,经过查找,源码中存在一个bug,就是eval的数据集的maxlen是设置了10000,但是在embedding中传入的查找表维度是从hparams传入的,两者不相同。不知道作者为什么会有这个bug。经过改正可以正常完成eval了。
量化第一步是需要将量化插入到tensorflow的图结构中,即在要量化的权重数据之后。这需要重新定义op和梯度,tensorflow中提供了tf.custom_gradient装饰函数来对梯度和op进行定义,所以我定义了如下梯度:
其中STE_clip中的y计算了对x的量化值,grad函数是对x和s进行梯度计算。X和s分别是传入的(d, d)权重和scale。dy是传入的上一个节点的梯度,所以完成和STE_cllip节点梯度的乘积,这是由函数梯度计算的传递性质决定的。这里需要注意的是,s是一个标量,q(x,s)对s梯度是一个矩阵向量,需要和dy进行点积和。
在tensorflow图构建中,将这个节点插入如下:
这里还添加了tf.print用于打印量化后的数据。
语法错误修正:
1 定义的custom_gradient函数中报NoneType object is not iterable,因为函数没有返回值,默认返回none。
2 TypeError: Input 'e' of 'Select' Op has type float32 that does not match type int32 of argument 't'. 因为使用tf.greater(x, y)x和y应该有相同数据类型。
3 ValueError: Shapes must be equal rank,tf.greater中数据必须具有相同的rank,即维度。
4 ValueError: Shape must be rank 1 but is rank 2,tf.tile(x, axis)中x必须是具有维度的,不能够是0维。
5 ValueError: Shape must be rank 2 but is rank 3,tf.matmul中两个矩阵维度必须相同。
6 TypeError: Failed to convert object of type
7 TypeError: Expected int32, got None of type '_Message' instead. 这是因为输入为[N, T, d_model],其中N开始是none的,所以当使用tf.constant([N,1,1])的时候就会出现错误,因为N是none类型。
8 Incompatible shapes between op input and calculated input gradient。输入的数据和对该输入数据的梯度维度不一致。
9 使用tf.print无法打印出数据。这是因为print是tensorflow中的一个节点,需要将这个节点加入图中,然后才能输出。而且只有计算流经这个print节点,其才会发挥作用。形象的描述应该是:
功能问题:
1 首先就是发现在训练过程中scale和量化数据都没有更新,一直保持不变,而且量化值和权重数据以及scale计算的数据不相同。目前还在查找当中。
引用文献
1 Learning Accurate Integer Transformer Machine-Translation Models,Ephrem Wu
2 Trained uniform quantization for accurate and efficient neural network inference on fixedpoint hardware,Sambhav R. Jain, Albert Gural, Michael Wu, and Chris Dick
3 Attention Is All You Need,Ashish Vaswani,Noam Shazeer,Niki Parmar
4 Quantized Neural Networks: Training Neural Networks with Low Precision Weights and Activations,Itay Hubara,Matthieu Courbariaux,Daniel Soudry