pytorch pruning训练自己的数据库（流程+BUG调试）

环境要求

pytorch 0.1
0.2版本会有问题，解决方法后文会介绍。
如何查看版本？在终端输入：
python
import torch
torch.version

数据集制作

本文代码使用的数据加载方法是datasets.ImageFolder，它要求数据集不同类别的图片放在不同文件夹下，文件格式如下：

将自己的数据集做成图中的形式，即可。

运行过程

train

python finetune.py –train –train_path path_to_your_dataset/train/ –test_path path_to_your_dataset/test/
这里会去掉VGG的后三层fc，而根据自己的num_class训练新的fc层，总共迭代20个epoch。

prune

python finetune.py –prune –train_path path_to_your_dataset/train/ –test_path path_to_your_dataset/test/
这里会开始对filter进行剪枝，每次剪512个filter，每次剪完都会迭代10个epoch以恢复模型的能力；然后继续下一次剪枝，直至将VGG模型的2/3的filter剪掉；最后进行15个epoch以得到最终的剪枝模型。

ps：这里的512是自己设定的，论文中是1，但每次只剪1个filter太慢，所以设定为512，加快剪枝过程。

上图是prune的过程，其中Layers that will be prunned是每一层剪枝的filter数量，可以看到更深的层减的更多。

---

以下是在调试过程中遇到的问题：

问题调试

1. ImportError cv2

解决办法：
将import cv2放在import torch前面
参考网址

2. RuntimeError THCTensorCopy.c:18

/b/wheel/pytorch-src/torch/lib/THCUNN/ClassNLLCriterion.cu:52: void cunn_ClassNLLCriterion_updateOutput_kernel(Dtype , Dtype , Dtype , long , Dtype *, int, int, int, int) [with Dtype = float, Acctype = float]: block: [0,0,0], thread: [0,0,0] Assertion t >= 0 && t < n_classes failed.

RuntimeError: cuda runtime error (59) : device-side assert triggered at /b/wheel/pytorch-src/torch/lib/THC/generic/THCTensorCopy.c:18
因为我的图片class为10类，而代码中只有2类，所以要将finetune.py中的self.classifier 的最后输出改为10
python
self.classifier = nn.Sequential(
nn.Dropout(),
nn.Linear(25088, 4096),
nn.ReLU(inplace=True),
nn.Dropout(),
nn.Linear(4096, 4096),
nn.ReLU(inplace=True),
nn.Linear(4096, 10))


3. out of memory

THCudaCheck FAIL file=/tmp/luarocks_cutorch-scm-1-3726/cutorch/lib/THC/generic/THCStorage.cu line=66 error=2 : out of memory

解决办法1：换一颗GPU
在代码中修改：
import torch
torch.cuda.set_device(id)

或者在终端命令中加入：CUDA_VISIBLE_DEVICES=1

解决办法2：减小batchsize
参考网址

4. RuntimeError: arguments are located on different GPUs

RuntimeError: arguments are located on different GPUs at /b/wheel/pytorch-src/torch/lib/THC/generated/../generic/THCTensorMathPointwise.cu:214

这种情况是属于一部分数据在CPU上，一部分数据在GPU上，导致无法进行pointwise add。

具体出错代码如下：

values = torch.sum((activation * grad), dim = 0).\ sum(dim=2).sum(dim=3)[0, :, 0, 0].data  \   #将activation与grad做点乘

    # Normalize the rank by the filter dimensions
if activation_index not in self.filter_ranks:
        self.filter_ranks[activation_index] = \
            torch.FloatTensor(activation.size(1)).zero_().cuda()  #得到和activation.size(1)一样长的一维0向量

    self.filter_ranks[activation_index] += values  #将算得的values赋给ranks

可以看到，values上icpu上的数据，而filter_ranks有.cuda()，所以是GPU上的数据，所以会出错，解决办法是将：
values = torch.sum((activation * grad), dim = 0).\ sum(dim=2).sum(dim=3)[0, :, 0, 0].data
换成GPU上的数据：
values = torch.sum((activation * grad), dim = 0).\ sum(dim=2).sum(dim=3)[0, :, 0, 0].data.cuda()

pytorch0.2版本兼容问题

1. keepdim

pytorch2版本中的sum函数多了keepdim这个参数，所以需要在每一个sum函数里加上这个参数。

values = \
torch.sum((activation * grad), dim = 0).\   
sum(dim=2).sum(dim=3)[0, :, 0, 0].data

换成

values = \
            torch.sum((activation * grad), dim = 0,keepdim=True).\   
            sum(dim=2,keepdim=True).sum(dim=3,keepdim=True)[0, :, 0, 0].data

2 bias

pytorch2中conv2d()的init函数的输入最后一项bias是bool型，而conv.bias是float tensor，所以源代码中将bias = conv.bias肯定会出错，解决办法就是去掉这一项：（原文中有两处需要删掉）

new_conv = \
    torch.nn.Conv2d(in_channels = conv.in_channels, \
        out_channels = conv.out_channels - 1, 
        kernel_size = conv.kernel_size, \
        stride = conv.stride,
        padding = conv.padding,
        dilation = conv.dilation,
        groups = conv.groups,
        bias = conv.bias)

改为：

new_conv = \
    torch.nn.Conv2d(in_channels = conv.in_channels, \
        out_channels = conv.out_channels - 1, 
        kernel_size = conv.kernel_size, \
        stride = conv.stride,
        padding = conv.padding,
        dilation = conv.dilation,
        groups = conv.groups)