[2024] RL-Pruner: Structured Pruning Using Reinforcement Learning for CNN Compression and Acceleration

一、论文说明

论文标题：使用强化学习进行结构化剪枝用于卷积神经网路压缩和加速

机构：伊利诺伊大学厄巴纳-香槟分校

论文链接：https://arxiv.org/pdf/2411.06463

代码链接：https://github.com/Beryex/RLPruner-CNN

论文简介： 卷积神经网络（ConvolutionalNeural Networks, CNNs）近年来表现出卓越的性能。压缩这些模型不仅减少了存储需求，使其在边缘设备上的部署变得可行，还加速了推理，从而降低了延迟和计算成本。结构化剪枝，它在层级上去除过滤器，直接修改了模型架构。这种方法实现了更紧凑的架构，同时保持目标准确性，确保压缩模型具有较好的兼容性和硬件效率。所提方法基于一个关键观察：

• 1、神经网络中不同层的过滤器对模型性能的重要性各不相同。

• 2、当修剪的过滤器数量固定时，不同层之间的最佳修剪分配是不均匀的，以最小化性能损失。

• 3、对修剪敏感的层应该占据更小的修剪分配比例。

为了利用这一洞察，文中提出了RL-Pruner，它使用强化学习来学习最佳修剪分配。RL-Pruner可以自动提取输入模型中过滤器之间的依赖关系并执行修剪，无需特定于模型的修剪实现。在GoogleNet、ResNet和MobileNet 等模型上进行了实验，将所提方法与其他结构化剪枝方法进行了比较，以验证其有效性。

二、原理

RL-Pruner 首先在模型中的层之间构建依赖图，然后分几个步骤进行剪枝。在每个步骤中：1) 基于基础分布生成一个新的剪枝稀疏分布作为动作 ，这作为策略；2)根据相应的稀疏度，使用泰勒准则（Taylorcriterion）对每一层进行剪枝；3) 评估压缩后的模型以获得奖励，并将动作和奖励存储在回放（replay）经验池中。每个步骤后，基础分布根据经验池更新，如果计算资源足够，则对压缩模型应用后训练阶段，使用知识蒸馏（knowledge distillation），其中原始模型作为教师。具体框图如图2所示。

三、实验与分析

1、环境配置

实验平台及软件

• Windows 10
• git bash
• conda环境

这里主要介绍如何在windows系统上让git bash链接conda环境。

在Windows配置git bash链接conda环境

由于工程代码中需要使用bash命令运行代码，因此需要保证git bash能调用conda环境运行对应的脚本文件。

在C:\Users\username\.bashrc文件内设置conda.sh位置（文中示例为：D:\\Anaconda3\\etc\\profile.d\\conda.sh），并激活配置。在git bash界面输入具体命令如下：

echo "D:\\Anaconda3\\etc\\profile.d\\conda.sh" >> ~/.bashrc  
source ~/.bashrc

然后关闭git bash界面，再重新打开一个git bash界面，最后输入命令激活conda环境conda activate 虚拟环境名字。如果命令提示中出现如下图所示的字样，即为配置成功，否则根据提示的要求进行配置，比如输入conda init，重新打开一个新的git bash界面。

2、项目代码运行

克隆项目文件，具体命令如下：

git clone https://github.com/Beryex/RLPruner-CNN.git --depth 1
cd RLPruner-CNN

安装python第三方包，具体命令如下（如果之前有conda环境，可以不用进行下面这一步，等报错了再根据提示安装对应的包即可）：

conda create -n RLPruner python=3.10 -y
conda activate RLPruner
pip install -r requirements.txt

官方代码提供了一步到位的运行脚本，从预训练模型、模型压缩到模型验证，仅需在命令行中输入如下代码：

./scripts/flexible.sh googlenet cifar100 0.20 taylor 0.00 0.00

为了更好地了解每一步的设置，下面内容将分为预训练模型、模型压缩、模型验证三个步骤进行介绍。

1、训练预训练权重

训练模型得到对应的预训练权重，这里以resnet32和googlenet为例，在git bash输入具体命令（默认使用cuda）如下：

./scripts/train.sh googlenet cifar100
./scripts/train.sh resnet32 cifar100

或者使用参考指定配置命令：

python -m train --model ${MODEL} --dataset ${DATASET} --device cuda \
                --output_dir ${PRETRAINED_MODEL_DIR} \
                --log_dir ${LOG}

其中，
${MODEL}代表backbone的类型（[“vgg11”, “vgg13”, “vgg16”, “vgg19”, “resnet18”, “resnet34”, “resnet50”, “resnet101”, “resnet152”, “resnet8”, “resnet14”, “resnet20”, “resnet32”, “resnet44”, “resnet56”, “resnet110”, “densenet121”, “densenet161”, “densenet169”, “densenet201”, “mobilenetv3_small”, “mobilenetv3_large”, “googlenet”]）；
${DATASET}代表数据集名称，如cifar10或者cifar100。
${PRETRAINED_MODEL_DIR}代表输出权重文件路径，默认在pretrained_model文件夹下；
${LOG}代表输出日志路径，默认在log文件夹下。

在CIFAR100数据集上训练resnet32的结果（最佳准确率：0.706）如下图所示。

在CIFAR100数据集上训练googlenet的结果（最佳准确率：0.774）如下图所示。

2、模型压缩

模型结构化剪枝这里以0.2的稀疏度，taylor剪枝策略和Q_FLOP_coef=0，Q_Para_coef=0的参数进行测试。在git bash输入具体命令（默认使用cuda）如下：

./scripts/flexible.sh googlenet cifar100 0.20 taylor 0.00 0.00

同理，也可以使用参考指定配置命令：

python -m compress --model ${MODEL} --dataset ${DATASET} --device cuda \
                   --sparsity ${SPARSITY} --prune_strategy ${prune_strategy} --ppo \
                   --Q_FLOP_coef ${Q_FLOP_coef} --Q_Para_coef ${Q_Para_coef} \
                   --pretrained_pth ${PRETRAINED_MODEL_PTH} \
                   --compressed_dir ${COMPRESSED_MODEL_DIR} \
                   --checkpoint_dir ${CKPT_DIR} \
                   --log_dir ${LOG} --save_model

测试结果如下图所示：

3、模型验证

在数据集上验证模型的识别性能，在git bash输入具体命令（默认使用cuda）如下：

./scripts/evaluate.sh googlenet cifar100

同理，也可以使用参考指定配置命令：

python -m evaluate --model ${MODEL} --dataset ${DATASET} --device cuda \
                   --pretrained_pth ${PRETRAINED_MODEL_PTH} \
                   --compressed_pth ${COMPRESSED_MODEL_PTH} \
                   --log_dir ${LOG}

测试结果如下图所示：

四、总结

本文提出了RL-Pruner，一种结构化剪枝方法，能够学习各层之间的最优稀疏性分布，并支持无模型特定修改的一般剪枝。希望所提方法能够认识到每一层对模型（model）性能的重要性不同，这将影响未来在神经网络压缩领域的工作，包括无结构剪枝和量化。