智能剪枝:RL-Pruner 优化 GoogleNet、ResNet 和 MobileNet 实现高效模型压缩 !

大模型向量数据库机器学习

点击下方卡片,关注 「AI视界引擎」 公众号

( 添加时备注:方向+学校/公司+昵称/姓名 )

picture.image

picture.image

近年来,卷积神经网络(CNNs)表现出了令人瞩目的性能。压缩这些模型不仅可以降低存储要求,使其能够在边缘设备上部署,还可以加速推理,从而减少延迟和计算成本。

结构化剪枝,即在层 Level 移除滤波器,直接修改模型架构。这种方法可以在保持目标准确性的同时实现更紧凑的架构,确保压缩后的模型具有良好的兼容性和硬件效率。

作者的方法基于一个关键观察:神经网络不同层中的滤波器对模型性能的重要性不同。当要剪枝的滤波器数量固定时,不同层的最优剪枝分布是不均匀的,以最小化性能损失。更敏感于剪枝的层应占较小比例的剪枝分布。

为了利用这一洞察,作者提出了一种基于强化学习的剪枝器,称为RL-Pruner。RL-Pruner可以自动提取输入模型中滤波器之间的依赖关系并执行剪枝,而无需特定于模型的剪枝实现。

作者在GoogleNet、ResNet和MobileNet等模型上进行了实验,将其与其它结构化剪枝方法进行了比较,以验证其有效性。

作者的代码已在 https://github.com/Beryex/RLPruner-CNN上公开。

1 Introduction

卷积神经网络(CNNs)在计算机视觉任务上表现出色,包括图像分类、检测和分割。随着这些架构变得更宽更深,它们从输入数据中提取复杂特征的能力得到增强。然而,参数数量的增长导致计算成本显著增加,推理既消耗资源又慢。此外,将CNN部署在边缘设备上的需求日益增长,这些设备通常计算能力和内存有限。因此,有效的方法来压缩CNN的需求越来越大,旨在减少参数数量并加速推理,同时保持原始模型性能。这些压缩技术对于创建高效可部署的CNN,以满足实际应用的挑战至关重要。

已有几种压缩卷积神经网络(CNNs)的技术被提出,这些技术大多可以归入四个类别:结构化和非结构化剪枝、量化、低秩分解和知识蒸馏。结构化剪枝从神经网络中移除整个滤波器,并直接修改模型的架构,从而实现压缩和实时的加速。非结构化剪枝,也称为权重剪枝,从权重矩阵中移除特定不重要的元素,这需要硬件或支持稀疏计算的库来实现实际加速。低秩分解通过分解权重矩阵将其近似为低秩矩阵。量化减少模型中权值数据的位宽,从而实现显著的压缩,但同时也需要硬件支持来实现低位量化的理论加速。知识蒸馏从较大的、更先进的教师模型中转移知识到较小的学生模型,可能不会生成新的压缩架构,但有助于恢复压缩导致的性能损失。在本论文中,作者主要关注结构化剪枝,但作者希望作者的方法的洞察力也将为其他类别的未来工作提供信息。

在结构化剪枝中,虽然从权重矩阵中识别出对最小性能降级最不重要的滤波器至关重要,但评估不同层对模型性能的重要性,并确定哪些层应剪裁更多或更少同样重要。这涉及学习不同层之间的稀疏分布。如图1所示,不同卷积和线性层对整体模型性能的重要性存在显著差异。此外,在一定程度的剪裁后,不同层之间的相对重要性也会发生变化。为了利用这一洞察,作者可以将更不重要的层的滤波器稀疏度分配较高,将更重要的层的稀疏度分配较低,并动态调整稀疏分布,以在剪裁过程中最小化性能下降。

picture.image在本文中,作者提出了一种新颖的方法 called RL-Pruner。与其他在稀疏训练期间学习层间稀疏分布的结构性剪枝方法不同,RL-Pruner 是一种在训练后进行结构性剪枝的方法,它利用强化学习采样来学习多个剪枝步骤的跨层最优稀疏分布。在每个步骤中,当前模型架构被定义为状态,而剪枝稀疏分布则作为策略。RL-Pruner 向策略分布中添加高斯噪声,以生成作为动作的实数剪枝稀疏分布,并生成下一个步骤的状态作为相应的压缩架构。每个世代都被视为一个采样过程。对于每个步骤,RL-Pruner 维护一个 replay buffer,其中存储剪枝稀疏分布动作及其相应的 Q 值,由奖励函数计算。奖励函数可以灵活定义,例如基于压缩模型的测试误差,测试误差和参数减少比例的组合,或根据特定使用场景的其它标准。在每次剪枝步骤后,RL-Pruner 更新稀疏分布策略,以学习最优稀疏分布。当计算资源允许时,在几次剪枝步骤后定期应用 post-training 阶段,以恢复剪枝造成的任何性能损失。在这些 post-training 阶段中,作者使用知识蒸馏,其中原模型作为教师,压缩模型作为学生。

