【論文筆記】AMC：AutoML for Model Compression and Acceleration on Mobile Devices

論文：https://arxiv.org/abs/1802.03494
代碼：https://github.com/mit-han-lab/amc-release

簡介

本篇論文是MIT韓松老師組提出的一種模型壓縮方法，其核心思想是使用強化學習技術來實現自動化壓縮模型。目前對於領域問題的解決方案通常採取人工壓縮模型，手工壓縮方法需要相關領域的專家知識，需要人工在推理速度，大小與準確率之間權衡，因此人工壓縮難以得到最優的壓縮策略。與人工設計模型相比使用強化學習的產生壓縮策略從數據出發能夠找到更合適的壓縮策略。該篇文章講FLOPs減少4倍的情況下，對VGG-16模型壓縮，比人工壓縮策略策略高2.7％的精度。並將這種方法應用到更適合移動端設備的Mobilenet模型上，使模型更加輕量化。

相關工作

與現有的技術相比 （NAS nerual architecture search [57]、NT network Transformation[6]、N2N network to network[2] ）本文提出的方法AMC不是對網絡結構的路徑搜索，而是採用強化學習中的DDPG（深度確定性策略梯度法）來產生連續空間上的具體壓縮比率，通過約束FLOPs和準確率設定獎勵值，因此得到的壓縮策略是適應當前硬件環境的。

方法

方法整體結構圖如下

圖 1：AMC 方法圖示。左：AMC 替代人工使模型壓縮完全自動化，而且表現優於人類。右：將 AMC 映射到強化學習問題場景。珠層壓縮 一個pretrained的 MobileNet。Agent（DDPG）從 t 層接收嵌入 $s_t$，輸出稀疏比率 $a_t$。用 $a_t$ 對層進行壓縮後，智能體移動到下一層 $L_{t+1}$。評估所有層都被壓縮的修剪後模型的準確率。最後，獎勵 R 作爲準確率和 FLOPs 的函數被返回給Agent。

本文使用DDPG方法在前問題場景下，agent的動作空間映射爲 待壓縮模型每一層的壓縮率 ，本文的目標是保持準確率的同時減少FLOPS與壓縮模型的大小，因此回報函數可以根據這些指標進行設計，論文中提到了兩種回報函數，基於FLOPs和參數量的，如下所示

$State Space: (t,n,e,h,w,stride,k,FLOPs[t],reduced,rest,a_{t-1})$
$Loss = \frac{1}{N} \sum_i{(y_i - Q(s_i,a_i| \theta^Q))^2}$
$y_i = r_i - b+\gamma Q(s_{i+1},\mu(s_{i+1})|\theta ^ Q)$
$R_{FLOPs} = - Error · log(FLOPs)$
$R_{Param} = - Error · log(Param)$

算法的僞代碼如圖所示
算法流程 主要通過DDPG算法爲每一層依次產生保留率，源碼中設定的動作邊界爲[0.2，1]，然後根據當前層原始通道數量來計算出需要保留channels的數量n，然後使用ranking的辦法爲通道排序然後 取出前n個通道索引作爲新model保留的通道，然後分別使用 $W_{new}$ 和 $W_{all}$ 兩組權重計算 top1 acc FLOPs等指標來計算獎勵值。

實驗

論文分別使用CIFAR10和ImageNet數據集分別進行實驗來驗證，主要壓縮了Resnet、Vgg和mobilenet驗證，手工壓縮和AMC方法調優之後模型的壓縮率和準確率，實驗結果AMC方法在各個維度上全面超越了手工調參的模型的指標

實驗結果提出了一個很有趣的結論，作者根據實驗結果提出3x3大小的權重包含了更多地冗餘信息因此在3x3層給出的壓縮率較高，而1x1大小的權重包含的冗餘信息小因此給出的壓縮率較低。這裏我認爲應該有更深層次的原因值得研究。

TIPs

總結一下復現過程中的一些坑，這篇論文的代碼已經開源了，但是開源項目中目前只保留了mobilenet模型的剪枝源碼。
1、官方開源代碼在進行通道裁剪過程中 採用的通道ranking的方式是直接將每個通道上的所有權值求和，似乎有點不科學，因此我嘗試使用更科學的方式（幾何中位數距離，範數距離，餘弦相似度等）替換ranking的方式，結果居然都沒有直接求和的效果好讓我很是失落。
2、同時這部分代碼僅能夠對mobilenet 的1x1 或 fc 層進行權重重構 reconstruction，其他的大卷積核的層寫了個 Todo 留了個小坑。
3、當需要壓縮其他模型時需要注意模型的多分支結構，例如resnet inception之類的，需要保持裁剪通道數的一致。另外在構建模型的時候一定要寫成基於配置參數構建model的形式，不然每剪枝一次產生的model都不一樣就需要手動重寫新的model，會要老命的！！

最後我發現了一個Intel的蒸餾工具包
https://github.com/NervanaSystems/distiller
這份代碼整合了時下比較經典的壓縮剪枝算法，當然也整合了本篇論文的AMC算法。設計模式思想體現的淋漓盡致，不愧爲工業界的正經代碼。拜讀源碼好幾天，真的是被秀到了。想要用起這份代碼還是需要一定功力的，如果只是paper做實驗推薦官方開源的代碼，在基礎上進行修改還是比較容易的，因爲功能比較單一註釋也很友好通俗易懂。