腾讯 AI Lab 开源世界首款自动化模型压缩框架PocketFlow

雷锋网AI科技评论按：腾讯 AI Lab 机器学习中心今日宣布成功研发出世界上首款自动化深度学习模型压缩框架—— PocketFlow，并即将在近期发布开源代码。根据雷锋网AI科技评论了解，这是一款面向移动端AI开发者的自动模型压缩框架，集成了当前主流的模型压缩与训练算法，结合自研超参数优化组件实现了全程自动化托管式的模型压缩与加速。开发者无需了解具体算法细节，即可快速地将AI技术部署到移动端产品上，实现用户数据的本地高效处理。

随着AI技术的飞速发展，越来越多的公司希望在自己的移动端产品中注入AI能力，但是主流的深度学习模型往往对计算资源要求较高，难以直接部署到消费级移动设备中。在这种情况下，众多模型压缩与加速算法应运而生，能够在较小的精度损失（甚至无损）下，有效提升 CNN 和 RNN 等网络结构的计算效率，从而使得深度学习模型在移动端的部署成为可能。但是，如何根据实际应用场景，选择合适的模型压缩与加速算法以及相应的超参数取值，往往需要较多的专业知识和实践经验，这无疑提高了这项技术对于一般开发者的使用门槛。

在此背景下，腾讯AI Lab机器学习中心研发了 PocketFlow 开源框架，以实现自动化的深度学习模型压缩与加速，助力AI技术在更多移动端产品中的广泛应用。通过集成多种深度学习模型压缩算法，并创新性地引入超参数优化组件，极大地提升了模型压缩技术的自动化程度。开发者无需介入具体的模型压缩算法及其超参数取值的选取，仅需指定设定期望的性能指标，即可通过 PocketFlow 得到符合需求的压缩模型，并快速部署到移动端应用中。

框架介绍

PocketFlow 框架主要由两部分组件构成，分别是模型压缩/加速算法组件和超参数优化组件，具体结构如下图所示。

开发者将未压缩的原始模型作为 PocketFlow 框架的输入，同时指定期望的性能指标，例如模型的压缩和/或加速倍数；在每一轮迭代过程中，超参数优化组件选取一组超参数取值组合，之后模型压缩/加速算法组件基于该超参数取值组合，对原始模型进行压缩，得到一个压缩后的候选模型；基于对候选模型进行性能评估的结果，超参数优化组件调整自身的模型参数，并选取一组新的超参数取值组合，以开始下一轮迭代过程；当迭代终止时，PocketFlow 选取最优的超参数取值组合以及对应的候选模型，作为最终输出，返回给开发者用作移动端的模型部署。

具体地，PocketFlow 通过下列各个算法组件的有效结合，实现了精度损失更小、自动化程度更高的深度学习模型的压缩与加速：

a) 通道剪枝（channel pruning）组件：在CNN网络中，通过对特征图中的通道维度进行剪枝，可以同时降低模型大小和计算复杂度，并且压缩后的模型可以直接基于现有的深度学习框架进行部署。在CIFAR-10图像分类任务中，通过对 ResNet-56 模型进行通道剪枝，可以实现2.5倍加速下分类精度损失0.4%，3.3倍加速下精度损失0.7%。

b) 权重稀疏化（weight sparsification）组件：通过对网络权重引入稀疏性约束，可以大幅度降低网络权重中的非零元素个数；压缩后模型的网络权重可以以稀疏矩阵的形式进行存储和传输，从而实现模型压缩。对于 MobileNet 图像分类模型，在删去50%网络权重后，在 ImageNet 数据集上的 Top-1 分类精度损失仅为0.6%。

c) 权重量化（weight quantization）组件：通过对网络权重引入量化约束，可以降低用于表示每个网络权重所需的比特数；团队同时提供了对于均匀和非均匀两大类量化算法的支持，可以充分利用 ARM 和 FPGA 等设备的硬件优化，以提升移动端的计算效率，并为未来的神经网络芯片设计提供软件支持。以用于 ImageNet 图像分类任务的 ResNet-18 模型为例，在8比特定点量化下可以实现精度无损的4倍压缩。