RL-Pruner可以自动提取输入模型不同层之间的依赖关系,通过跟踪张量计算。目前,作者的方法支持几种流行的卷积神经网络(CNN)架构,包括残差连接、拼接连接和 Shortcut 。因此,作者的方法不需要特定的模型剪枝实现,可以自主执行剪枝,提高方法的通用性。

为了验证RL-Pruner的有效性,作者将所提出的算法应用于多种用于图像分类的CNN,包括 VGGNet(Simonyan和Zisserman,2014年),ResNet ,GoogLeNet 和MobileNet ,使用CIFAR-10和CIFAR-100数据集。根据实验结果,RL-Pruner在CIFAR-100上为VGG-19实现了60%的通道稀疏性,为GoogLeNet和MobileNetV3-Large在CIFAR-100上实现了40%的通道稀疏性,同时性能下降均小于1%。

2 Related Work

剪枝剪枝是神经网络压缩的主流方法之一,并在各个方面取得了显著的进展。正式地,给定一个包含个卷积层和个全连接层的神经网络,即,其中和分别表示第层卷积层的输出和输入通道数。每个层包含个滤波器,其中卷积层的形状为,全连接层的形状为。非结构剪枝去除单个权重,导致稀疏滤波器矩阵,除非有专门的硬件,否则不会导致推理速度提升。相比之下,结构剪枝在更大的粒度上工作,移除整个滤波器,并在标准硬件上实现压缩和真实加速。已经提出了几种剪枝标准,包括基于范数的方法,该方法剪除权重或范数最小的滤波器[17, 16],基于激活的方法,该方法使用激活图识别不重要的权重或滤波器,以及基于正则化的方法,该方法通过添加不同的正则化器学习结构稀疏网络。在本论文中,作者采用泰勒标准,该标准将权重滤波器的梯度(来自校准数据集)与其范数相加,以确定权重滤波器的重要性。

结构化剪枝通过移除滤波器并生成压缩子网络,可以被视为一种神经架构搜索(NAS)。由于模型中的不同层对整体性能的重要性不同,因此已有许多工作专注于自动分配层级的稀疏性。在这些工作中,开发强化学习 Agent 来自动选择适当的层级剪枝比例,无需进行手动敏感性分析。基于梯度的方法修改了梯度更新规则,使具有稀疏性约束的优化问题相对于权重是可微分的。进化方法使用进化算法来探索和寻找稀疏子网络。作者的方法主要使用强化学习来确定层级的最佳稀疏分布。

知识蒸馏通常会降低网络的性能,因此通常会在训练后进行后处理以恢复这种损失。知识蒸馏[13]将来自较大教师模型的知识传递到较小学生模型。基于响应的知识蒸馏允许学生模型通过在它们的预测之间添加蒸馏损失来模拟教师模型最终输出层的预测,从而将它们添加到优化函数中。另一方面,基于特征的知识蒸馏使学生模型能够从教师模型中学习中间层表示。在作者这种方法中,作者在恢复训练阶段采用基于响应的知识蒸馏,因为剪枝会改变中间层的维度,使得这些层之间的知识传递不那么直接。

3 Method

RL-Pruner首先在模型中的层之间构建一个依赖图,然后分几个步骤进行剪枝。在每个步骤中:

  1. 根据基础分布生成一个新的剪枝稀疏分布,作为策略;

  2. 每个层根据相应的稀疏度使用泰勒准则进行剪枝;

  3. 压缩模型被评估以获得奖励,动作和奖励被存储在 replay buffer 中。在每个步骤之后,根据 replay buffer 更新基础分布,如果计算资源充足,则使用知识蒸馏对压缩模型进行后训练,其中原始模型作为老师。图2说明了作者的方法。

picture.image### Build Dependency Graph

定义:给定卷积神经网络, 作者将卷积和线性层定义为可剪枝层, 同时,将激活层和池化层视为不可剪枝层。批量归一化层也被视为不可剪枝,因为它们只有在其前面的卷积或线性层被剪枝时才会被剪枝。正式地,作者通过,

picture.image

