X

滴滴、东北大学提自动结构化剪枝压缩算法框架,性能提升高达120倍

机器之心发布

机器之心编辑部

2020 年 2 月 7 日-2 月 12 日,AAAI 2020 将于美国纽约举办。不久之前,大会官方公布了今年的论文收录信息:收到 8800 篇提交论文,评审了 7737 篇,接收 1591 篇,接收率 20.6%。本文介绍了滴滴 AI Labs 与美国东北大学合作的一篇论文《AutoCompress: An Automatic DNN Structured Pruning Framework for Ultra-High Compression Rates》

近年来,随着深度神经网络模型性能不断刷新,模型的骨干网络参数量愈发庞大,存储和计算代价不断提高,从而导致难以部署在资源受限的嵌入式平台上。滴滴 AI Labs 与美国东北大学王言治教授研究组合作,联合提出了一种基于 AutoML 思想的自动结构化剪枝的算法框架 AutoCompress,能自动化的去寻找深度模型剪枝中的超参数,去除模型中不同层的参数冗余,替代人工设计的过程并实现了超高的压缩倍率。从而满足嵌入式端上运行深度模型的实时性能需求。
相较之前方法的局限性,该方法提出三点创新性设计:

(1)提出混合型的结构化剪枝维度;(2)采用高效强大的神经网络剪枝算法 ADMM(交替乘子优化算法)对训练过程中的正则项进行动态更新;(3)利用了增强型引导启发式搜索的方式进行行为抽样。在 CIFAR 和 ImageNet 数据集的大量测试表明 AutoCompress 的效果显著超过各种神经网络压缩方法与框架。在相同准确率下,实际参数量的压缩相对之前方法最大可以提高超 120 倍。

图 1. 自动化超参数决策框架的通用流程,以及性能提升来源
为了改进以上的不足,我们提出了神经网络权重剪枝问题超参数设置自动化过程的通用流程(generic flow),如图 1 所示。整个自动化通用流程主要可以分为 4 步。步骤 1 为行为抽样,步骤 2 为快速评估,步骤 3 为确定决策,步骤 4 为剪枝操作。
由于超参数的巨大搜索空间,步骤 1 和步骤 2 应该快速进行,因此无法进行再训练(re-training)后去评估其效果。因此根据量级最小的一部分权重直接进行剪枝评估。步骤 3 根据工作抽样和评估的集合对超参数进行决策。步骤 4 利用剪枝核心算法对模型进行结构化剪枝并生成结果。
基于上述通用流程,并针对之前方法的局限性,进一步提出如下三点创新性设计,通过综合现有的深度神经网络与机器学习框架首次实现了目前最高效的深度神经网络自动化结构化剪枝的通用框架 AutoCompress。该框架在滴滴已经得到了实际有效应用。
基于神经网络自动化结构化剪枝框架
三点创新性设计为:(1)提出混合型的结构化剪枝维度;(2)采用高效强大的神经网络剪枝算法 ADMM(交替乘子优化算法)对训练过程中的正则项进行动态更新;(3)利用了增强型引导启发式搜索的方式进行行为抽样。

图 3. AutoCompress 框架示意图
基于上述三点创新性设计,我们搭建了 AutoCompress 框架,如图 3 所示。通过基于启发式搜索算法的自动化代理模块的指导,AutoCompress 框架进行模型自动化剪枝主要分为两个步骤,步骤 1:通过基于 ADMM 算法的结构化剪枝,得到权重分布结构性稀疏化的模型;步骤 2:通过网络结构净化(Purification)操作,将 ADMM 过程中无法完全删除的一小部分冗余权重找到并删除。值得注意的是,这两个步骤都是基于相同的启发式搜索机制。
启发式搜索机制
针对 AutoCompress 中最核心的搜索算法设计,我们利用了搜索算法中的模拟退火算法为搜索算法的核心。举例来讲,给定一个原始模型,我们会设置两种目标函数 — 根据权重数量设置或根据运算量(FLOPs)设置。搜索过程进行若干轮,比如第一轮目标为压缩两倍权重数量,第二轮为压缩四倍权重数量。在每一轮搜索过程中,首先初始化一个行为(超参数),然后每次对行为进行一个扰动(超参数的小幅变化)生成新的行为,根据模拟退火算法原理,评估两个行为,如果新的行为评估结果优于原结果则接受该行为,如果新的行为评估结果劣于原结果则以一定概率接受该行为。每一轮算法中的温度参数 T 会下降,直到 T 下降到某个阈值后即停止搜索,该结果即为图 1 中的步骤 3 输出。最后,根据搜索得到的超参数,对神经网络进行结构化剪枝操作。

表 2. 在 ResNet-18 (NISP 和 AMC 结果为 ResNet-50) 上基于 CIFAR-10 数据集的权重剪枝对比结果。

表 4. 在 ResNet-18/50 上基于 ImageNet 数据集的结构化权重剪枝对比结果。