d) 网络蒸馏（network distillation）组件：对于上述各种模型压缩组件，通过将未压缩的原始模型的输出作为额外的监督信息，指导压缩后模型的训练，在压缩/加速倍数不变的前提下均可以获得0.5%-2.0%不等的精度提升。

e) 多GPU训练（multi-GPU training）组件：深度学习模型训练过程对计算资源要求较高，单个GPU难以在短时间内完成模型训练，因此团队提供了对于多机多卡分布式训练的全面支持，以加快使用者的开发流程。无论是基于 ImageNet 数据的Resnet-50图像分类模型还是基于 WMT14 数据的 Transformer 机器翻译模型，均可以在一个小时内训练完毕。[1]

f) 超参数优化（hyper-parameter optimization）组件：多数开发者对模型压缩算法往往不甚了解，但超参数取值对最终结果往往有着巨大的影响，因此团队引入了超参数优化组件，采用了包括强化学习等算法以及 AI Lab 自研的 AutoML 自动超参数优化框架来根据具体性能需求，确定最优超参数取值组合。例如，对于通道剪枝算法，超参数优化组件可以自动地根据原始模型中各层的冗余程度，对各层采用不同的剪枝比例，在保证满足模型整体压缩倍数的前提下，实现压缩后模型识别精度的最大化。

性能展示

通过引入超参数优化组件，不仅避免了高门槛、繁琐的人工调参工作，同时也使得 PocketFlow 在各个压缩算法上全面超过了人工调参的效果。以图像分类任务为例，在 CIFAR-10 和 ImageNet 等数据集上， PocketFlow 对 ResNet 和 MobileNet 等多种 CNN 网络结构进行有效的模型压缩与加速。[1]

在 CIFAR-10 数据集上，PocketFlow 以 ResNet-56 作为基准模型进行通道剪枝，并加入了超参数优化和网络蒸馏等训练策略，实现了2.5倍加速下分类精度损失0.4%，3.3倍加速下精度损失0.7%，且显著优于未压缩的ResNet-44模型； [2] 在 ImageNet 数据集上，PocketFlow 可以对原本已经十分精简的 MobileNet 模型继续进行权重稀疏化，以更小的模型尺寸取得相似的分类精度；与 Inception-V1 、ResNet-18 等模型相比，模型大小仅为后者的约20~40%，但分类精度基本一致（甚至更高）。

相比于费时费力的人工调参，PocketFlow 框架中的 AutoML 自动超参数优化组件仅需10余次迭代就能达到与人工调参类似的性能，在经过100次迭代后搜索得到的超参数组合可以降低约0.6%的精度损失；通过使用超参数优化组件自动地确定网络中各层权重的量化比特数，PocketFlow 在对用于 ImageNet 图像分类任务的ResNet-18模型进行压缩时，取得了一致性的性能提升；当平均量化比特数为4比特时，超参数优化组件的引入可以将分类精度从63.6%提升至68.1%（原始模型的分类精度为70.3%）。

深度学习模型的压缩与加速是当前学术界的研究热点之一，同时在工业界中也有着广泛的应用前景。随着PocketFlow的推出，开发者无需了解模型压缩算法的具体细节，也不用关心各个超参数的选择与调优，即可基于这套自动化框架，快速得到可用于移动端部署的精简模型，从而为AI能力在更多移动端产品中的应用铺平了道路。

参考文献

[1] Zhuangwei Zhuang, Mingkui Tan, Bohan Zhuang, Jing Liu, Jiezhang Cao, Qingyao Wu, Junzhou Huang, Jinhui Zhu, “Discrimination-aware Channel Pruning for Deep Neural Networks", In Proc. of the 32nd Annual Conference on Neural Information Processing Systems, NIPS '18, Montreal, Canada, December 2018.

[2] Jiaxiang Wu, Weidong Huang, Junzhou Huang, Tong Zhang, “Error Compensated Quantized SGD and its Applications to Large-scale Distributed Optimization”, In Proc. of the 35th International Conference on Machine Learning, ICML ’18, Stockholm, Sweden, July 2018.

雷锋网AI科技评论