首款自动化深度学习模型压缩框架——PocketFlow

【导语】腾讯 AI Lab 机器学习中心今日宣布成功研发出世界上首款自动化深度学习模型压缩框架——PocketFlow，并即将在今年10-11月发布开源代码。这是一款面向移动端 AI 开发者的自动模型压缩框架，与其他现有的模型压缩框架相比，PocketFlow 优势在于通过集成多种深度学习模型压缩算法，并引入超参数优化组件，提升了模型压缩技术的自动化程度。开发者无需了解具体算法细节，即可快速地将AI技术部署到移动端产品上，实现用户数据的本地高效处理。接下来 AI科技大本营就为大家解读一下其背后的理论与技术。

PocketFlow 算法概述

在模型压缩算法方面，团队提出了一种基于判别力最大化准则的通道剪枝算法，在性能基本无损的前提下可以大幅度降低CNN网络模型的计算复杂度，相关论文发表于 NIPS 2018。该算法在训练过程中引入多个额外的损失项，以提升CNN网络中各层的判别力，然后逐层地基于分类误差与重构误差最小化的优化目标进行通道剪枝，去除判别力相对较小的冗余通道，从而实现模型的无损压缩。

在超参数优化算法方面，团队研发了 AutoML 自动超参数优化框架，集成了包括高斯过程（Gaussian Processes, GP）和树形结构 Parzen 估计器（Tree-structured Parzen Estimator, TPE）等在内的多种超参数优化算法，通过全程自动化托管解决了人工调参耗时耗力的问题，大幅度提升了算法人员的开发效率。

另一方面，考虑到深度学习模型的训练周期普遍较长，团队对基于 TensorFlow 的多机多卡训练过程进行优化，降低分布式优化过程中的梯度通信耗时，研发了名为 TF-Plus 的分布式优化框架，仅需十几行的代码修改即可将针对单个GPU的训练代码扩展为多机多卡版本，并取得接近线性的加速比。此外，团队还提出了一种误差补偿的量化随机梯度下降算法，通过引入量化误差的补偿机制加快模型训练的收敛速度，能够在没有性能损失的前提下实现一到两个数量级的梯度压缩，降低分布式优化中的梯度通信量，从而加快训练速度，相关论文发表于 ICML 2018 。

PocketFlow 框架介绍

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

      

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

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

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

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

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

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

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

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

      

PocketFlow 性能

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

在 CIFAR-10 数据集上，PocketFlow 以 ResNet-56 作为基准模型进行通道剪枝，并加入了超参数优化和网络蒸馏等训练策略，实现了 2.5 倍加速下分类精度损失 0.4%，3.3 倍加速下精度损失 0.7%，且显著优于未压缩的 ResNet-44 模型； 在 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 助力移动端业务落地。在腾讯公司内部，PocketFlow 框架正在为多项移动端实际业务提供了模型压缩与加速的技术支持。例如，在手机拍照 APP 中，人脸关键点定位模型是一个常用的预处理模块，通过对脸部的百余个特征点（如眼角、鼻尖等）进行识别与定位，可以为后续的人脸识别、智能美颜等多个应用提供必要的特征数据。团队基于 PocketFlow 框架，对人脸关键点定位模型进行压缩，在保持定位精度不变的同时，大幅度地降低了计算开销，在不同的移动处理器上取得了 25%-50% 不等的加速效果，压缩后的模型已经在实际产品中得到部署。

      

对话腾讯 AI Lab

在发布的首日，我们和腾讯 AI Lab 就 PocketFlow 相关问题进行了更深入的交流。以下内容为采访实录。

▌目前有哪些主流的模型压缩算法？你们采用的基于判别力最大化准则的通道剪枝算法的原理是什么？有哪些优势？

腾讯 AI Lab：目前主流的模型压缩算法包括低秩分解（low-rank decomposition）、通道剪枝（channel pruning）、权重稀疏化（weight sparsification）、权重量化（weight quantization）、网络结构搜索（network architecture search）等。PocketFlow 中不仅包含了 AI Lab 自研的模型压缩算法（例如基于判别力最大化准则的通道剪枝算法），还提供了对当前主流的多种模型压缩算法（包括通道剪枝、权重稀疏化、权重量化）的支持，并且基于自研的 AutoML 框架提供了分布式网络结构搜索功能，功能完备且易于扩展。

我们提出基于判别力最大化准则的通道剪枝算法的一个出发点是目前的通道剪枝算法大多没有考虑通道的判别力信息，仅考虑了重构误差的最小化，导致高压缩倍数下的精度损失较大。我们算法首先在训练过程中引入了多个额外的有监督损失项，以提升CNN网络中各层的判别力，再逐层地基于分类误差与重构误差最小化的优化目标进行通道剪枝，去除判别力相对较小的冗余通道，从而实现模型的无损压缩。

▌PocketFlow框架支持通道剪枝、权重稀疏化、权重量化算法，每次模型压缩的过程是随机选择一种算法吗？这些算法各有哪些优势？