如果一个张量从 流向 ,且它们之间没有中间层,这意味着如果作者在 中剪裁特定输出通道 ,那么 中相应的输入通道 也必须剪裁。如果 和 之间没有依赖关系,作者有:

picture.image

作者旨在构建所有层之间的依赖图 ,并记录它们之间的通道索引映射关系。

为了实现这一目标,作者将一个示例输入图像输入到模型中,并提取所有层()及其输入张量()和输出张量(),其中表示所有层的数量。

然后作者通过匹配每个层的输入和输出张量来构建依赖图。当一个层的输出张量()直接作为另一个层的输入()时,会发生基本依赖,这意味着基本依赖。

picture.image

作者首先遍历和,以检查是否存在任何两个张量匹配,从而检测所有基本依赖关系。接下来,作者检查是否有未使用的输入张量,例如,这可能是由于张量修改(例如,压缩)或层之间的特殊连接(例如,残差或连接)导致的。在卷积和线性层之间通常会使用压缩。为了检测张量压缩的使用情况,作者还需要将所有输出张量压缩,并将其与未使用的输入张量进行比较,以构建依赖关系。作者将检测到的输出张量表示为,它表示在压缩之前的状态。为了确定压缩后的映射关系,作者需要检查产生的前一层输出区域。对于压缩,作者需要检查产生的前一层输出区域。

picture.image

该结果表明,在中,输出通道必须与从到的输入通道同时剪裁。为了检测 ConCat 连接,作者遍历,检查是否存在,使得

picture.image

可以通过将 沿着 进行截取,使得 来检测。然后,作者可以构建依赖关系。

picture.image

为了检测残差连接,作者遍历{T^out},检查是否存在T_b^out,T_c^out,使得

picture.image

然后作者可以构建依赖关系。

picture.image

这也适用于挤压与激发模块,其中

picture.image

对于残差连接和Squeeze-and-Excitation模块,来自层和的张量要么相加,要么相乘,然后传递给。这意味着剪枝的输出通道也需要同时剪枝的输出通道,反之亦然。因此,和必须同时剪枝。在构建依赖图后,作者将需要同时剪枝的层划分为互不相同的集合,确保同一集合内的层一起剪枝。

Assigning Layer-wise Sparsity

在自动提取层依赖性和相应的通道索引映射后,作者仍需确定如何将稀疏性分配到每个剪枝步骤中的层以最小化剪枝导致的性能下降,同时实现预期的整体模型稀疏性。用剪枝分布表示可剪枝层的稀疏性,其中作者为第个可剪枝层分配稀疏度,其中为目标整体模型稀疏性。由于作者只能以整数数量剪枝输出通道,因此寻找最优的最优化问题不是可微分的。

为了应对这个问题,作者采用了一种结合强化学习中的Q学习算法的蒙特卡罗采样策略[16]。具体来说,蒙特卡罗采样用于探索不同的剪枝分布,而Q学习则更新由基础剪枝分布表示的政策。因此,每个剪枝步骤包括多个采样阶段和剪枝阶段。

作者将每个压缩模型的架构定义为状态,将策略(PD)定义为动作。在每次采样中,作者将高斯噪声向量 添加到 PD 以生成真实剪枝动作 :

picture.image

where 是高斯噪声的方差,控制着在 周围的探索体积。较大的 表示更大的探索,允许剪枝策略探索更广泛的潜在架构。用 表示当前的模型架构,然后,受到贝尔曼方程(Bellman, 1966)的启发,作者评估每个 的 Q 值以及相应的压缩模型的架构 :picture.image

γ代表贴现因子,是在每个时间步的采样数量。为了实现采样效率,在每个采样阶段,作者采样个时间步以近似方程13中的期望:

picture.image

奖励函数由压缩模型的架构决定,通过以下方式进行建模:picture.image

其中 分别表示输入模型的测试准确性、FLOPs 压缩比和参数数量压缩比。超参数 可以根据具体目标设置为优先加速推理或减少参数数量。

在每个采样阶段之后,作者将每个动作 及其对应的 Q 值 存储在 replay buffer 中,用 中最低的 Q 值替换该条目。如果最低 Q 值超过 ,则对应的采样数据被丢弃。

然后,作者从 replay buffer 中选择一个动作 来更新策略 。为了平衡探索和利用,作者使用一个 -greedy 策略:以概率 ,作者从 中随机选择 。否则,作者选择 Q 值最高的动作:

picture.image

随着模型架构最初更冗余且通过剪枝变得更紧凑,作者在初始剪枝步骤中使用较高的探索值,而在后续步骤中逐渐降低值。考虑了几种衰减策略,包括常数、线性、余弦衰减。然后,通过步长调整策略分布以指向所选动作分布。为确保更新过程的稳定性,作者采用了近端策略优化的思想(PPO)[16],它限制策略分布的变化:

