基于Swin Transformer的家居垃圾分類系統

瞿定垚，王學,

(1.河北工程大學 數理科學與工程學院，河北邯鄲，056038; 2.河北省計算光學成像與光電檢測技術創新中心，河北邯鄲，056038)

0 引言

中國是全世界人口最多的國家，在日常生活的過程中，會制造大量垃圾。在由生態環境部所公布的《2019年全國大、中城市固體廢物污染環境防治年報》中顯示，在2018年里，在全國中的200個大、中城市里生活垃圾產生量為21147 3萬噸，處置量為21028 9萬噸，處置率達到了99 4%[1]。雖然由數據顯示的處置率己經很高了，只有其中少部分未被處理。但是這并不是全國所有城市的數據，除了生態環境部統計的這200個大、中城市，還有許多城市和落后的中小城鎮的數據并未被統計到，可以推斷出未被處理的生活垃圾仍有很多。在年報中，2013~2018年中國200個大、中城市生活垃圾產生量統計及增長情況，可以得出生活垃圾產生量以5-10%的增長速率逐年遞增[2]。在垃圾產量巨大的同時，處置垃圾的方法仍然十分落后。

我國解決垃圾分類主要依靠人工二次分揀，其弊端如下：

（1）二次分揀模式的參與主體是分類指導員，居民其實還是沒有真正地參與進來，所以它解決不了居民參與率低的問題[3]；要讓居民真正地參與垃圾分類，從根本上還是要讓他們從源頭開始做好分類。

（2）二次分揀的對象是廚余垃圾，像可回收垃圾、有害垃圾等其他垃圾，分類指導員們都不分揀，所以沒有真正地達到生活垃圾的減量化、資源化、無害化[4]。

（3）此模式操作簡單，方法原始，效率低，安全性差，更談不上精細化管理運營。因此靠“二次分揀”支撐的垃圾分類是走不長遠的。

使用計算機視覺技術實現自動的垃圾分類識別，存在的問題是效率低，回收利用效率降低，安全性差[5]。精度差，價格昂貴。

針對居民參與率低的問題，我國應該從垃圾產生源頭入手，以家居垃圾為研究環境，走出二次分揀的困境，使全民參與進來，切實解決分類問題。針對垃圾數據集數據少、質量較差、背景單一、存在標注錯誤的情況，神經網絡無法全面、準確分析、理解分類特征的問題，構建大規模垃圾圖像數據庫是重要途徑。同時數據量的大幅增加與計算機硬件性能的提升，使得深度學習技術開始在計算機視覺領域大放異彩。由于其速度快、準確率高的優點，使得其在各個領域都取得了一定的成果。圖像分類在深度學習領域主要依靠CNN卷積神經網絡和視覺Transformer圖卷積神經網絡，由于家居環境對分類精度要求高，對推理速度要求較低，使用單圖像分類而不是目標檢測。

1 系統組成

家居垃圾分類系統主要由三部分組成，構建垃圾分類數據集、設計垃圾分類網絡和搭建安卓app。

首先，本文構建垃圾分類數據集。數據集來源為：網絡爬蟲技術爬取生活中常見的各種家居垃圾圖片，將不符合要求的垃圾圖片剔除[6]；使用華為垃圾數據集中垃圾圖片，選擇符合要求的垃圾圖片數據；通過攝像設備收集邯鄲市河北工程大學生活中的常見家居垃圾圖片。通過數據增強方式，旋轉，鏡像，對比度，進一步擴大數據集。其次，本文設計垃圾分類網絡。比較CV領域中的CNN經典圖像分類模型的優缺點，設計實驗選出適用于本課題的圖像分類模型[7][8][9]。通過NLP領域對CV領域的遷移，搭建了適用于本課題的圖像分類模型。通過遷移學習[10]，數據增強和添加全局注意力的方式，對網絡進行了優化。最后，搭建了安卓端口app,對訓練好了的模型成功進行了轉換和量化，將自己訓練好了的深度學習模型移植到了手機安卓app上。垃圾分類系統框架圖如圖1所示。

圖1 垃圾分類系統框架圖

■ 1.1 Swin transformer簡介

