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的文章中表明，全连接层容易引起过拟合，去掉全连接层反而有助于精度提升，所以这么看来压缩模型其实是个不吃力又讨好的活，获得的好处显然是双倍的。但运用到特定的网络中，还需要不断反复试验，因地制宜，寻找适合该网络的压缩方式。