據機器之心了解，騰訊 AI Lab 機器學習中心今日宣布成功研發自動化深度學習模型壓縮框架——PocketFlow，并即將在近期發布開源代碼。據介紹，這是一款面向移動端 AI 開發者的自動模型壓縮框架，集成了當前主流（包括騰訊 AI Lab 自研）的模型壓縮與訓練算法，結合自研超參數優化組件實現了全程自動化托管式的模型壓縮與加速。開發者無需了解具體算法細節，即可快速地將 AI 技術部署到移動端產品上，實現用戶數據的本地高效處理。目前該框架正在為騰訊的多項移動端業務提供模型壓縮與加速的技術支持，在多款手機 APP 中得到應用。值得一提的是，不久之前 MIT 韓松等人發表論文，提出了用于移動端模型自動壓縮與加速的AutoML方法AMC。
隨著 AI 技術的飛速發展，越來越多的公司希望在自己的移動端產品中注入 AI 能力，以優化用戶使用體驗。以深度學習為代表的 AI 技術極大地提升了圖像理解、語音識別等諸多應用領域中的識別精度，但是主流的深度學習模型往往對計算資源要求較高，難以直接部署到消費級移動設備中。常用的解決方案是將復雜的深度學習模型部署在云端，移動端將待識別的數據上傳至云端，再等待云端返回識別結果，但這對網絡傳輸速度的要求較高，在網絡覆蓋不佳地區的用戶使用體驗較差，同時數據上傳至云端后的隱私性也難以保證。
在這種情況下，眾多模型壓縮與加速算法應運而生，能夠在較小的精度損失（甚至無損）下，有效提升 CNN 和 RNN 等網絡結構的計算效率，從而使得深度學習模型在移動端的部署成為可能。但是，如何根據實際應用場景，選擇合適的模型壓縮與加速算法以及相應的超參數取值，往往需要較多的專業知識和實踐經驗，這無疑提高了這項技術對于一般開發者的使用門檻。
在此背景下，騰訊 AI Lab 機器學習中心研發了 PocketFlow 開源框架，以實現自動化的深度學習模型壓縮與加速，助力 AI 技術在更多移動端產品中的廣泛應用。通過集成多種深度學習模型壓縮算法，并創新性地引入超參數優化組件，極大地提升了模型壓縮技術的自動化程度。開發者無需介入具體的模型壓縮算法及其超參數取值的選取，僅需指定設定期望的性能指標，即可通過 PocketFlow 得到符合需求的壓縮模型，并快速部署到移動端應用中。
AI Lab 在 PocketFlow 上的研究進展
據了解，近期騰訊 AI Lab 機器學習中心在深度學習模型壓縮以及超參數優化算法方面持續投入，并取得了諸多研究進展。在模型壓縮算法方面，團隊提出了一種基于判別力最大化準則的通道剪枝算法，在性能基本無損的前提下可以大幅度降低 CNN 網絡模型的計算復雜度，相關論文已被 NIPS 2018 接收 [1]。該算法在訓練過程中引入多個額外的損失項，以提升 CNN 網絡中各層的判別力，然后逐層地基于分類誤差與重構誤差最小化的優化目標進行通道剪枝，去除判別力相對較小的冗余通道，從而實現模型的無損壓縮。在超參數優化算法方面，團隊研發了 AutoML 自動超參數優化框架，集成了包括高斯過程（Gaussian Processes, GP）和樹形結構 Parzen 估計器（Tree-structured Parzen Estimator, TPE）等在內的多種超參數優化算法，通過全程自動化托管解決了人工調參耗時耗力的問題，大幅度提升了算法人員的開發效率。
另一方面，考慮到深度學習模型的訓練周期普遍較長，團隊對基于 TensorFlow 的多機多卡訓練過程進行優化，降低分布式優化過程中的梯度通信耗時，研發了名為 TF-Plus 的分布式優化框架，僅需十幾行的代碼修改即可將針對單個 GPU 的訓練代碼擴展為多機多卡版本，并取得接近線性的加速比。此外，團隊還提出了一種誤差補償的量化隨機梯度下降算法，通過引入量化誤差的補償機制加快模型訓練的收斂速度，能夠在沒有性能損失的前提下實現一到兩個數量級的梯度壓縮，降低分布式優化中的梯度通信量，從而加快訓練速度，相關論文發表于 ICML 2018 [2]。
在 PocketFlow 框架的開發過程中，團隊加入了對上述多個自研算法的支持，并且有效降低了模型壓縮的精度損失，提升了模型的訓練效率，并極大地提高了超參數調節方面的自動化程度。
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 機器翻譯模型，均可以在一個小時內訓練完畢。
f) 超參數優化（hyper-parameter optimization）組件：多數開發者對模型壓縮算法往往不甚了解，但超參數取值對最終結果往往有著巨大的影響，因此團隊引入了超參數優化組件，采用了包括強化學習等算法以及 AI Lab 自研的 AutoML 自動超參數優化框架來根據具體性能需求，確定最優超參數取值組合。例如，對于通道剪枝算法，超參數優化組件可以自動地根據原始模型中各層的冗余程度，對各層采用不同的剪枝比例，在保證滿足模型整體壓縮倍數的前提下，實現壓縮后模型識別精度的最大化。
PocketFlow 性能展示
通過引入超參數優化組件，不僅避免了高門檻、繁瑣的人工調參工作，同時也使得 PocketFlow 在各個壓縮算法上全面超過了人工調參的效果。以圖像分類任務為例，在 CIFAR-10 和 ImageNet 等數據集上，PocketFlow 對 ResNet 和 MobileNet 等多種 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 表示，隨著 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.