Transformer在NLP領域應用十分廣泛，但是在CV領域的應用存在很多困難，這源自于兩類任務的本質區別，例如CV里物體尺寸變化很大，模型需要處理不同尺度的同類物體，而在NLP領域，把word tokens作為基本要素，并不存在以上問題。同時圖片分辨率高，自注意力如果逐像素計算成本過高[11]。

Swin Transformer 是一種新型的transformer，可以用作視覺和語言處理的統一模型， 引出了一種具有層級的特征表達方式和具有線性的計算復雜度。Swin Transformer有兩個特征： (1) 層級特征結構；(2) 線性復雜度。層級特征結構使得本模型適用于 FPN 或者 U-Net 模型；線性復雜度是因為使用local window self-attention[12]。這些特點使得本模型可以作為一種通用模型用于各種視覺任務。在目標檢測、圖像分割等視覺任務中，Swin Transformer 均得到了SOTA 的結果。

Swin Transformer在上游任務取得了極好的成績，可拿來可直接做下游任務，并且效果很好。本文選擇了swin_base_patch4_window12_384_in22k作為預訓練模型，在保持高精度的同時推理速度也適用于家居垃圾分類任務，并在后續的任務中加速了推理速度。

■ 1.2 訓練流程

首先輸入像素384×384圖片到Patch Partition模塊中進行分塊，即每4×4相鄰的像素為一個Patch，然后在channel方向展平（flatten）。本文輸入的是RGB三通道圖片，每個patch就有4×4=16個像素，然后每個像素有R、G、B三個值，展平后是16×3=48，所以通過Patch Partition后圖像shape由 [H, W, 3]變成了 [H/4, W/4, 48]。然后在通過Linear Embeding層對每個像素的channel數據做線性變換，由48變成C，即圖像shape再由 [H/4, W/4, 48]變成了 [H/4, W/4, C]。其實在源碼中Patch Partition和Linear Embeding就是直接通過一個卷積層實現的。

然后就是通過四個Stage構建不同大小的特征圖，除了Stage1中先通過一個Linear Embeding層外，剩下三個stage都是先通過一個Patch Merging層進行下采樣，然后都是重復堆疊Swin transformer Block。Swin transformer 網絡結構如圖2所示。

圖2 Swin transformer網絡結構

這里的Block有兩種結構，如圖3所示，這兩種結構的不同之處僅在于一個使用了W-MSA結構，一個使用了SW-MSA結構[13]。而且這兩個結構是成對使用的，先使用一個W-MSA結構再使用一個SW-MSA結構，堆疊Swin Transformer Block的次數都是偶數。

圖3 Swin transformer 兩種blocks結構

最后對于圖像分類網絡，后面還會接上一個Layer Norm層、全局池化層以及全連接層得到最終輸出。

2 模型優化與算法改進

■ 2.1 模型優化

2.1.1 數據增強

數據增強：通過旋轉，鏡像，對比度的方法擴充了約54000張圖片用于訓練，對擴充后的數據集圖像數據進行預處理，增加模型的魯棒性和模型泛化能力[14]。

2.1.2 遷移學習

微調 Fine-tuning:將在大數據集上訓練得到的weights作為特定任務（小數據集）的初始化權重，重新訓練該網絡（根據需要，修改全連接層輸出）[15]。訓練的方式可以是：(1)微調所有層，(2)固定網絡前面幾層權重，只微調網絡的后面幾層，這樣做有兩個原因：A 避免因數據量小造成過擬合現象；B CNN前幾層的特征中包含更多的一般特征（比如，邊緣信息，色彩信息等），這對許多任務來說是非常有用的，但是CNN后面幾層的特征學習注重高層特征，也就是語義特征，這是針對于數據集而言的，不同的數據集后面幾層學習的語義特征也是完全不同的[16]，本文選擇訓練方式2。

網絡微調是一項非常強大的技術，無需從頭開始訓練整個網絡，“快速啟動”學習，最終得到更高精度的遷移學習模型，同時節省了計算資源[17]。

■ 2.2 模型改進