picture.image

在每一层的每个元素的改变比例都限制在范围内。因此,在经过几次采样阶段后,策略将收敛到各层的最优剪枝分布。

Layer Pruning

现在作者已经有了每个层的依赖图和每个层的稀疏度,下一个问题是如何在每个层中选择要剪裁的输出通道,以最小化性能下降。对于每个卷积层 和线性层 ,其中包含总共 个输出通道,每个对应一个权重矩阵。作者使用从训练数据集中提取的一个子集来帮助确定哪个权重矩阵更重要。根据[11],层 中第 个权重矩阵 的权重值 如下所示:

picture.image

由于计算Hessian矩阵需要O(N^2)的计算资源,作者省略了它以加速剪枝过程。同时,作者忽略了O(|W_i|^3),因为其值通常比第一项小得多。因此,估计的权重结构重要性变为:

picture.image

在评估每个输出通道的重要性后,作者对它们进行排序,并根据指定的稀疏性剪除最不重要的通道。

Post-train with Knowledge Distillation

鉴于模型经过多次剪枝以达到所需的整体稀疏性,剪枝必然会降低性能,因此在每个剪枝步骤后定期应用后训练以恢复失去的性能是有益的。知识蒸馏[10]有助于从更先进的教师模型向较小的或压缩的学生模型转移知识。

为了保留压缩模型的性能,作者使用原始模型作为教师,相应的压缩模型作为学生。具体来说,在后训练期间,作者将一个惩罚项引入损失函数,该惩罚项衡量教师模型和学生在输入图像上的概率分布之间的距离。修改后的损失函数为:

picture.image

因此,压缩模型能够从原始模型那里获得更多的知识,有助于保持其性能。

4 Experiment

Settings

作者从头使用PyTorch构建了一个名为RL-Pruner的AI模型。作者的实验在配备24GB内存的NVIDIA 4090D GPU和Intel(R) Xeon(R) Platinum 8481C CPU的系统上进行。

对于实现作者的方法,作者使用默认的超参数值:噪声方差 ,策略更新步长 ,折扣因子 ,采样步骤 ,每个采样阶段的样本数 ,每个剪枝步骤包含10个采样阶段。采用近似策略优化,。作者在实验中使用了三种奖励策略:基于准确度的、基于FLOPs的、基于参数的,分别对应方程17中的 、 和 。对于使用泰勒方法剪枝每一层,作者从训练数据集中提取100个样本作为校准数据集。使用教师模型和学生模型之间的损失系数 进行快速后训练。此外,如果在剪枝过程中发现具有更高准确度的压缩模型,作者将切换教师模型。对于探索策略,作者初始化探索参数 ,它在剪枝步骤的前10%以余弦方式减小。4.3节提供了超参数分析的详细信息。

在本研究中,作者关注分类任务,尽管作者提出的算法可以轻松地扩展到其他任务,例如分割。为了证明RL-Pruner对CNNs的广泛适用性,作者在几个广泛使用的分类模型上评估它,包括VGG-19 [29],ResNet-56 [17],GoogLeNet [20],DenseNet121 [16],以及MobileNetV3-Large [15]。这些模型涵盖了多种架构特征,如残差连接、 ConCat 连接和Squeeze-and-Excitation模块。

对于数据集,作者使用CIFAR-10和CIFAR-100 [18]来评估性能。浮点运算(FLOPs)确定计算复杂度和推理速度,而参数数量反映了推理过程中的内存使用。在作者的实验中,作者主要使用三个标准来评估压缩模型:测试Top-1准确率,FLOPs压缩比和参数数量压缩比。这些比例分别计算为和。

为了验证作者的方法的有效性,作者将它与其他结构化剪枝技术进行了比较,包括DepGraph [14]和GReg [14]。此外,作者还与其他神经架构搜索方法进行了比较,例如GNN-RL [21]。

作者在三个通道稀疏度比例(25%,50%,75%)下评估所有方法的表现,以评估它们在不同程度稀疏度下的性能。

Results

性能作者首先在CIFAR-100流行的CNN架构上评估作者的方法。如表1所示,作者可以将VGG-19压缩到60%的通道稀疏度,GoogleNet和MobileNetV3-Large压缩到40%的通道稀疏度,性能降低小于1%。这表明在这些架构中存在大量的冗余可以修剪,而不会造成显著的性能损失,因为在这些稀疏度下,FLOPs和参数数量都有了显著的减少。

