據機器之心了解,，騰訊 AI Lab 機器學習中心今日宣布成功研發(fā)自動化深度學習模型壓縮框架——PocketFlow,，并即將在近期發(fā)布開源代碼。據介紹,，這是一款面向移動端 AI 開發(fā)者的自動模型壓縮框架,，集成了當前主流（包括騰訊 AI Lab 自研）的模型壓縮與訓練算法，結合自研超參數優(yōu)化組件實現了全程自動化托管式的模型壓縮與加速,。開發(fā)者無需了解具體算法細節(jié),，即可快速地將 AI 技術部署到移動端產品上，實現用戶數據的本地高效處理,。目前該框架正在為騰訊的多項移動端業(yè)務提供模型壓縮與加速的技術支持,，在多款手機 APP 中得到應用。值得一提的是,，不久之前 MIT 韓松等人發(fā)表論文,，提出了用于移動端模型自動壓縮與加速的AutoML方法AMC。
隨著 AI 技術的飛速發(fā)展,，越來越多的公司希望在自己的移動端產品中注入 AI 能力,，以優(yōu)化用戶使用體驗。以深度學習為代表的 AI 技術極大地提升了圖像理解,、語音識別等諸多應用領域中的識別精度,，但是主流的深度學習模型往往對計算資源要求較高，難以直接部署到消費級移動設備中,。常用的解決方案是將復雜的深度學習模型部署在云端,，移動端將待識別的數據上傳至云端，再等待云端返回識別結果,，但這對網絡傳輸速度的要求較高,，在網絡覆蓋不佳地區(qū)的用戶使用體驗較差，同時數據上傳至云端后的隱私性也難以保證,。
在這種情況下,，眾多模型壓縮與加速算法應運而生，能夠在較小的精度損失（甚至無損）下,，有效提升 CNN 和 RNN 等網絡結構的計算效率,，從而使得深度學習模型在移動端的部署成為可能。但是,，如何根據實際應用場景,，選擇合適的模型壓縮與加速算法以及相應的超參數取值，往往需要較多的專業(yè)知識和實踐經驗,，這無疑提高了這項技術對于一般開發(fā)者的使用門檻,。
在此背景下,，騰訊 AI Lab 機器學習中心研發(fā)了 PocketFlow 開源框架，以實現自動化的深度學習模型壓縮與加速,，助力 AI 技術在更多移動端產品中的廣泛應用,。通過集成多種深度學習模型壓縮算法，并創(chuàng)新性地引入超參數優(yōu)化組件,，極大地提升了模型壓縮技術的自動化程度,。開發(fā)者無需介入具體的模型壓縮算法及其超參數取值的選取，僅需指定設定期望的性能指標,，即可通過 PocketFlow 得到符合需求的壓縮模型,，并快速部署到移動端應用中。
AI Lab 在 PocketFlow 上的研究進展
據了解,，近期騰訊 AI Lab 機器學習中心在深度學習模型壓縮以及超參數優(yōu)化算法方面持續(xù)投入,，并取得了諸多研究進展。在模型壓縮算法方面,，團隊提出了一種基于判別力最大化準則的通道剪枝算法,，在性能基本無損的前提下可以大幅度降低 CNN 網絡模型的計算復雜度，相關論文已被 NIPS 2018 接收 [1],。該算法在訓練過程中引入多個額外的損失項,，以提升 CNN 網絡中各層的判別力，然后逐層地基于分類誤差與重構誤差最小化的優(yōu)化目標進行通道剪枝,，去除判別力相對較小的冗余通道,，從而實現模型的無損壓縮。在超參數優(yōu)化算法方面,，團隊研發(fā)了 AutoML 自動超參數優(yōu)化框架,，集成了包括高斯過程（Gaussian Processes, GP）和樹形結構 Parzen 估計器（Tree-structured Parzen Estimator, TPE）等在內的多種超參數優(yōu)化算法，通過全程自動化托管解決了人工調參耗時耗力的問題,，大幅度提升了算法人員的開發(fā)效率,。
另一方面，考慮到深度學習模型的訓練周期普遍較長,，團隊對基于 TensorFlow 的多機多卡訓練過程進行優(yōu)化,，降低分布式優(yōu)化過程中的梯度通信耗時，研發(fā)了名為 TF-Plus 的分布式優(yōu)化框架,，僅需十幾行的代碼修改即可將針對單個 GPU 的訓練代碼擴展為多機多卡版本，并取得接近線性的加速比,。此外,，團隊還提出了一種誤差補償的量化隨機梯度下降算法，通過引入量化誤差的補償機制加快模型訓練的收斂速度,，能夠在沒有性能損失的前提下實現一到兩個數量級的梯度壓縮,，降低分布式優(yōu)化中的梯度通信量,，從而加快訓練速度，相關論文發(fā)表于 ICML 2018 [2],。
在 PocketFlow 框架的開發(fā)過程中,，團隊加入了對上述多個自研算法的支持，并且有效降低了模型壓縮的精度損失,，提升了模型的訓練效率,，并極大地提高了超參數調節(jié)方面的自動化程度。
PocketFlow 框架介紹
PocketFlow 框架主要由兩部分組件構成,，分別是模型壓縮/加速算法組件和超參數優(yōu)化組件,，具體結構如下圖所示。
開發(fā)者將未壓縮的原始模型作為 PocketFlow 框架的輸入,，同時指定期望的性能指標,，例如模型的壓縮和/或加速倍數；在每一輪迭代過程中,，超參數優(yōu)化組件選取一組超參數取值組合,，之后模型壓縮/加速算法組件基于該超參數取值組合，對原始模型進行壓縮,，得到一個壓縮后的候選模型,；基于對候選模型進行性能評估的結果，超參數優(yōu)化組件調整自身的模型參數,，并選取一組新的超參數取值組合,，以開始下一輪迭代過程；當迭代終止時,，PocketFlow 選取最優(yōu)的超參數取值組合以及對應的候選模型,，作為最終輸出，返回給開發(fā)者用作移動端的模型部署,。
具體地,，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 等設備的硬件優(yōu)化,，以提升移動端的計算效率,，并為未來的神經網絡芯片設計提供軟件支持。以用于 ImageNet 圖像分類任務的 ResNet-18 模型為例,，在 8 比特定點量化下可以實現精度無損的 4 倍壓縮,。
d) 網絡蒸餾（network distillation）組件：對于上述各種模型壓縮組件，通過將未壓縮的原始模型的輸出作為額外的監(jiān)督信息,，指導壓縮后模型的訓練,，在壓縮/加速倍數不變的前提下均可以獲得 0.5%-2.0% 不等的精度提升。
e) 多 GPU 訓練（multi-GPU training）組件：深度學習模型訓練過程對計算資源要求較高,，單個 GPU 難以在短時間內完成模型訓練,，因此團隊提供了對于多機多卡分布式訓練的全面支持，以加快使用者的開發(fā)流程,。無論是基于 ImageNet 數據的 Resnet-50 圖像分類模型還是基于 WMT14 數據的 Transformer 機器翻譯模型,，均可以在一個小時內訓練完畢。
f) 超參數優(yōu)化（hyper-parameter optimization）組件：多數開發(fā)者對模型壓縮算法往往不甚了解,，但超參數取值對最終結果往往有著巨大的影響,，因此團隊引入了超參數優(yōu)化組件，采用了包括強化學習等算法以及 AI Lab 自研的 AutoML 自動超參數優(yōu)化框架來根據具體性能需求,，確定最優(yōu)超參數取值組合,。例如，對于通道剪枝算法,，超參數優(yōu)化組件可以自動地根據原始模型中各層的冗余程度,，對各層采用不同的剪枝比例，在保證滿足模型整體壓縮倍數的前提下,，實現壓縮后模型識別精度的最大化,。
PocketFlow 性能展示
通過引入超參數優(yōu)化組件，不僅避免了高門檻,、繁瑣的人工調參工作,，同時也使得 PocketFlow 在各個壓縮算法上全面超過了人工調參的效果。以圖像分類任務為例，在 CIFAR-10 和 ImageNet 等數據集上,，PocketFlow 對 ResNet 和 MobileNet 等多種 CNN 網絡結構進行有效的模型壓縮與加速。
在 CIFAR-10 數據集上,，PocketFlow 以 ResNet-56 作為基準模型進行通道剪枝,，并加入了超參數優(yōu)化和網絡蒸餾等訓練策略，實現了 2.5 倍加速下分類精度損失 0.4%,，3.3 倍加速下精度損失 0.7%,，且顯著優(yōu)于未壓縮的 ResNet-44 模型； 在 ImageNet 數據集上,，PocketFlow 可以對原本已經十分精簡的 MobileNet 模型繼續(xù)進行權重稀疏化,，以更小的模型尺寸取得相似的分類精度；與 Inception-V1,、ResNet-18 等模型相比,，模型大小僅為后者的約 20~40%，但分類精度基本一致（甚至更高）,。
相比于費時費力的人工調參,，PocketFlow 框架中的 AutoML 自動超參數優(yōu)化組件僅需 10 余次迭代就能達到與人工調參類似的性能，在經過 100 次迭代后搜索得到的超參數組合可以降低約 0.6% 的精度損失,；通過使用超參數優(yōu)化組件自動地確定網絡中各層權重的量化比特數,，PocketFlow 在對用于 ImageNet 圖像分類任務的 ResNet-18 模型進行壓縮時，取得了一致性的性能提升,；當平均量化比特數為 4 比特時,，超參數優(yōu)化組件的引入可以將分類精度從 63.6% 提升至 68.1%（原始模型的分類精度為 70.3%）。
PocketFlow 助力移動端業(yè)務落地
據了解,，在騰訊公司內部,，PocketFlow 框架正在為多項移動端實際業(yè)務提供了模型壓縮與加速的技術支持。例如,，在手機拍照 APP 中,，人臉關鍵點定位模型是一個常用的預處理模塊，通過對臉部的百余個特征點（如眼角,、鼻尖等）進行識別與定位,，可以為后續(xù)的人臉識別、智能美顏等多個應用提供必要的特征數據,。團隊基于 PocketFlow 框架,，對人臉關鍵點定位模型進行壓縮，在保持定位精度不變的同時,，大幅度地降低了計算開銷,，在不同的移動處理器上取得了 25%-50% 不等的加速效果，壓縮后的模型已經在實際產品中得到部署。
結語
深度學習模型的壓縮與加速是當前學術界的研究熱點之一,，同時在工業(yè)界中也有著廣泛的應用前景,。騰訊 AI Lab 表示，隨著 PocketFlow 的推出,，開發(fā)者無需了解模型壓縮算法的具體細節(jié),，也不用關心各個超參數的選擇與調優(yōu)，即可基于這套自動化框架,，快速得到可用于移動端部署的精簡模型,，從而為 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.