基于Caffe的CNN剪枝

背景

传统的CNN网络训练完之后，全连接层的权值矩阵动辄就几十万、几百万个参数值，可见CNN模型的庞大，但是仔细观察CNN的权值矩阵就会发现，里面有很多的参数的绝对值都很小，比如在-0.001到0.001之间，也就是说这些连接对CNN的训练或者测试结果作用很小，因此我们就可以尝试将这些小值参数去掉，既可以减小模型的规模又可以减少计算量，最重要的前提是要保证CNN的有效性，也即正确率。

主要思路

• 修改blob的结构，将原来的矩阵改写成稀疏矩阵的存储方式
• 采用新的方法来计算稀疏矩阵和向量的相乘

具体实现

blob的修改

在这里需要对blob.hpp和blob.cpp进行修改：

1.blob.hpp的修改（include/caffe/blob.hpp）

在原来的blob中，存有data_、diff_、shape_data_、shape_、count_、capacity_这6个属性。因为我们要将原来的矩阵（后文为了区分称为密集矩阵）存储为稀疏矩阵，所以要添加新的属性来存储稀疏后的矩阵参数。稀疏矩阵的存储方式可以参考这里，在这里我们添加了3个向量csrval_、csrrowptr_、csrcolind_，这三个变量分别存储着所有非零元素值、非零元素行指针、非零元素列索引。除了这三个新的变量外，还需要添加三个变量nnz_、sparse_、mask_，nnz_用来存储非零元素的个数，sparse_用来表征data是否需要进行稀疏存储，第三个变量mask_需要重点说一下。在我们剪枝的过程中会把data中的一些元素置为零，大量的元素值为零之后势必会影响网络的准确性，所以需要重新训练，将剩下的非零权值进行一次再训练，为了保证在再训练过程中非零元素不会被反馈过程更改掉，我们需要加一个mask_,用来标示该元素是否需要进行梯度更新，该mask_在最初的初始化时应该全为1，在剪枝阶段进行更新。
除了给blob添加新的属性之外，还需要给新加入的属性添加相应的set和get方法，添加方法时参考blob中data和diff的方法(由于源码太长，在此就不粘贴了，具体查看源码)。

2.blob.cpp的修改（src/caffe/blob.cpp）

首先将新添加的变量的get和set方法实现，这部分比较简单，基本上都是复制粘贴修改变量名。除此之外还有三个比较重要的函数：Update(),FromProto()和ToProto().

Update()

该函数主要用来在每次后向反馈之后对blob中的data参数进行更新，因为我们添加了mask_矩阵，所以需要在正常反馈之后将更新值屏蔽掉，于是我们在更新之后将data_和mask_的对应位相乘，屏蔽掉更新，在这里我们调用了caffe中的caffe_gpu_mul()方法,代码如下：

if(sparse_&&FLAGS_step!="three")
    caffe_gpu_mul<Dtype>(count_,
        static_cast<const Dtype*>(mask_->gpu_data()),
        static_cast<const Dtype*>(data_->gpu_data()),
        static_cast<Dtype*>(data_->mutable_gpu_data()));

上面代码最上方有一个if判断，sparse_用来判断当前blob是否是需要进行稀疏压缩的blob,FLAGS_step用来表征当前是第几阶段，如果是第三阶段，则不进行该过程。

FLAGS_Step

在介绍ToProto()和FromProto()之前，先介绍一下FLAGS_step。如果从整体上去观察我们的剪枝过程，可以将其分成三步：
1. 常规训练CNN网络，并保存训练后的模型，然后将小值参数置为零
2. 对置零后的网路进行再训练，保存最终的caffemodel
3. 读入caffemodel进行测试
我们将步骤一中最先保存的caffemodel记为origin，小值置为零的caffemodel记为fixed，将步骤二中再次训练好的caffemodel记为retune。这三个不同的caffemodel除了名字之外，还有很多不同，下面通过表格列举一下。

model	data	diff	mask	csr	是否为稀疏矩阵
origin	保存	不保存	保存	保存	否
fixed	保存	不保存	保存	保存	是
retune	不保存	不保存	不保存	保存	是

注：上图中的’保存’代表：该caffemodel中保存了该项参数值，’csr’代表：csrval、csrrowptr和csrcolind这三个向量的总称。
因为在剪枝的不同阶段生成的caffemodel是不同的，所以在将训练好的网络保存下来和读入时需要根据不同阶段区别对待。为了区别不同阶段，我们引入了FLAGS_step这个全局变量。该变量可以通过命令行读入，关于FLAGS_name形式的全局变量，可以参考这篇博文。

ToProto()

该函数定义了如何将网络训练的权值参数保存进caffemodel中，比如是否将diff_保存进caffemodel中。
在该函数中最主要的修改是实现了对密集矩阵的稀疏处理，生成csrval_、csrrowptr_、csrcolind_和nnz_,将稀疏矩阵进行保存（关于如何生成的csr相关向量，我们单独放在下面一节说）。在将参数矩阵保存进caffemodel时，主要通过sparse_和FLAGS_step这两个变量进行控制。只有sparse_为true时，我们才会对当前blob进行稀疏化处理，否则只进行常规处理。当需要对该blob进行稀疏化处理时，只有FLAGS_step等于one的时候，才会保存data和mask，否则不保存这两个参数矩阵。

稀疏矩阵的存储

稀疏矩阵的存储可以说是CNN剪枝的重点，在实现中，我们调用了CUDA的cuSPARSE库，该库主要是为了优化稀疏矩阵的运算，提供了很多方便易用的接口，在这里我们用到了它的cusparseSnnz(),cusparseSdense2csc(),cusparseScsrmv()这几个函数接口，cusparseSnnz()主要是用来求出矩阵的非零元素个数，cusparseSdense2csc主要是生成矩阵的csrval、csrrowptr和csrcolind这几个特征，cusparseScsrmv()主要是计算稀疏矩阵和向量相乘的，这个才是我们的最终目的，在这个函数中我们需要传入上面生成的那几个csr向量。在此有几个坑，我简单说一下。首先，在caffe中矩阵是行主序的，但是在cuda中矩阵式列主序的，行主序就是把矩阵一行一行的存入内存，列主序是把矩阵一列一列的存入内存，这也是为什么我用的是cusparseScsc()而不是cusparseScsr()；第二点是要注意CPU和GPU之间的数据交换，需要用相应的函数cudaMemcpy()去复制一份，否则会报错；第三点是cuSPARSE库的异步性，要想同步执行各个函数，需要明确指定，可以用cudaDeviceSynchronize();来完成。

FromProto()

与ToProto()相反，该函数主要是将权值矩阵从caffemodel中读出来，根据FLAGS_step和sparse_的不同，有选择的读出csr、data、mask等。在这个地方需要注意的是，因为blob的reshape()中没有对csr进行初始化，所以在进行读出csr时，需要先给csr申请空间，然后再读出。

proto的修改

caffe.proto