PCONV是目前最快的通用性移动端神经网络加速框架,由美国东北大学(Northeastern University),威廉玛丽学院(William & Mary),和清华大学交叉信息研究院、交叉信息核心技术研究院共同提出。
该框架创新性地提出了神经网络稀疏性结构与编译器的协同优化,提出了新的适合编译器与移动端加速的稀疏化结构:模式化剪枝与连通性剪枝,并且使用交替乘子算法(ADMM)达到稀疏结构。
通过自己编写的编译器自动代码生成,在Samsung手机上达到了19ms的VGG-16推理速度与5.4ms的MobileNet-V2推理速度(均为ImageNet数据集,无精度损失),首次在通用移动端达到实时性要求。该结果显著超过现有的移动端通用加速框架。我们的研究使得在移动端对任意神经网络进行实时运算变为可能。
随着深度神经网络近年来的发展,神经网络模型随之增大。巨大的模型尺寸带来了更复杂的计算需求,在应用端也占用了更多的内存空间,使得深度神经网络在移动平台上的实时推理能力受到很大的挑战。即便是当今的移动平台已经拥有非常强大的计算与存储能力,借助现有的移动端软件深度神经网络加速框架(例如TVM,MNN和TensorFlow-Lite)依然很难实现实时推理。
为了解决上述问题,深度神经网络模型剪枝技术成为模型压缩中一个直接有效的方法。通过模型剪枝,神经网络的权重结构变得稀疏化。早期的非结构化剪枝采用了一种迭代型探索方法,将模型中绝对值较小的权重删除,得到一种权重不规则分布的神经网络模型。这种方法虽然取得了一定的压缩率与精度,但是却并没有改善模型的硬件执行效率,对硬件平行计算结构与内存访问并没有任何优化,而且也不适用于并行度较高的硬件平台。结构化剪枝的提出在一定程度上解决了上述问题,但是其过于激进的剪枝方式(剪枝整个filter或channel)导致模型的精度和泛化能力的大大下降,从而影响了结构化剪枝的广泛使用。
我们通过总结发现上述两种模型剪枝方法实质上代表了两个极端,限制了神经网络模型剪枝在软件层面和硬件层面的表现。因此,我们提出了一种新型的剪枝稀疏性结构PCONV,包含了卷积核(convolution kernel)模式化剪枝(pattern pruning)与连通性剪枝(connectivity pruning),如图1所示。卷积核模式化剪枝作为一种获取卷积核内部稀疏性的方法,其剪枝维度定位在每一个卷积核内部,通过删除卷积核内部固定数量的权重,使得卷积核呈现不同的形状,我们称这些不同的形状为卷积核模式(kernel pattern)。
值得注意的是,特定的卷积核模式能够融合计算机视觉概念中图像滤波器的特性,使得深度神经网络对图像的特征提取能力增强,从而得到更高的推理精度。另一方面,连通性剪枝作为一种获取卷积核之间稀疏性的方法,其剪枝维度定位在卷积核层面,通过删除深度神经网络中每一个卷积层输入与输出之间较弱的联通关系,我们进一步加大了模型的压缩率,从而提高了模型加速的可行性空间。
通过联合模式化剪枝与连通性剪枝,我们提出的新型模型压缩兼顾了非结构化剪枝与结构化剪枝的优势,将高精度,高压缩率,权重的结构化空间分布完美融合在同一个深度神经网络模型中。为了利用这种新型的稀疏性结构,我们同时设计了一种面向移动平台的编译器,能够部署并高效执行PCONV模型。
这种编译器构架基于编译器概念中非常重要的代码生成原理,将深度神经网络模型转化为底层静态执行代码,并配合适合模型特点的编译优化。由于PCONV模型的特点与编译器的结合,我们实现了在移动平台上指令级和线程级的高并行性,以及对于大规模神经网络在通用性移动设备首次实现了实时推理。
图1. 模式化剪枝与连通性剪枝示意图。
如何设计卷积核模式,如何分配这些模式给神经网络中每一个卷积核,同时还能判别神经网络中较弱的联通关系并实现连通性剪枝是算法层面的关键问题。首先,对于卷积核模式的设计,我们采用了一种独特的角度,将剪枝转化成对神经网络模型权重加入一层二进制掩膜(binary masking)。
通过将这些二进制掩膜与原权重进行矩阵点乘,得到稀疏性。我们将加入掩膜看作一种对于神经网络的插值操作(interpolation),通过不同模式的掩膜插值,得到一些功能性图像滤波器特征,能够实现图像的锐化与降噪,提高图像质量。
图2展示了我们设计的四种卷积核模式,通过n次插值,得到增强型拉普拉斯高斯滤波器。其中n代表神经根网络层数,p为0.75,在正则化运算后没有实际意义。其次,我们设计了基于ADMM(交替乘子优化算法)算法的深度神经网络剪枝框架,通过将剪枝问题转向为数学优化问题,同时解决了卷积核模式分配与连通性剪枝的问题,并在剪枝的同时训练剩余权重。
ADMM将原始剪枝问题分解为两个子问题,用传统梯度下降法求解子问题一,并引入一个二次项迭代求解子问题二。通过实验我们可以发现,该方法在不同量级的神经网络上,比如VGG-16, ResNet-50, MobileNet-V2等代表性网络模型在ImageNet与CIFAR-10数据集上,均取得更好的训练效果(准确率)。
图2. 卷积核模式设计。
上述算法实现为我们提供了PCONV模型,如何利用模型的特点实现编译器优化成为编译器与硬件层面的研究重点。我们设计了适合PCONV的移动端推理框架,如图3所示。首先我们分析并提取PCONV模型中的稀疏性信息。随后,卷积核模式与连通性信息变为已知,对应于每个卷积核的计算方式也就变成已知的。
图3. PCONV移动端基于编译器的推理框架设计与优化示意图。
我们通过将相同的卷积核模式排列在一起,同时将拥有相似模式的输出通道(filter)排列在一起形成一个计算组(group),如图4所示。这种新型排列后的卷积核生成的静态代码能够消除所有代码分支,意味着高指令级平行性,同时,相似的输出通道保证了相当高的负载均衡,从而得到了高线程级平行性。
图4. 卷积核与输出通道重排示意图
卷积核与输出通道重排对于内存方面的优化效果也非常明显,重排过后的PCONV模型权重排布非常规则,意味着数据访问的频率降低,硬件上不必在每次做卷积计算时都访问一次内存,而是在卷积核模式变化时才访问一次。数据访问频率降低意味着更低的内存开销。
图5. 模型模式化剪枝后在VGG-16上的梯度图像可视化效果图(左),模型压缩精度与压缩率在ImageNet上与ADMM-NN (Ren et al. 2019), NeST (Dai, Yin, and Jha 2017), Deep Compression (Han, Mao, and Dally 2015), Fine-grained pruning (Mao et al. 2017), APoZ (Hu et al. 2016) and ThiNet (Luo, Wu, and Lin 2017) 对比图。(右)
实验结果表明,卷积核模式化剪枝对于图像有明显的增强作用,如图5(左)所示。这解释了在图5(右)中,深度神经网络在模式化剪枝后的精度增加的表现。模式化剪枝后的模型能够提取图像中更多的关键特征,并降低图像中的噪声。总体来讲,相比于其他方法,PCONV能够得到更高的压缩率,并得到更高的模型精度。
图6. PCONV模型结合编译器在移动端的推理速度对比图(无精度损失,16位浮点数表示),可以看到PCONV速度远超现有的加速框架,实现了在通用性移动端的实时推理。
PCONV不仅能得到更高的压缩率和模型精度,更重要的是与编译器的协同优化极大地提高了移动端的推理速度。在移动端,我们使用了Samsung Galaxy S10智能手机来测试PCONV的推理速度。图6展示了PCONV与编译器结合后在移动端的推理速度与在现有的一些深度神经网络加速器(TVM,MNN,TensorFlow-Lite)上的速度对比。
我们可以看到,PCONV在移动端的推理速度显著超过现有的加速器。事实上,在每一种网络结构下,PCONV在移动端都能在没有精度损失的情况下实现实时计算的要求(30 frames/second,i.e.,33ms/second)。比如对于大型神经网络VGG-16,我们在ImageNet数据集上达到了19ms的推理时间。对于轻量级神经网络MobileNet-V2,我们更是达到了5.4ms的推理时间。我们的研究使得在移动端对任意神经网络进行实时运算变为可能。
论文已被AAAI 2020接收,题目《PCONV: The Missing but Desirable Sparsity in DNN Weight Pruning for Real-time Execution on Mobile Devices》,主要撰稿:马晓龙。
论文下载地址: