騰訊 AI Lab提出自動化模型壓縮框架PocketFlow:將深度學習裝進口袋

機器之心發表於2018-09-17

據機器之心瞭解,騰訊 AI Lab 機器學習中心今日宣佈成功研發自動化深度學習模型壓縮框架——PocketFlow,並即將在近期釋出開原始碼。據介紹,這是一款面向移動端 AI 開發者的自動模型壓縮框架,整合了當前主流(包括騰訊 AI Lab 自研)的模型壓縮與訓練演算法,結合自研引數優化元件實現了全程自動化託管式的模型壓縮與加速。開發者無需瞭解具體演算法細節,即可快速地將 AI 技術部署到移動端產品上,實現使用者資料的本地高效處理。目前該框架正在為騰訊的多項移動端業務提供模型壓縮與加速的技術支援,在多款手機 APP 中得到應用。值得一提的是,不久之前 MIT 韓鬆等人發表論文,提出了用於移動端模型自動壓縮與加速的AutoML方法AMC

Update:11月1日,騰訊 AI Lab 在南京舉辦的騰訊全球合作伙伴論壇上宣佈正式開源「PocketFlow」專案。專案地址:https://github.com/Tencent/PocketFlow

隨著 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 框架主要由兩部分元件構成,分別是模型壓縮/加速演算法元件和引數優化元件,具體結構如下圖所示。

騰訊 AI Lab提出自動化模型壓縮框架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 自動引數優化框架來根據具體效能需求,確定最優引數取值組合。例如,對於通道剪枝演算法,引數優化元件可以自動地根據原始模型中各層的冗餘程度,對各層採用不同的剪枝比例,在保證滿足模型整體壓縮倍數的前提下,實現壓縮後模型識別精度的最大化。

騰訊 AI Lab提出自動化模型壓縮框架PocketFlow:將深度學習裝進口袋

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%,但分類精度基本一致(甚至更高)。

騰訊 AI Lab提出自動化模型壓縮框架PocketFlow:將深度學習裝進口袋

騰訊 AI Lab提出自動化模型壓縮框架PocketFlow:將深度學習裝進口袋

相比於費時費力的人工調參,PocketFlow 框架中的 AutoML 自動引數優化元件僅需 10 餘次迭代就能達到與人工調參類似的效能,在經過 100 次迭代後搜尋得到的引數組合可以降低約 0.6% 的精度損失;通過使用引數優化元件自動地確定網路中各層權重的量化位元數,PocketFlow 在對用於 ImageNet 影象分類任務的 ResNet-18 模型進行壓縮時,取得了一致性的效能提升;當平均量化位元數為 4 位元時,引數優化元件的引入可以將分類精度從 63.6% 提升至 68.1%(原始模型的分類精度為 70.3%)。

騰訊 AI Lab提出自動化模型壓縮框架PocketFlow:將深度學習裝進口袋

騰訊 AI Lab提出自動化模型壓縮框架PocketFlow:將深度學習裝進口袋

PocketFlow 助力移動端業務落地

據瞭解,在騰訊公司內部,PocketFlow 框架正在為多項移動端實際業務提供了模型壓縮與加速的技術支援。例如,在手機拍照 APP 中,人臉關鍵點定位模型是一個常用的預處理模組,通過對臉部的百餘個特徵點(如眼角、鼻尖等)進行識別與定位,可以為後續的人臉識別、智慧美顏等多個應用提供必要的特徵資料。團隊基於 PocketFlow 框架,對人臉關鍵點定位模型進行壓縮,在保持定位精度不變的同時,大幅度地降低了計算開銷,在不同的移動處理器上取得了 25%-50% 不等的加速效果,壓縮後的模型已經在實際產品中得到部署。

騰訊 AI Lab提出自動化模型壓縮框架PocketFlow:將深度學習裝進口袋

結語

深度學習模型的壓縮與加速是當前學術界的研究熱點之一,同時在工業界中也有著廣泛的應用前景。騰訊 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.

相關文章