作者还注意到,ResNet-56的性能在修剪过程中会急剧下降。这可能是由于ResNet-56具有许多层,相对参数较少,导致每个层的通道数较少。因此,在修剪过程中,作者的算法学习的最优稀疏分布很容易被修剪为零,从而限制了其表示每个层相对重要性的能力,导致性能下降。

picture.image结构化剪枝方法比较基于前人的工作[13],作者将作者的方法与其他结构化剪枝和神经架构搜索方法进行比较,使用了两种流行的架构:VGG-19在CIFAR-100上,ResNet-56在CIFAR-10上。如表2所示,在作者的方法在所有三个测试的通道稀疏度下,VGG-19的性能更高。对于VGG-19架构,作者的方法认为线性层不那么重要,并将其分配更多的通道稀疏度,从而在75%稀疏度下实现了显著的FLOPs压缩比。

请注意,GNN-RL方法在作者的重现中不支持在VGG-19的75%稀疏度下剪枝。对于ResNet-56,作者的方法在25%稀疏度下实现了更高的性能,同时具有更大的FLOPs和参数压缩比。然而,作者观察到在50%和75%稀疏度下的结果不理想,这可能是因为ResNet-56中每个层通道数较少。

picture.image### More Analysis

在本节中,作者探讨了在剪枝过程中采用的各种方法的效果,例如奖励策略和探索(ε)。为了评估这些方法,作者以较简单的VGG-19架构为基础进行实验,以最小化复杂架构引入的错误。

首先,作者评估了不同的奖励策略如何影响剪枝结果。作者在VGG-19上跟踪了75个剪枝步骤后,具体层级的参数奖励策略倾向于在卷积层中剪裁更多的滤波器,而准确度奖励策略则倾向于在线性层中剪裁更多的神经元。确定奖励策略与哪些层剪裁更多之间的关系具有确定性原则是具有挑战性的。相反,作者的方法基于给定的奖励策略确定最优架构。如图3所示,参数奖励策略倾向于在卷积层中剪裁更多的滤波器,而准确度奖励策略则倾向于在线性层中剪裁更多的神经元。

picture.image探索ε选择作者评估了不同的ε初始值,并在修剪步骤的前10%中使用各种衰减策略将ε减少到0,在CIFAR-100上的VGG-19模型上进行评估。如表3所示,采用探索策略可以提高性能,通过探索更广泛的架构范围。然而,它也可能偶尔收敛到一个次优的架构,导致性能降低。

picture.image相较于Scratch训练

接下来作者验证了作者的方法的去重效果。具体来说,作者从零开始训练一个压缩版本的VGG-19模型,该模型的架构与压缩模型相同,并比较其性能与压缩模型的性能。

如表4所示,从压缩架构的零开始训练模型并未达到与压缩模型相同的性能。此外,随着通道稀疏度增长,性能差距也在增大。这表明剪枝与后训练不仅可以加速推理并减少参数数量,还可以指导模型更高效地学习特征,从而在神经网络中实现更高的知识密度,这是仅通过训练无法实现的。随着稀疏度的增加,这种现象变得更加明显。

picture.image由于作者的方法通过跟踪和比较张量来构建依赖图,因此对于复杂架构(如ResNet和DenseNet)构建依赖图需要几分钟,而对于简单架构(如VGG)则不到一分钟。

考虑到作者的方法涉及数十个剪枝步骤,每个步骤都包含多个采样阶段,因此在作者的实验设置下,整个剪枝过程需要几个小时。通过并行化采样过程,可以加速采样过程,因为在采样阶段内的每个采样操作都是独立的。

5 Conclusion

在本文中,作者提出了RL-Pruner,这是一种结构化剪枝方法,它学习跨层的最佳稀疏度分布,并支持无需模型特定修改的一般剪枝。

作者希望作者的方法,认识到每一层对模型性能的相对重要性不同,将影响未来神经网络压缩领域的工作,包括非结构化剪枝和量化。

参考文献

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

点击上方卡片,关注 「AI视界引擎」 公众号

0
0
0
0
关于作者
关于作者

文章

0

获赞

0

收藏

0

相关资源
字节跳动 XR 技术的探索与实践
火山引擎开发者社区技术大讲堂第二期邀请到了火山引擎 XR 技术负责人和火山引擎创作 CV 技术负责人,为大家分享字节跳动积累的前沿视觉技术及内外部的应用实践,揭秘现代炫酷的视觉效果背后的技术实现。
相关产品
评论
未登录
看完啦,登录分享一下感受吧~
暂无评论