Qian's Blog HOME CATEGORY REFERENCE ABOUT

神经网络模型压缩——软硬件的相爱相杀

一句废话

笔者作为一个芯片人,有幸被这个时代洪流卷入其中,既和神经网络模型压缩创始人喝过酒唱过歌,又和各路xPU同过场P过K,做了几年下来,对AI芯片这个领域有些思考。这里先从神经网络模型压缩这件事儿来看看AI芯片设计中软硬件的相爱相杀。篇幅有限,浅谈一二。

网络压缩的手段

提到压缩,其核心就是要减小数据量,从而获得以下几个好处:

那到底有多少手段可以达到模型压缩的目的呢?在这些手段中又蕴含了哪些相爱相杀呢?我们不妨从以下几个纬度简单探讨一下。

数字量化

我们知道,计算机系统中常用的数据格式有双精度(FP64)、单精度(FP32)、整型(INT)。早年间AI的应用都是在CPU/GPU上进行开发部署的,因为它们对浮点数的支持较好,且浮点数的精度表现不错,所以FP32一度成为了算法工程师的首选。 从硬件设计的角度来看,设计师更关心的是面积、功耗以及设计复杂度,显然单精度优于双精度,定点数优于浮点数。但从软件算法的角度来看,设计师更关心的是精度,显然得出的结论与硬件设计师相反。如何在精度和硬件复杂度之间找到一个平衡,这就是量化需要完成的工作。

最常见的量化,是把FP32量化到FP16或者INT8,无数的例子已经证实,采用合适的量化以及Fine-Tune策略,FP32量化成FP16或INT8是基本可行的,精度不损失的情况下,模型的尺寸可以下降到一半甚至四分之一。我们不禁会想,既然FP16和INT8可行,那FP7可不可以?INT4可不可以?激进一些,INT2甚至BOOL类型可不可以?以笔者的经验,或许真可以,但很难推进。之所以说可以,是因为在学术界已经有不少论文在讨论低比特网络,但在产业界,如何协同算法工程师、业务工程师和芯片工程师共推进成这件事情,是个难点。不可否认,现阶段AI芯片还不成熟,算法遥遥领先硬件,很难要求业务逻辑迁就硬件架构。可喜的是,NVIDIA的最新GPU中已经开始支持INT4,相信不久的将来,软件工程师会更有动力和硬件工程师配合推进低比特网络的部署。

再来谈谈量化方式,我们粗暴的把量化方式分成两种,均匀量化和非均匀量化。前者又称线性编码,也就是说各个量化区间的宽度相等,从硬件角度来说,好处不言而喻,计算起来方便,但是从软件角度来看,这未必是一个很好的解决方案,因为神经网络中的数据动态范围比较广,对于数值比较小的区间,丢失的精度比较高。而非均匀量化可以很好的解决这个问题,无疑从软件角度来看,非均匀量化才是平衡精度和数据尺寸的最佳解决方案,然而这对硬件的设计又提出了挑战,甚至得用查找表之类的结构来实现运算,对于神经网络这种并行度巨大的硬件架构而言,查找表未必是一个可以接受的方案。

最后,不得不提一个学术界关于网络量化的研究成果。对于一个已知的多层网络而言,往往前几层和后几层的量化精度要求比较高,而中间层量化精度要求比较低,因此为了保证神经网络的精度,一种量化长度可变的网络被提了出来。遗憾的是,这种结构对于硬件设计并不太友好,而且对于算法工程师而言要求也比较高,因此,并没有在业界广泛使用。

模型剪枝

五年前,松哥的一篇《Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding》开创了神经网络模型压缩的先河,其中最重要的一个贡献就是提供了一个叫做“剪枝”的方法,用以去掉网络中冗余的连接。用松哥自己的话说,这个灵感来源有点仿生学的意味。因为研究发现,婴儿大脑的神经元个数比成年人多,也就是说随着人的成长,大脑的功能区越来越模块化,各种能力趋于成熟,很多神经元被“剪枝”了(一个笔者的猜想:是否意味着一个人如果保留尽可能多的神经元就会更聪明?)。

回到神经网络,神经元之间的连接就是权值,也就是神经网络的参数,如果用一些手段,可以在不掉精度的前提下,尽可能的减少神经元之间的连接,那么神经网络的参数和运算量就会相应的减小。以传统的ResNet50为例,现在的技术手段分分种就可以将运算量减少到剪枝前的五分之一甚至更多。可想而知,这种巨大的收益是多么的诱人,相信不少算法工程师都会用剪枝的技术显著提高自己的模型推理性能。