在swin transformer基礎上加入了NLP領域的全局注意力Global Attention[18~20]。Global Attention的目的：Gobal Attention的目的是為了在生成上下文向量，用來處理文本信息的,單詞和句子存在聯系，局部像素點也和全圖存在聯系[21]。transformer是用來處理文本信息的，現在可以用來處理CV領域的圖像信息。swin transformer在CV領域起到了突破性的重大作用，選擇在swin transformer的head前 邊插 入Global Attention，經過訓練，精度得到了提升。全局注意力如圖4所示。

圖4 全局注意力

改進實驗分為三組，插入一層Gobal Attention，命名模型為Swin transformer-A；插入兩層Gobal Attention，命名模型為Swin transformer-B；插入三層Gobal Attention命名為Swin transformer-C。使用基于微調的遷移學習，固定部分層數的權重，插入一層Gobal Attention模型訓練層數如下代碼。

“training layers 4 norm1 weight

training layers 4 norm1 bias

training layers 4 attn qkv weight

training layers 4 attn qkv bias

training layers 4 attn proj weight

training layers 4 attn proj bias

training layers 4 norm2 weight

training layers 4 norm2 bias

training layers 4 mlp fc1 weight

training layers 4 mlp fc1 bias

training layers 4 mlp fc2 weight

training layers 4 mlp fc2 bias

training head weight

training head bias”

插入兩層和三層就是將代碼中的layers 4進行堆疊。

3 垃圾分類實驗

■ 3.1 獲取數據集

實驗整理的原始數據集共20022張，約13000張來自華為云比賽數據集，7022張自己收集標注。提前做好四分類，O代表other waste，H代表harm garbage，K代表kitchen waste，R 代表recyclable garbage，小類別共計40類。途徑來自于拍攝，網絡爬蟲爬取，網上下載以及截取等方式。將其標注好，放在同一類別的文件夾中，小類別共計40類。后來考慮到部分種類垃圾數量較少，使用數據增強的手段，通過鏡像，旋轉，對比度三種方式，擴充了數據集約54000張用于訓練。由于旋轉角度有的圖像會缺失信息，所以對擴充的數據進行清洗，然后重新歸類，初始數據集類別和數量如圖5所示。

圖5 初始數據集類別以及種類

■3.2 實驗結果與分析

將家居垃圾數據集按照9:1的比例分為訓練集和驗證集，18035張用于訓練，1987張用于驗證。對訓練集使用數據增強后，72035張用于訓練，驗證數量不變。使用預訓練 權 重 swin_base_patch4_window12_384_in22k pth，模型準確率如表1所示。

表1 不同模型準確率

9 ConvNeXt-Base+遷移學習+數據增強 有 64 4 100 97 2%97 2%10 VIsion tranfformer-Base+遷移學習+數據增強 有 64 4 100 99 2%96 8%

通過該表，得出遷移學習和數據增強和全局注意力機制都能有效提升模型的精度。消融實驗可以看出，在數據量小的時候，全局注意力提升的效果較好，使用數據增強后，仍然可以提升模型精度，提升幅度會變小。說明全局注意力和數據增強以及遷移學習改善網絡性能有部分重合了，組合在一起時效果最好。實驗8，9，10是和前四個實驗進行對比，在同量級參數下看模型效果，Efficientnetv2選擇m模型[22]，pytorch架構下無預訓練權重，使用工具將Tensorflow下ImageNet-21K預訓練權重轉換為Pytorch架構下可以使用；ConNeXt選 擇 convnext_base_22k_224 pth 預 訓 練權重[23]；vision transformer選擇vit_base_patch16_224_in21k pth作為預訓練權重[24]。

垃圾40分類最優模型混淆矩陣如圖6所示。

圖6 垃圾40分類混淆矩陣

通過對每個類別的分析，誤差主要集中在可回收物瓶罐這些物品中，實際生活中只需要一級分類標準就可以了。使用一級標準，重新訓練4分類。選擇上邊最優模型和策略，72084張用于訓練，2001張用于驗證，Train acc99 8%,Val acc99 2%，再次驗證模型性能很強。

垃圾40分類每個類別的精確率，召回率，特異度如表2所示。

表2 每個類別的精確率、召回率、特異度

