面向垃圾图像分类的改进注意力机制剪枝算法*

高文晓

北方民族大学电气信息工程学院 宁夏 银川 750021

引言

正确进行垃圾分类和分类处理可以促进资源再利用，避免垃圾处理造成二次危害。传统的垃圾分类主要依靠人工分拣，需要耗费大量人力物力，而且不利于人们的身体健康。随着深度学习技术的不断进步与发展，借助计算机视觉进行垃圾分类的研究得到了广泛研究与应用。但是由于神经网络复杂的结构和庞大的参数量，使模型在边缘设备上的部署和应用变得十分困难。研究表明，神经网络的内部存在大量的冗余参数，去除这些参数可以对网络进行精简，在保证识别精度的基础上降低模型的参数和计算量，从而实现在嵌入式设备上的部署[1]，对网络进行压缩成为研究热点。

模型剪枝具有原理简单、易实现、压缩效果显著的特点，被广泛应用于模型压缩。剪枝主要分为非结构化剪枝和结构化剪枝，其中非结构化剪枝的模型需要专门的算法或硬件结构才能实现加速[2]；结构化剪枝是对通道、卷积核等结构进行剪枝，不依靠特定的软硬件平台，更有利于嵌入式设备的部署，目前得到了广泛应用。

如何准确评价网络结构的重要程度是结构化剪枝过程中的主要问题。Li等人提出将滤波器的L1范数作为衡量标准[3]；Hassibi等人利用误差函数的二阶导数去对权重的重要性进行衡量[4]；Wen等人将不同粒度的剪枝作为正项添加进损失函数中进行学习[5]；Hu等人提出利用激活函数的输出值对神经元重要性进行评价[6]；Molchanov等人提一种基于泰勒展开式的通道评价标准[7]；He等人提出了一种基于Lasso回归和最小重构误差的方法衡量滤波器的重要性程度[8]。

本文受注意力机制对数据的重要部分进行增强，不重要部分进行减弱的启发，提出一种基于注意力机制的剪枝算法。可以从空间和通道两个维度对特征图进行标定，同时生成每个通道的权值，该值代表着通道在分类任务中的重要性程度，进而根据该值指导模型剪枝。

1 改进的注意力机制

本文提出的混合注意力模块包含空间注意力模块和通道注意力模块两部分。

图1 改进的空间注意力模块结构图

将与分别与输入特征图F逐元素相乘实现对特征图行和列方向的重构得到和。然后如（3）所示将和逐元素相加得到包含两个方向注意力的特征图。

图2是改进的通道注意力模块结构图，使用全局平均池化和全局最大池化得到两个一维向量和，同时使用这两种池化方式可以有效提高通道注意力的评价能力。使用1×k的一维卷积对每个通道的k个近邻进行交互得到两个通道注意力特征向量和，然后将这两个向量进行逐元素相加和sigmoid运算得到通道注意力向量。改进的通道注意力部分计算过程如（4）所示。

图2 改进的通道注意力模块结构图

2 基于改进注意力机制的剪枝算法

经过混合注意力模块不仅能够增强卷积神经网络的性能，同时也能生成同一卷积层中不同通道的重要性评价向量。考虑到不同输入经过混合注意力模块生成的通道描述向量不同，为了避免个别样本造成的偏差这里对不同输入图像产生的通道描述向量进行逐元素求均值，得到一个汇集不同输入特征信息的通道评价向量，从而得到更准确的各个通道的权重值。具体每一层的通道重要性得分计算算法如表1所示。对所有卷积层i层包含的全部j个通道基于第d张图像生成的注意力权重进行累加，然后除以D得到平均值，计算过程如（5）所示。即通过注意力机制生成的每一个通道的重要性评价得分，根据该值可以对通道进行重要性排序和去除。

其中，i∈{1,2,…,l}, j∈{1,2,…,C}, d∈{1,2,…,D}, C代表通道数量，D代表数据集的图像数量。

考虑到不同的卷积层所提取的信息特征不同，如果设置统一的阈值，可能会使某一层通道剪枝过度造成较大的精度损失，所以本文对每一层设置相同的剪枝比例p。根据剪枝比例p、每层通道数目和通道重要性排序结果得到各层的阈值，然后根据该阈值进行剪枝。

首先将混合注意力模块插入预剪枝的网络中需要剪枝卷积层的后面，对网络进行微调及推理得到每层各个通道的重要性排序；其次，根据预设剪枝率p对网络逐层剪枝，并移除对应层插入的混合注意力模块，每完成一层剪枝对网络进行一次微调；最后，完成全部剪枝后对网络进行整体微调提高精度。

3 实验结果与分析

本文模型训练和剪枝实验在Ubuntu 18.04系统下进行，CPU为Intel(R) Xeon(R) E-2236 ，GPU为NVIDIA Tesla P100。

3.1 数据集构建

本文实验数据集为TrashNet和自行构建的生活垃圾数据集。

TrashNet数据集包含6类，共计2507张图片。本文对公开数据集TrashNet和garbage_classify进行筛选和补充，如表1所示选取了30种常见垃圾构建了生活垃圾数据集。本文构建的生活垃圾数据集总计38000张图片，将数据集进行划分得到训练集样本26600张，验证集样本7600张，测试集样本3800张。

表1 生活垃圾数据集构成

3.2 实验结果

3.2.1 在TrashNet数据集上的实验。在公开垃圾数据集TrashNet上对ResNet50和VGG16网络进行不同剪枝率及不同剪枝方法的对比实验。实验结果如图3所示，图案面积大小和标注分别代表剪枝后模型的FLOPs（M）和参数量（M）。随着剪枝率增加，模型FLOPs和参数量减小的同时准确率缓慢降，整体来说基于混合注意力机制指导下的剪枝方法带来的准确率损失较小。本文方法与基于通道L1范数剪枝算法（L1nrom） ，网络瘦身剪枝算法（slimming pruning）进行实验对比。对于ResNet50，压缩率为0.3时，本文方法准确率要比基于通道L1范数剪枝算法和网络瘦身剪枝算法分别高0.72%和0.03%；当压缩率为0.5时，本文方法准确率要比基于通道L1范数剪枝算法和网络瘦身剪枝算法分别高1.18%和1.03%。本文方法整体表现要优于其他两种方法。

图3 ResNet50和VGG16在TrashNet数据集上的剪枝实验结果

3.2.2 在自建垃圾数据集上的剪枝实验。在本文构建的垃圾数据集上对ResNet50和VGG16网络进行训练和剪枝实验。如图4所示ResNet50和VGG16基准模型准确率分别为93.75%和91.98%。压缩率为0.3时，ResNet50和VGG16的准确率分别为93.31%和91.73%，参数量非别为14.92M和96.2M。压缩率为0.5时，ResNet50和VGG16的准确率分别为92.35%和91.12%，参数量分别为9.98M和71.97M。结果显示本文方法对ResNet50和VGG16模型可以进行有效的压缩同时仅损失较少的准确率。

图4 VGG16与ResNet50在本文生活垃圾数据集上的剪枝实验结果对比

4 结束语

为使神经网络模型更便于在边缘设备上部署和应用，本文提出了一种基于改进注意力机制的剪枝算法。基于该混合注意力模块的剪枝算法能够在较少准确率损失的前提下对网络进行有效压缩。本文剪枝方法可以进一步与量化等方法结合进一步压缩模型。