Deep Compression閱讀理解及Caffe源碼修改

作者：may0324

更新：?
沒想到這篇文章寫出后有這么多人關注和索要源碼，有點受寵若驚。說來慚愧，這個工作當時做的很粗糙，源碼修改的比較亂，所以一直不太好拿出手。最近終于有時間整理了一下代碼并開源出來了。關于代碼還有以下幾個問題：?
~1.在.cu中目前仍然是調用cpu_data接口，所以可能會增加與gpu數據交換的額外耗時，這個不影響使用，后面慢慢優化。~(已解決)?
2.目前每層權值修剪的比例仍然是預設的，這個比例需要迭代試驗以實現在盡可能壓縮權值的同時保證精度。所以如何自動化選取閾值就成為了后面一個繼續深入的課題。?
3.直接用caffe跑出來的模型依然是原始大小，因為模型依然是.caffemodel類型，雖然大部分權值為0且共享，但每個權值依然以32float型存儲，故后續只需將非零權值及其下標以及聚類中心存儲下來即可，這部分可參考作者論文，寫的很詳細。?
4.權值壓縮僅僅壓縮了模型大小，但在前向inference時不會帶來速度提升。因此，想要在移動端做好cnn部署，就需要結合小的模型架構、模型量化以及NEON指令加速等方法來實現。?
代碼開源在github?
https://github.com/may0324/DeepCompression-caffe?

最近又轉戰CNN模型壓縮了。。。（我真是一年換N個坑的節奏），閱讀了HanSong的15年16年幾篇比較有名的論文，啟發很大，這篇主要講一下Deep Compression那篇論文，因為需要修改caffe源碼，但網上沒有人po過，這里做個第一個吃螃蟹的人，記錄一下對這篇論文的理解和源碼修改過程，方便日后追本溯源，同時如果有什么紕漏也歡迎指正，互相交流學習。?
這里就從Why-How-What三方面來講講這篇文章。

Why

首先講講為什么CNN模型壓縮刻不容緩，我們可以看看這些有名的caffe模型大小:?
1. LeNet-5?1.7MB?
2. AlexNet?240MB?
3. VGG-16?552MB?
LeNet-5是一個簡單的手寫數字識別網絡，AlexNet和VGG-16則用于圖像分類，刷新了ImageNet競賽的成績，但是就其模型尺寸來說，根本無法移植到手機端App或嵌入式芯片當中，就算是想通過網絡傳輸，較高的帶寬占用率也讓很多用戶望塵莫及。另一方面，大尺寸的模型也對設備功耗和運行速度帶來了巨大的挑戰。隨著深度學習的不斷普及和caffe,tensorflow,torch等框架的成熟，促使越來越多的學者不用過多地去花費時間在代碼開發上，而是可以毫無顧及地不斷設計加深網絡，不斷擴充數據，不斷刷新模型精度和尺寸，但這樣的模型距離實用卻仍是望其項背。?
在這樣的情形下，模型壓縮則成為了亟待解決的問題，其實早期也有學者提出了一些壓縮方法，比如weight prune(權值修剪)，權值矩陣SVD分解等，但壓縮率也只是冰山一角，遠不能令人滿意。今年standford的HanSong的ICLR的一篇論文Deep Compression: Compressing deep neural networks with pruning, trained quantization and Huffman coding一經提出，就引起了巨大轟動，在這篇論文工作中，他們采用了3步，在不損失（甚至有提升）原始模型精度的基礎上，將VGG和Alexnet等模型壓縮到了原來的35~49倍，使得原本上百兆的模型壓縮到不到10M，令深度學習模型在移動端等的實用成為可能。

How

Deep Compression 的實現主要有三步，如下圖所示：?

包括Pruning（權值修剪），Quantization（權值共享和量化），Huffman Coding（Huffman編碼）。

1.Prunning

如果你調試過caffe模型，觀察里面的權值，會發現大部分權值都集中在-1~1之間，即非常小，另一方面，神經網絡的提出就是模仿人腦中的神經元突觸之間的信息傳導，因此這數量龐大的權值中，存在著不可忽視的冗余性，這就為權值修剪提供了根據。pruning可以分為三步：?
step1. 正常訓練模型得到網絡權值；?
step2. 將所有低于一定閾值的權值設為0；?
step3. 重新訓練網絡中剩下的非零權值。?
經過權值修剪后的稀疏網絡，就可以用一種緊湊的存儲方式CSC或CSR（compressed sparse column or compressed sparse row）來表示。這里舉個栗子來解釋下什么是CSR?
假設有一個原始稀疏矩陣A?
?
CSR可以將原始矩陣表達為三部分，即AA,JA,IC?
?
其中，AA是矩陣A中所有非零元素，長度為a，即非零元素個數；?
JA是矩陣A中每行第一個非零元素在AA中的位置，最后一個元素是非零元素數加1，長度為n+1, n是矩陣A的行數；?
IC是AA中每個元素對應的列號，長度為a。?
所以將一個稀疏矩陣轉為CSR表示，需要的空間為2*a+n+1個，同理CSC也是類似。?
可以看出，為了達到壓縮原始模型的目的，不僅需要在保持模型精度的同時，prune掉盡可能多的權值，也需要減少存儲元素位置index所帶來的額外存儲開銷，故論文中采用了存儲index difference而非絕對index來進一步壓縮模型，如下圖所示：?
?
其中，第一個非零元素的存儲的是他的絕對位置，后面的元素依次存儲的是與前一個非零元素的索引差值。在論文中，采用固定bit來存儲這一差值，以圖中表述為例，如果采用3bit，則最大能表述的差值為8，當一個非零元素距其前一個非零元素位置超過8，則將該元素值置零。（這一點其實也很好理解，如果兩個非零元素位置差很多，也即中間有很多零元素，那么將這一元素置零，對最終的結果影響也不會很大）?
做完權值修剪這一步后，AlexNet和VGG-16模型分別壓縮了9倍和13倍，表明模型中存在著較大的冗余。