K_Fruit pulp 0 95 0 983 0 998 K_Leaves edible orange roots 0 964 1 0 0 998 K_Tea leaves 1 0 1 0 1 0 K_The shell 1 0 0 977 1 0 K_leftovers 0 936 0 957 0 998 O_Broken flowerpots and dishes 0 978 0 978 0 999 O_Cigarette butts 0 977 0 977 0 999 O_Disposable snack box 0 977 0 955 0 999 O_Stained plastic 0 978 0 938 0 999 O_chopstick 0 976 0 976 0 999 O_toothpick 0 9 1 0 0 999 R_Charging treasure 1 0 0 98 1 0 R_Cosmetics bottles 0 98 0 98 0 999 R_Courier bags 1 0 1 0 1 0 R_Edible oil drum 0 957 1 0 0 999 R_Leather shoes 1 0 0 98 1 0 R_Metal food can 0 974 0 949 0 999 R_Old clothes 0 977 0 956 0 999 R_Plastic bowl tub 0 986 0 986 0 999 R_Plastic hangers 1 0 0 98 1 0 R_Plastic toys 1 0 0 973 1 0 R_Plug wire 1 0 0 976 1 0 R_Plush toys 0 987 0 987 0 999 R_Shampoo bottle 1 0 0 976 1 0 R_Spice bottles 1 0 0 911 1 0 R_The bottle 0 929 1 0 0 998 R_The glass 0 966 0 983 0 999 R_cans 0 959 0 959 0 999 R_carton 1 0 1 0 1 0 R_cutting board 1 0 1 0 1 0 R_drink bottle 0 947 0 947 0 999 R_package 0 948 1 0 0 998 R_pan 1 0 1 0 1 0 R_pillow 0 915 0 956 0 998

40垃圾分類最優模型訓練過程如圖7所示。

通過遷移學習，從曲線中可以看出模型收斂的很快，損失下降的很快并且到達了較低值，驗證集準確率很快得到了最優解。訓練集損失穩定下降，訓練集準確率穩定提升。在數據集復雜度更高的情況下，已經超過了華為2019年垃圾分類挑戰杯的方案，40垃圾分類的準確率達到了97 63%，超過了它的96 96%，并且參數量只占他的14，flops只占12。

Python垃圾圖片測試結果如圖8所示，推理結果全部正確。

圖8 python可視化測試結果

4 應用開發

■ 4.1 模型轉換

4.1.1 模型轉換預處理

Pytorch的模型文件一般會保存為 pth文件，C++接口一般讀取的是 pt文件，因此，C++在調用Pytorch訓練好的模型文件的時候就需要進行一個轉換，轉換為 pt文件，才能夠讀取。

預處理如下：

在轉換的時候，首先就需要先將模型文件讀取進來，然后利用pytorch提供的函數torch jit trace進行轉換。PyTorch 一個算子可能會導出ONNX很多的膠水算子，計算圖變復雜，導致推理效率下降。使用onnx simply庫對onnx文件簡化，使網絡結構更加精簡，消除細碎op。

4.1.2 模型轉換方式以及結果

PyTorch -> ONNX -> 第三方后端（主流方式），第三方后端有 TensorRT，MNN，NCNN等，這樣的優勢是不同的后端在目標平臺都有相對應的優化，可以獲得最快的推理速度，但問題是第三方后端對onnx算子的支持優化，會有運行不了的問題[25]。

方案：Pytorch -> MNN：pytorch -> onnx -> MNN。

結果：模型轉換成功，python和C++調用結果一致。

■ 4.2 模型量化

轉化成mnn模型雖然可以進行推理，模型文件較大或者運行較慢的情況特別是在移動設備等邊緣設備上，算力和儲存空間受限因此壓縮模型是一個急需的工作，本文所做的工作為離線量化[26,27]。

原模型轉換后數據類型為Fp32,由于模型太大370MB，使用轉換工具做了量化處理，壓縮模型的大小，同時可以加速模型的推理速度。模型量化如表3所示。

表3 模型量化

使用python和c++對同一張圖片進行結果測試，輸出一致，說明模型轉換和模型量化沒有問題。Fp16量化和Int8量化后的模型可以正常推理，并且可以顯示出置信度。Int2量化后的模型圖像推理時候無任何反應，結構壓縮太多，出現了精度失真的情況。最終量化選擇Fp16量化和Int8量化。

■ 4.3 app 開發

