面向垃圾圖像分類的改進注意力機制剪枝算法*

高文曉

北方民族大學電氣信息工程學院 寧夏 銀川 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 結束語

為使神經網絡模型更便于在邊緣設備上部署和應用，本文提出了一種基于改進注意力機制的剪枝算法。基于該混合注意力模塊的剪枝算法能夠在較少準確率損失的前提下對網絡進行有效壓縮。本文剪枝方法可以進一步與量化等方法結合進一步壓縮模型。