2.Weight Shared & Quantization

為了進一步壓縮網絡，考慮讓若干個權值共享同一個權值，這一需要存儲的數據量也大大減少。在論文中，采用kmeans算法來將權值進行聚類，在每一個類中，所有的權值共享該類的聚類質心，因此最終存儲的結果就是一個碼書和索引表。?
1.對權值聚類?
論文中采用kmeans聚類算法，通過優化所有類內元素到聚類中心的差距（within-cluster sum of squares ）來確定最終的聚類結果：?
?
式中，W ={w1,w2,…wn}是n個原始權值，C={c1,c2,…ck}是k個聚類。?
需要注意的是聚類是在網絡訓練完畢后做的，因此聚類結果能夠最大程度地接近原始網絡權值分布。?
2. 聚類中心初始化?
常用的初始化方式包括3種：?
a) 隨機初始化。即從原始數據種隨機產生k個觀察值作為聚類中心。?
b) 密度分布初始化。現將累計概率密度CDF的y值分布線性劃分，然后根據每個劃分點的y值找到與CDF曲線的交點，再找到該交點對應的x軸坐標，將其作為初始聚類中心。?
c) 線性初始化。將原始數據的最小值到最大值之間的線性劃分作為初始聚類中心。?
三種初始化方式的示意圖如下所示：?

由于大權值比小權值更重要（參加HanSong15年論文），而線性初始化方式則能更好地保留大權值中心，因此文中采用這一方式，后面的實驗結果也驗證了這個結論。?
3. 前向反饋和后項傳播?
前向時需要將每個權值用其對應的聚類中心代替，后向計算每個類內的權值梯度，然后將其梯度和反傳，用來更新聚類中心，如圖：?

共享權值后，就可以用一個碼書和對應的index來表征。假設原始權值用32bit浮點型表示，量化區間為256，即8bit，共有n個權值，量化后需要存儲n個8bit索引和256個聚類中心值，則可以計算出壓縮率compression ratio:?
r = 32*n / (8*n + 256*32 )≈4?
可以看出，如果采用8bit編碼，則至少能達到4倍壓縮率。

3.Huffman Coding

Huffman 編碼是最后一步，主要用于解決編碼長短不一帶來的冗余問題。因為在論文中，作者針對卷積層統一采用8bit編碼，而全連接層采用5bit，所以采用這種熵編碼能夠更好地使編碼bit均衡，減少冗余。

4.Evaluation

實驗結果就是能在保持精度不變（甚至提高）的前提下，將模型壓縮到前所未有的小。直接上圖有用數據說話。?

5.Discussion

不同模型壓縮比和精度的對比，驗證了pruning和quantization一塊做效果最好。

不同壓縮bit對精度的影響，同時表明conv層比fc層更敏感， 因此需要更多的bit表示。

不同初始化方式對精度的影響，線性初始化效果最好。

卷積層采用8bit，全連接層采用5bit效果最好。

What

此部分講一講修改caffe源碼的過程。其實只要讀懂了文章原理，修改起來很容易。?
對pruning過程來說，可以定義一個mask來“屏蔽”修剪掉的權值，對于quantization過程來說，需定義一個indice來存儲索引號，以及一個centroid結構來存放聚類中心。?
在include/caffe/layer.hpp中為Layer類添加以下成員變量：?

以及成員函數：?

由于只對卷積層和全連接層做壓縮，因此，只需修改這兩個層的對應函數即可。

在include/caffe/layers/base_conv_layer.hpp添加成員函數?

這兩處定義的函數都是基類的虛函數，不需要具體實現。

在include/caffe/layers/conv_layer.hpp中添加成員函數聲明：?

類似的，在include/caffe/layers/inner_product_layer.hpp也添加該函數聲明。?
在src/caffe/layers/conv_layer.cpp 添加該函數的聲明，用于初始化mask和對權值進行聚類。?

同時，修改前向和后向函數。?
在前向函數中，需要將權值用其聚類中心表示，紅框部分為添加部分：?

在后向函數中，需要添加兩部分，一是對mask為0，即屏蔽掉的權值不再進行更新，即將其weight_diff設為0，另一個則是統計每一類內的梯度差值均值，并將其反傳回去，紅框內為添加部分。?