(1)模型準備，留下以下兩個模型model-87_384_sim_fp16和model-87_384_sim_int8。

(2)創建工程。創建Android工程，打開Android Studio新建一個Native C++工程，配置好JDK 、SDK 、NDK。

(3)模型導入。將步驟(1)中生成的兩個文件復制到步驟(2)中的新工程中，存放的目錄為“F:swintransformerappsrcmainassets”。這兩個模型都可以使用，在后續的“MainActivity java”可以自由選擇切換。

(4)添加標簽文件。在“app”目錄下創建名為“F:swintransformerappsrcmainjni”的文件，將本文需要分類的類別在這個“swimtransformer_jni cpp”文件中，共40類。編寫MNN進行模型推理，根據自己的模型，編寫jni調用MNN進行模型推理

(5)模型分類。在“app”目錄下創建名為“swintrans former java”的文件，在里面添加模型分類的代碼，主要包含了以下功能：目標圖像讀取、目標圖像預處理、模型加載、模型前向推理、模型預測結果。在圖像的預處理中，需要對圖像進行歸一化處理，同時統一縮放為 384×384 大小。

(6) 界面設計和操作邏輯設計。在“app”目錄下創建名為“MainActivity java”的文件，寫入調用相冊、文件管理器、瀏覽器的函數以及一些操作的邏輯。創建“ xml”的文件，編輯 App 的界面顯示，完成 App的開發[28,29]。

■ 4.4 app 測試

4.4.1 app 界面以及操作方式

圖形界面主要分為三個按鈕，選圖、CPU識別、GPU識別。進入界面后，下方會彈出“Copy model to data”和“Copy model to success”。之后就可以從相冊選圖或者文件管理或者瀏覽器選擇圖片進行識別。

4.4.2 app 單張圖片測試分析

使用單張圖片測試模型量化和模型轉換有無推理精度損失，選出最優搭配下的app。具體測試方案和結果如下：選取Mix2s、Note 10 pro、Redmi K40 在不同量化模型下，分別進行同一張圖像用cpu 和gpu識別，評價指標分為準確率和推理速度。測試原圖與python測試圖如圖9所示，方便與表4進行結果對比。

圖9 測試原圖與python測試可視化

python推理圖9(a)花費時間400ms，結果如圖9(b)，使用設備 GPU2080TI。

App具體測試圖9(a)結果如表4所示。

表4 APP測試圖a結果

結論：

（1）在單一量化時，硬件設備越好，推理速度越快，精度越高。

（2） int8量化效果比Fp16量化效果更好，用同一款手機測試同一張圖片時候，精度和推理速度都得到部分提升；相比原模型Fp32數據推理同一張圖片時，無精度損失。

（3） 使用Mnn加速優化后，用同一款手機測試同一張圖片時候，推理速度得到了再一次的提升，但是精度對于不同的機型表現略有不同，但總體趨勢是下降的。

（4）在不同安卓版本和不同分辨率下下的不同機型，該App都能完美適用。手機移動邊緣設備與顯卡算力差距很大，手機推理速度經過優化滿足日常操作時長需要。

4.4.3 app 可視化

本次app可視化測試選擇機型Mix2s，選擇模型為Int8量化，加速推理使用MNN優化，圖形測試選擇三張圖片，分別使用cpu和gpu推理。評價指標分為準確率和推理速度，圖像推理會顯示出前TOP5的標簽和概率，標簽按照概率大小依次排列。可視化結果如圖10所示。

App測試可視化推理結果全部正確，推理速度滿足日常使用。

5 結論

通過對各種深度學習圖像算法進行研究，加上了大量的訓練策略，并且對網絡進行了改進，最終得到超過現有技術的精度。在保持高精度的前提下，盡量提升推理速度，對模型進行量化，使之可以移植到安卓嵌入式設備的智能垃圾分類應用。通過模型轉換，使其用手機的cpu和gpu推理，徹底擺脫通信的束縛，在深度學習的多種部署方式中具有高度便攜性 、通用性 、經濟適用性性的特點。除此之外，使用MNN優化加速了推理速度。由于構建并使用一個龐大的數據集進行訓練，模型的泛化性能較好，適用于生活中的各種場景。