但遗憾的是,当运算量减小到五分之一的网络真正部署到加速芯片(CPU除外)的时候,结果多少有些让人沮丧,因为性能的提升远没有5倍之多,能有三倍的提升就已经做得不错了。原因也很简单,模型剪枝对硬件设计来说,并不友好。简单的来说,硬件加速芯片为了加速神经网络的计算,往往采用大并行度的计算阵列,而问题就出现在“并行度”三个字上,有的时候即使网络变得稀疏了,但由于并行度的存在,依然会有些已经被剪掉的连接会参与运算,即使它们的输入输出都是零。

为了缓解这个矛盾,软硬件工程师终于坐到一起商量出了一个折中的方案——结构化剪枝。简单说,就是根据硬件的并行度,有意的操作剪枝的形状,使得剪枝后网络的形状对硬件更友好。虽然牺牲了部分压缩比例,但是却可以有效降低网络的计算时间。因此结构化剪枝,是目前更为广泛的压缩方式。

数据压缩

经过量化和剪枝后的网络参数依然是有一个比较大的数据量,是不是可以进一步压缩一下?毕竟都不富裕,最好能省就省。相信所有同学应该都听说过数据压缩,毕竟网络上的各种图片、视频、文件,基本都是压缩过的。毫无疑问,所有的数据压缩方式(只考虑无损压缩)都可以用以神经网络的数据压缩,但由于神经网络的数据还是有一些自己的特点,所以还有一些特殊的压缩算法可以采用。接下来我们举几个典型的例子,来看看这些压缩算法到底适不适合用在我们的AI芯片里。

先说说变长编码中的典型——霍夫曼编码,其核心思想是使用较少的比特数表示出现频繁的数据而使用较多的比特数表示出现较少的数据。从硬件角度来看,这种压缩方式存在的问题是只有解析出前一个数据的值,才能确定下一个数据在哪,如果芯片想从存储器中读取特定的数据,实现起来非常麻烦。另外,如果神经网络参数的概率分布不集中,那压缩比率也不高。

再说说游程编码,其核心思想是用一个符号代替具有相同值的连续数据。这种编码方式很适合神经网络,这是由于神经网络的参数受一些非线性操作的影响会有很多零值,经过剪枝的网络更是如此。在如此稀疏的参数中,连续零的个数非常多,如果采用游程编码,压缩比率很高,而且压缩和解压相对比较简单,硬件实现起来很友好。但是和变长编码类似,想读取特定位置的数据比较困难。

最后说说一些专门用于稀疏矩阵压缩的算法。以Compressed Row Storage(CRS)算法或者称为 Compressed Sparse Row(CSR)算法为例,由名称可见,该格式是保留了稀疏矩阵中每一行的非零元素,以及非零元素的位置信息。虽然多了位置信息,但如果稀疏度达到一定程度,压缩还是有正收益的。这一类算法的压缩解压都比较简单,压缩比率不如游程编码,但由于配合一些特殊的运算结构,可以实现不解压直接运算的神奇疗效。所以这类算法在AI芯片中相对广泛的得到了应用。

但是,站在软件的角度,尤其是编译工程师眼里,压缩对他们产生指令来说,无疑是个灾难,毕竟在压缩后的一堆数里访问特定数据…好难。因此,对于具有多层存储结构的加速芯片而言,压缩后的数据在从最外层存储向芯片内搬移的时候就需要实现解压缩,这样芯片内部的寻址和计算,就都只需要针对非压缩的数据进行就可以了。虽然数据压缩的好处不能惠及片内存储,但多多少少也减少了些芯片和外部存储芯片之间的带宽需求和功耗。

对网络压缩未来的思考

剪枝还有必要么?

我们回想一下刚才提到的人脑神经元的发育过程,以及笔者的一个猜想:是否意味着一个人如果保留尽可能多的神经元就会更聪明呢?事实上,对于人工神经网络,笔者有相似的猜想:

为了实现同样的精度,是否存在一种网络结构,小而致密,没有冗余连接,无法剪枝,无法压缩? 对于一个剪枝后的网络,如果我们再增加一些神经元的连接,是否会得到精度更高的网络? 针对猜想一,我们可以看到一些趋势,新结构的网络往往比老结构的网络更难压缩,这是由于网络结构更合理,冗余连接更少造成的。针对猜想二,学术界已经有不少研究成果证实了这个结论。那么,我们是不是可以大胆的预测,未来的网络或许无需考虑模型剪枝。

量化到多少比特才是尽头?

如果一个网络随随便便从FP32量化到INT4都没有精度的损失,那只能说明这个网络结构太冗余了,与其用各种小技巧提高量化精度,不如改进一下网络结构;如果一个网络设计之初就采用低比特的架构,而不是通过量化+Fine-Tune来转化一个网络,那么相信前者这个网络的性能大概率要优于后者;从硬件设计角度来看,定点优于浮点,比特数越低越好。

因此,我猜想,在有限的未来一段时间,INT4或许会成为一个比较合适的数据结构,毕竟NVIDIA的支持也会助推软件工程师更关注INT4的网络。