kmeans的實現如下，當然也可以用Opencv自帶的，速度會更快些。

<span style="color:#565f69"><span style="color:#333333"><code>template<typename Dtype> void kmeans_cluster(vector<int> &cLabel, vector<Dtype> &cCentro, Dtype *cWeights, int nWeights, vector<int> &mask, /*Dtype maxWeight, Dtype minWeight,*/ int nCluster, int max_iter /* = 1000 */) {//find min maxDtype maxWeight=numeric_limits<Dtype>::min(), minWeight=numeric_limits<Dtype>::max();for(int k = 0; k < nWeights; ++k){if(mask[k]){if(cWeights[k] > maxWeight)maxWeight = cWeights[k];if(cWeights[k] < minWeight)minWeight = cWeights[k];}}// generate initial centroids linearlyfor (int k = 0; k < nCluster; k++)cCentro[k] = minWeight + (maxWeight - minWeight)*k / (nCluster - 1);//initialize all label to -1for (int k = 0; k < nWeights; ++k)cLabel[k] = -1;const Dtype float_max = numeric_limits<Dtype>::max();// initializeDtype *cDistance = new Dtype[nWeights];int *cClusterSize = new int[nCluster];Dtype *pCentroPos = new Dtype[nCluster];int *pClusterSize = new int[nCluster];memset(pClusterSize, 0, sizeof(int)*nCluster);memset(pCentroPos, 0, sizeof(Dtype)*nCluster);Dtype *ptrC = new Dtype[nCluster];int *ptrS = new int[nCluster];int iter = 0;//Dtype tk1 = 0.f, tk2 = 0.f, tk3 = 0.f;double mCurDistance = 0.0;double mPreDistance = numeric_limits<double>::max();// clusteringwhile (iter < max_iter){// check convergenceif (fabs(mPreDistance - mCurDistance) / mPreDistance < 0.01) break;mPreDistance = mCurDistance;mCurDistance = 0.0;// select nearest clusterfor (int n = 0; n < nWeights; n++){if (!mask[n])continue;Dtype distance;Dtype mindistance = float_max;int clostCluster = -1;for (int k = 0; k < nCluster; k++){distance = fabs(cWeights[n] - cCentro[k]);if (distance < mindistance){mindistance = distance;clostCluster = k;}}cDistance[n] = mindistance;cLabel[n] = clostCluster;}// calc new distance/inertiafor (int n = 0; n < nWeights; n++){if (mask[n])mCurDistance = mCurDistance + cDistance[n];}// generate new centroids// accumulation(private)for (int k = 0; k < nCluster; k++){ptrC[k] = 0.f;ptrS[k] = 0;}for (int n = 0; n < nWeights; n++){if (mask[n]){ptrC[cLabel[n]] += cWeights[n];ptrS[cLabel[n]] += 1;}}for (int k = 0; k < nCluster; k++){pCentroPos[ k] = ptrC[k];pClusterSize[k] = ptrS[k];}//reduction(global)for (int k = 0; k < nCluster; k++){cCentro[k] = pCentroPos[k];cClusterSize[k] = pClusterSize[k];cCentro[k] /= cClusterSize[k];}iter++;// cout << "Iteration: " << iter << " Distance: " << mCurDistance << endl;}//gather centroids//#pragma omp parallel for//for(int n=0; n<nNode; n++)// cNodes[n] = cCentro[cLabel[n]];delete[] cDistance;delete[] cClusterSize;delete[] pClusterSize;delete[] pCentroPos;delete[] ptrC;delete[] ptrS; }</code></span></span>

連接層的修改和卷積層的一致不再贅述。同樣的，可以把對應的.cu文件中的gpu前向和后向函數實現也修改了，方便gpu訓練。?
最后，在src/caffe/net.cpp的CopyTrainedLayersFrom(const NetParameter& param)函數中調用我們定義的函數，即在讀入已經訓練好的模型權值時，對每一層做需要的權值mask初始化和權值聚類。?

至此代碼修改完畢，編譯運行即可。

Reference

[1] SqueezeNet: AlexNet-level accuracy with 50x fewer parameters and <0.5MB model size?
[2] Deep Compression: Compressing deep neural networks with pruning, trained quantization and Huffman coding?
[3] Learning both Weights and Connections for Efficient Neural Networks?
[4] Efficient Inference Engine on Compressed Deep Neural Network

總結：

最后再提一句，幾乎所有的模型壓縮文章都是從Alexnet和VGG下手，一是因為他們都采用了多層較大的全連接層，而全連接層的權值甚至占到了總參數的90%以上，所以即便只對全連接層進行“開刀”，壓縮效果也是顯著的。另一方面，這些論文提出的結果在現在看來并不是state of art的，存在可提升的空間，而且在NIN的文章中表明，全連接層容易引起過擬合，去掉全連接層反而有助于精度提升，所以這么看來壓縮模型其實是個不吃力又討好的活，獲得的好處顯然是雙倍的。但運用到特定的網絡中，還需要不斷反復試驗，因地制宜，尋找適合該網絡的壓縮方式。

總結

以上是生活随笔為你收集整理的Deep Compression阅读理解及Caffe源码修改的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。