腾讯 AI Lab：PocketFlow 会基于具体的模型压缩需求，选择适当的模型压缩算法。这三类算法各有优势：通道剪枝算法的加速性能较好，易于部署；权重稀疏化算法可以达到更高的压缩倍数；权重量化算法的内存带宽开销较小，并且可以充分利用 ARM、FPGA 等设备在定点运算方面的硬件优化。举个例子，如果用户希望得到加速倍数较高的压缩模型，并且直接在移动端部署，那么PocketFlow 会自动地选择通道剪枝算法进行模型压缩。

▌对网络权重引入稀疏化约束会增加模型的训练难度，降低模型的收敛性吗？你们是如何解决的？

腾讯 AI Lab：是的，在加入权重稀疏化约束的过程中，会引入数个超参数，如果取值不当，会导致模型训练不稳定甚至完全不收敛。PocketFlow 中超参数优化组件的引入正是为了解决这一问题（同时也适用于其他模型压缩算法），通过自动化地搜索最优超参数取值组合，仅经过 10余 次迭代就可以达到与人工调参类似的效果，在 100 次迭代后可以使得压缩后的模型降低约 0.6% 的精度损失。

▌在超参数优化方面，你们提出的 AutoML-自动超参数优化框架和 Google 的 AutoML 有什么异同？

腾讯 AI Lab：AutoML 自动超参数优化框架是我们 AI Lab 之前已有的一套自研工具，适用于任意模型的超参数优化任务，PocketFlow 借助于该框架实现了模型压缩与加速场景下的超参数优化，以提高训练流程的自动化程度，降低使用门槛。Google 的 AutoML 则更多地针对视觉领域的应用场景，实现了从网络结构搜索到迁移学习和超参数优化的整体流程的自动化。

▌为什么 PocketFlow 可以在压缩模型的基础上甚至让模型的性能提升？

腾讯 AI Lab：目前深度学习模型大多存在一定的冗余度，通过在一定范围内降低模型的冗余度，往往不会对模型的拟合能力造成影响；另一方面，模型压缩算法通过限制模型的解空间，可以起到防止模型过拟合的作用，尤其是在一些训练数据和标注信息有限的应用场景下，反而能带来模型的性能提升。

▌压缩过程是在云端还是离线完成的？如何保护使用者技术隐私和数据隐私？

腾讯 AI Lab：模型压缩的训练过程没有对云端服务的依赖，可以直接基于我们即将发布的开源代码在使用者的本地环境中运行，无需担心技术和数据的隐私性问题。

▌模型压缩后如何部署到移动端（各种不同的移动设备）？对硬件的需求如何？

腾讯 AI Lab：基于不同模型压缩算法训练得到的模型的部署方式不尽相同，例如通道剪枝后的模型可以直接基于目前常见的移动端深度学习框架进行部署，权重定点量化后的模型可以基于 TensorFlow Lite进行部署。经过 PocketFlow 压缩或者加速后，无论是模型大小还是计算量都大大精简，因此在移动端部署模型时对硬件的需求并不高，例如我们正在为一款手机拍照 APP 提供模型压缩支持，压缩后的模型在华为 Mate 10 和小米 5S Plus 等设备上的运行耗时均在10ms 以内。

▌你们的实验都是用的 CV 里ResNet 和 MobileNet 这种比较成熟的模型，以及被广泛使用的公开数据集CIFAR-10和ImageNet，推广到其他任意的模型上的效果是否也可以保证？泛化效果如何？

腾讯 AI Lab：我们在腾讯公司内部也支持了诸如人脸关键点定位和姿态估计等任务中的模型压缩需求。以人脸关键点定位任务为例，我们可以在保持定位精度不变的同时，显著降低计算开销，根据在不同手机上的实测数据，压缩后的模型可以取得 1.4-2.0 倍不等的加速效果。

▌PocketFlow目前有没有对语音识别、NLP等领域的任务模型进行压缩前后的对照试验？效果如何？

腾讯 AI Lab：我们当前主要针对视觉领域中的模型进行压缩，下一阶段的主要研发目标之一就是为语音和 NLP 等模型提供压缩技术支持，届时也欢迎有相关需求的开发者和研究人员来使用并提出改进意见。

▌用户需要设定哪些期望性能指标？是预先设定好迭代次数？迭代次数如何设定是合适的？如果一直达不到设定的指标怎么办？

腾讯 AI Lab：用户主要需要设置的期望性能指标是目标的模型压缩和/或加速倍数，之后 PocketFlow 会自动搜索符合要求的模型中精度最高者作为输出；我们会给出建议的压缩/加速倍数的设置范围（例如基于权重量化的压缩倍数不能超过 32 倍），因此基本不会出现达不到性能指标的问题。迭代次数不需要用户设置（对于进阶用户我们提供了设置迭代次数的接口），根据实验评估结果，我们预设的迭代次数适用于多个模型压缩算法的超参数优化任务。

▌开发者什么时候可以真正使用上？

腾讯 AI Lab：我们预计在今年10-11月完成 PocketFlow 的开源发布工作，届时普通开发者即可开始使用，同时我们也欢迎有兴趣的开发者向 PocketFlow 贡献代码。