基于多尺度注意力機制和知識蒸餾的茶葉嫩芽分級方法

黃海松 陳星燃 韓正功 范青松 朱云偉 胡鵬飛

(1.貴州大學現代制造技術教育部重點實驗室， 貴陽 550025; 2.貴州裝備制造職業學院科研處， 貴陽 551400; 3.清鎮紅楓山韻茶場有限公司， 貴陽 551400)

0 引言

在茶葉生產過程中，茶葉嫩芽的分級從根本上決定了產品品質，目前主要借助人工感官評審法完成，所得結果易受人為主觀性及環境因素影響[1-2]。如今，計算機視覺技術雖不斷與農業工程學科交叉[3-4]，但針對茶葉嫩芽分級問題，國內外學者所開展的研究仍十分有限。許高建等[5]將采用VGG-16進行特征提取的Faster R-CNN模型在COCO數據集上進行預訓練后用于四等級茶葉嫩芽的分級問題，陳妙婷[6]考慮茶葉采摘時遮擋情況的干擾添加相應的圖像樣本，利用YOLO算法處理改進PSO-SVM提取的特征信息實現茶葉嫩芽分級，兩者雖均能在較短時間內實現茶葉分級，但準確率不高且所處理的數據需事先利用圖像工具進行標注，時間和人工成本較大。高震宇等[7]為實現茶葉嫩芽分級，基于多層卷積神經網絡搭建識別模型，并借助局部連接的訓練方式極大加快了收斂速度，但因較為簡單的網絡結構使其處理復雜特征的能力匱乏，在實際應用中性能難以保證。毛騰躍等[8]融合SVM與特殊角點檢測兩種方法對四等級茶葉嫩芽進行分級，分級準確率為94.24%，但多分類器會產生大量計算冗余，加劇設備計算負擔。吳正敏等[9]基于茶葉動態下落過程中所采集到的圖像數據，結合隨機森林算法和分類算法，利用前者預先選出具有強可分性的復雜特征，經分類算法處理后實現茶葉嫩芽的準確分級，但模型參數量龐大，將加劇設備的計算負擔。QIAN等[10]將預訓練后的Inception-V3模型的權重矩陣遷移至自建茶葉嫩芽數據集之上，并利用仿射變換對數據進行增強，最終對三等級茶葉的分級準確率達到了98%，但模型較大，不利于投入實際應用。XU等[11]提出了一種聯合電子鼻和計算機視覺技術的茶葉品質鑒定方法，采用KNN、SVM等完成識別模型的建模，并通過融合策略對兩者信號進行分析，綜合兩維度的信息實現精準分級，但使用成本較高且環境要求較為嚴苛。在實際應用中茶葉嫩芽圖像采集面臨著諸多困難，而上述分級方法普遍存在冗余參數多和模型規格大的缺點，在小規模數據集上易出現模型過擬合現象,嚴重影響分類性能，且對設備儲存和計算能力要求高，難以部署到生產現場中。

綜上，本文基于采摘自貴州紅楓山韻茶場的茶葉嫩芽自建圖像數據集，構建以多尺度注意力單元為核心的分級模型，聚焦強分級特征并從感知野和網絡縱深兩個層面充分提取茶葉嫩芽圖像中的多尺度特征信息；提出一種結合雙遷移學習和知識蒸餾的分級方法，在保證模型輕量高效的前提下，提升抵抗過擬合的能力和分級性能。

1 遷移學習和知識蒸餾

為解決現有方法在小規模茶葉嫩芽數據集上因模型過擬合現象嚴重影響分級性能，及模型所占內存空間大難以應用到實際生產中的問題，本文在構建一種輕量型分級網絡的基礎上，提出結合雙遷移學習(Transfer learning, TL)和知識蒸餾(Knowledge distillation, KD)的分級方法，進一步提升模型穩健性和分級性能。本文中TL從兩個源域數據中提取知識并將該部分知識應用于解決茶葉嫩芽分級任務之中。與多任務學習不同，TL會更加傾向于關注目標任務，而不是同時關注源任務和目標任務[12]。如圖1所示，TL可以在目標任務中高質量訓練數據較少的情況下，避免從頭開始學習每一項任務，充分利用預訓練時的源域信息，將從以往任務中習得的知識遷移至該目標任務之中。KD則是由多倫多大學的HINTON等[13]提出的一種知識遷移技術，借助混合損失函數的反向傳播將教師模型(Teacher module)所習得的先驗知識傳授給學生模型(Student module)，以達到壓縮教師模型和增強學生模型性能的目的。

圖1 學習過程示意圖Fig.1 Learning process diagrams

2 基于多尺度注意力機制和知識蒸餾的方法

2.1 卷積塊注意力模塊

卷積塊注意力模塊(Convolutional block attention module, CBAM)由WOO等[14]于2018年提出，其借助通道注意力模塊(Channel attention module, CAM)和空間注意力模塊(Spatial attention module, SAM)推理注意力特征圖并與輸入特征圖相乘，使模型提取出更具區分性的特征，并抑制無關特征，達到自適應優化特征和提高模型性能的目的。CBAM憑借輕量的特性，可以嵌入到任何卷積神經網絡(CNN)結構。

2.1.1通道注意力模塊

CAM主要關注特征矩陣中有價值的信息，提取分辨性較強的特征。通過對輸入特征進行最大池化(MaxPool)和平均池化(AvgPool)操作壓縮空間維度上的特征信息，生成兩種不同空間上的特征圖，經多層感知器(MLP)處理后逐元素相加，最后通過Sigmoid激活函數進行非線性轉換，其結構如圖2所示。

圖2 通道注意力模塊Fig.2 Convolutional block attention module

2.1.2空間注意力模塊

SAM主要關注輸入特征矩陣映射的位置信息，是對CAM的一種補充。輸入特征圖在經過沿空間維度的平均池化和最大池化操作后被拼接為一個通道數為2的三維空間特征圖，后經卷積操作及Sigmoid激活函數處理生成空間注意力特征圖，其結構如圖3所示。

圖3 空間注意力模塊Fig.3 Spatial attention module

2.2 多尺度卷積塊注意力模塊

為使模型更加關注茶葉嫩芽圖像的感興趣區域，聚焦有利于分級的重要特征信息并抑制無關特征，本文引入注意力機制對茶葉嫩芽圖像中的強分級特征信息進行提取與處理。由于茶葉嫩芽圖像采集較為困難，為在有限的數據集上盡可能地挖掘特征信息，本文提出了一種多尺度空間注意力模塊(Multiscale spatial attention module, MSAM)，并通過與CAM結合構建了多尺度卷積塊注意力模塊(Multiscale convolutional block attention module，MCBAM)，其結構如圖4所示。在MSAM中用卷積核尺寸為5×5和9×9所構成的多尺度特征提取卷積層替換原普通卷積層，對經池化操作后的特征進行處理，并利用1×1卷積層實現多尺度特征信息的整合，兩種不同的感知野較原始固定大小的卷積層更有利于模型提取茶葉圖像中的多尺度特征信息[15]。對MCBAM而言，輸入特征經CAM處理得到加權結果后再經MSAM加權得到輸出特征圖。

圖4 MCBAM注意力模塊Fig.4 MCBAM attention module

2.3 模型網絡設計

ShuffleNet-V2網絡為MA等于2018年提出的一種輕量型網絡，引入了通道拆分(Channel split)、通道清洗(Channel shuffle)以及用保持通道數不變的分支拼接代替逐元素相加等操作，完全符合高效網絡結構設計應該遵守的G1～G4準則，在計算復雜度和精度上有優異的綜合表現。

結合ShuffleNet-V2 0.5x基本單元、MCBAM、多尺度深度捷徑(Multiscale depth shortcut, MDS)構建Shuffle多尺度注意力單元(Shuffle multiscale attention unit, SMAU)，如圖5所示。在ShuffleNet-V2 0.5x基本單元的3組卷積層之后以串聯的方式嵌入MCBAM，對經卷積層處理后的特征信息進行不同大小感知野下多尺度信息的提取及強分級復雜特征信息的增強。通過MDS(圖5中①②③)將不同網絡縱深下的特征信息進行融合，避免信息流傳遞過程中有效特征信息的損失，提高分級性能。且MDS可充當梯度反向流動的高速通道，使梯度不受阻礙的流動，緩解模型退化的問題，增強模型的穩健性[16]，并在MCBAM、卷積層和整個SMAU之間構造成恒等映射，實現MCBAM與卷積層間的自適應組合[17]，避免因引入MCBAM而引起精度下降。最終使模型聚焦茶葉嫩芽圖像中有利于分級的特征，并從感知野大小和網絡縱深兩個層面提取多尺度特征信息用于分級，充分利用有限的數據，強化模型處理小規模數據集的能力。基于ShuffleNet-V2 0.5x的網絡結構，以SMAU為核心構建一種融合多尺度注意力機制的輕量型茶葉分級模型ShuffleNet-V2 0.5x-SMAU，結構見表1。

圖5 多尺度注意力單元Fig.5 Shuffle multi-scale attention unit

表1 ShuffleNet-V2 0.5x-SMAU網絡結構Tab.1 Structure of ShuffleNet-V2 0.5x-SMAU

2.4 結合雙遷移學習和知識蒸餾的茶葉嫩芽分級方法

本文引入遷移學習以充分利用源域的豐富特征信息，ShuffleNet-V2 0.5x-SMAU通過在大規模數據集上的預訓練可以增強其在茶葉嫩芽數據集上抵抗模型過擬合的能力，并提高分類性能。結合知識蒸餾技術，提出一種結合雙遷移學習和知識蒸餾的茶葉嫩芽分級方法，將預訓練后的教師模型所習得的暗知識蒸餾到網絡結構更簡單、參數量更少的ShuffleNet-V2 0.5x-SMAU中，以增強其分類性能和處理小規模數據集的能力。

在茶葉嫩芽分級領域，由于目前尚無公共數據集，故常將Oxford-17 flower和Oxford-102 flower數據集應用于模型的預訓練[10]，但綜合考慮計算量和時間成本等因素，基于Oxford-102 flower數據集中的50種花卉圖像數據建立flower-50數據集。在flower-50上通過引入Focal Loss損失函數對學生模型進行預訓練，可減少易分類樣本的權重，使模型更關注難分類的樣本[18]。

結合雙遷移學習和知識蒸餾的茶葉嫩芽分級方法，訓練出能實現不同等級新鮮茶葉嫩芽劃分的模型。將在ImageNet上預訓練后的Resnet32以及在flower-50上預訓練后的ShuffleNet-V2 0.5x-SMAU的模型參數矩陣遷移至HF-G3數據集上，以解決茶葉嫩芽分級問題，在利用該方法訓練模型時依次以二者構造教師模型及學生模型。利用該方法進行訓練的流程如圖6所示，設置溫度T為10，比例系數α為0.8。利用結合雙遷移學習和知識蒸餾的茶葉嫩芽分級方法進行模型訓練的步驟為：

圖6 基于雙遷移學習和知識蒸餾的茶葉嫩芽分級方法結構圖Fig.6 Structure diagram of tea bud grading method based on dual transfer learning and knowledge distillation

(1)以像素矩陣的形式將茶葉嫩芽圖像信息輸入到教師模型中得到對應各等級茶葉的概率分布，并通過逐次除以溫度T進行平滑處理，后經softmax函數得到教師模型的軟標簽(Soft labels)qSLi，計算公式為

(1)

式中xTi——教師模型最后一層輸出的特征圖

n——特征圖像通道數

(2)將茶葉嫩芽圖像信息輸入學生模型，并按步驟(1)中溫度進行平滑處理，經softmax函數得到學生模型的軟預測(Soft predictions)qSPi, 計算公式為

(2)

式中xSi——學生模型最后一層輸出的特征圖

(3) 將茶葉嫩芽圖像信息輸入學生模型，所獲概率分布直接由softmax函數處理獲取學生模型的硬預測(Hard predictions)qHPi，計算公式為

(3)

(4)由相對熵損失函數(Kullback-Leibler Divergence)對步驟(1)中獲得的軟標簽和步驟(2)中獲得的軟預測進行計算得到軟損失函數值LSOFT，計算公式為

(4)

式中B——批處理圖像數量

C——茶葉嫩芽等級數

qSPj——學生模型軟預測

(5)由交叉熵損失函數(Cross entropy)對步驟(3)中獲得的硬預測和真實標簽(True lables)進行計算得到硬損失函數LHARD，計算公式為

(5)

式中yTRUE——真實標簽

qHPj——學生模型硬預測

(6)最后通過比例系數α調整兩損失值的比例后得到混合損失函數LKD，計算公式為

LKD=(1-α)LHARD+αT2LSOFT

(6)

其中，LKD的反向傳播是教師模型將暗知識授予學生模型并大幅提高學生模型分類性能的關鍵。

3 實驗結果與分析

3.1 數據處理與實驗設置

本文所用茶葉嫩芽圖像采集平臺由中科微創ZW-C3600型工業相機、變焦鏡頭及LED環形補光燈等組成，如圖7所示。所拍攝茶葉均采摘于貴州省清鎮市紅楓山韻茶場，由茶場工人將原始茶葉嫩芽圖像樣本分為單芽、一芽一葉、一芽兩葉共3個等級，如圖8所示。選用翻轉、平移、旋轉及添加隨機噪聲4種數據增強方式對茶葉圖像進行擴容處理，如圖9所示。最終獲得茶葉嫩芽圖像數據共計2 136幅，其中對應3種不同等級的茶葉嫩芽各712幅，并按照比例3∶1將所得數據劃分為訓練集與測試集，作為本文實驗采用的HF-G3圖像數據集。

圖7 茶葉嫩芽圖像采集平臺Fig.7 Tea bud image collection platform1.中科微創ZW-C3600型工業相機 2.變焦鏡頭 3.LED環形補光燈

圖8 HF-G3數據集中3種不同等級茶葉的原始圖像Fig.8 Three original images of tea at different levels in HF-G3 dataset

圖9 經數據增強后的茶葉圖像Fig.9 Tea image after data enhancement

本文設置了消融實驗和對比實驗。在消融實驗中，以準確率變化曲線、茶葉嫩芽分級混淆矩陣等指標對多尺度注意力機制、結合雙遷移學習和知識蒸餾的分級方法的合理性與有效性進行驗證。在對比實驗中，則是面向其他現有技術，從模型抗過擬合能力、分級能力、分級均衡性及模型規格4個角度對本文方法在綜合性能上的優越性進行論證。

3.2 消融實驗結果與分析

為驗證本文方法的有效性，依次設置采用本文茶葉分級方法的實驗組A(ShuffleNet-V2 0.5x-SMAU+TL+KD)、基于傳統CBAM注意力模塊的實驗組B(ShuffleNet-V2 0.5x-CBAM+TL+KD)、學生模型未在flower-50上進行預訓練的實驗組C(ShuffleNet-V2 0.5x-SMAU+KD)以及未采用知識蒸餾的實驗組D(ShuffleNet-V2 0.5x-SMAU+TL)。在HF-G3測試集數據上驗證4個實驗組的分級性能，通過Origin2019b繪制各實驗組對應的準確率變化曲線，如圖10所示。為進一步分析針對不同等級茶葉嫩芽進行分級時性能的差異，繪制各實驗組對應的混淆矩陣，如圖11所示。各級茶葉分級準確率及平均準確率如表2所示。

圖10 各消融實驗組在測試集上的準確率變化曲線Fig.10 Accuracy rate change curves of each ablation experimental group on test set

圖11 各消融實驗組在測試集上的混淆矩陣Fig.11 Confusion matrices of each ablation experiment group on test set

表2 各消融實驗組在測試集上的分級準確率Tab.2 Classification accuracy of each ablation experiment group on test set %

從圖10可知，實驗組A與實驗組B準確率分別為95.88%、95.13%，表現出了近似的分類性能，但后者對應的變化曲線震蕩嚴重，說明SMAU相較于傳統CBAM注意力模塊，可以向模型提供更豐富的多尺度特征，并借助MDS構建恒等映射，有效緩解模型過擬合現象。對比實驗組A、C，實驗組A的分類性能明顯優于實驗組C，且實驗組C出現了明顯的過擬合現象，說明實驗組A中通過學生模型在flower-50上的預訓練，能在提高模型分類性能的同時，增強模型抵抗過擬合的能力。對比實驗組A、D，前者因以預訓練后的Resnet32作為教師模型進行知識蒸餾，較實驗組D的最高分級準確率提高了3.87個百分點，并借助教師模型所習得的龐大權重矩陣在一定程度上緩解了過擬合現象，該組對比驗證了知識蒸餾的有效性。

由表2可知，由于一級茶葉(單芽)區分特征明顯，各實驗組對其進行劃分的準確率都十分理想，而在對另外兩等級茶葉(一芽一葉、一芽兩葉)的劃分中實驗組A的準確率最高，分別為92.70%、89.89%。對比實驗組A、B，實驗組A在一芽一葉及一芽兩葉分級中，準確率較實驗組B分別提高了3.94、1.13個百分點，說明SMAU更能使模型聚焦到利于分級的多尺度特征。對比實驗組A、C、D，在對一芽兩葉的識別準確率上，實驗組A比另外兩組高10.68、9.56個百分點，說明基于雙遷移學習和知識蒸餾的茶葉等級分級方法能使模型處理茶葉復雜特征的能力得到顯著提升。

在此部分所設置的消融實驗中，本文方法對教師模型與學生模型均進行預訓練，并將所獲取到的權重參數矩陣遷移至茶葉嫩芽分級問題上，借助梯度反向傳播將教師模型所習得暗知識向學生模型進行傳授，結果表明該方法極大增強了學生模型獲取和分析茶葉復雜特征的能力，使分類性能取得大幅提升。且通過相關實驗的設置亦證明了SMAU較傳統CBAM在茶葉等級劃分問題上的優越性。

3.3 對比實驗結果與分析

針對茶葉分級問題，為衡量本文ShuffleNet-V20.5x-SMAU+TL+KD與其他方法間性能的差異，本部分設置實驗從抗過擬合能力、分級能力、分級均衡性、模型規格4方面依次對本文模型與FI-DenseNet[19]、Alter-Second Model[15]、ResNet18-Green tea[20]、AlexNet-Camelia[21]在HF-G3測試集上的表現進行對比。利用Origin2019b繪制各對比實驗組的準確率與損失值變化曲線，如圖12所示。基于圖13所示的混淆矩陣，計算不同等級茶葉嫩芽的分級準確率及平均準確率，并繪制分級準確率柱狀圖，如圖14所示。為從模型規格的角度深度探究對比模型間差異，以計算量、參數量、模型規格為指標制定如圖15所示的模型比較柱狀圖。

圖12 各對比模型在測試集上的準確率及損失值變化曲線Fig.12 Accuracy and loss curves of each comparison model on test set

圖13 各對比模型在測試集上的混淆矩陣Fig.13 Confusion matrices of each comparison model on test set

圖14 各對比模型在測試集上的分級準確率Fig.14 Classification accuracy of each comparison model on test set

圖15 各對比模型比較Fig.15 Comparison charts of each comparison model

從茶葉嫩芽分級準確率及損失值變化曲線看，本文的初始準確率最高。與FI-DenseNet及ResNet18-Green tea進行比較，兩者分別基于DenseNet169及ResNet18搭建，最高分級準確率分別為91.57%、95.51%，與本文模型最高準確率94.94%具有相似的分類性能，但進一步觀察變化曲線，可看出后兩者震蕩劇烈，出現了所有模型中最為嚴重的過擬合現象，說明基于復雜網絡模型的方法并不適用于解決小規模茶葉等級劃分問題。與Alter-Second Model對比，本文模型中的SMAU使模型不只限于感知野下的多尺度特征，還能挖掘不同深度下的多尺度特征，同時在注意力機制作用下，對有利于分級的特征進行增強，在同樣擁有極強抵抗過擬合能力的條件下最高準確率較Alter-Second Model上升了14.34個百分點，收斂時損失值自0.415 1下降至0.150 7。AlexNet-Camelia采用了遷移學習的方法，表現出了一定的抵抗模型過擬合的能力，損失值曲線收斂迅速，震蕩微弱，但準確率曲線仍存在一定幅度的震蕩，本文模型不同于AlexNet-Camelia只是簡單的結合AlexNet和遷移學習技術進行特征提取和分析，本文模型是將預訓練后的ShuffleNet-V2 0.5x-SMAU和ResNet32模型參數遷移至茶葉分級中，并依次作為知識蒸餾中的學生和教師模型，這種方式不但使學生模型的穩健性得到大幅提升，也通過暗知識的傳授，使分級性能得到很大程度的增強，使其最高準確率較AlexNet-Camelia上升了2.06個百分點，收斂時損失值自0.626 7下降至0.150 7。

從對不同等級茶葉嫩芽的分級準確率及平均準確率看，本文憑借MCBAM與MDS，從不同感知野和網絡深度聚焦多尺度特征信息，且通過結合遷移學習與知識蒸餾，充分利用在源域中所得的權重矩陣，在避免模型過擬合的情況下兼具了優異的茶葉分級能力，對一芽一葉的劃分準確率達到了92.70%，較其他方法分別提高了6.18、34.27、5.62、6.74個百分點。在一芽兩葉分級問題上，ResNet18-Green tea雖憑借其復雜的網絡結構表現出了最高的分類準確率，但本文模型的分級精度較其余對比實驗組仍上升顯著，且在平均準確率上本文模型的準確率為94.19%，優于其余模型的分級均衡性。

從模型規格看, 除Alter-Second Model以外，本文模型無論是在計算量、參數量還是模型內存占用量上都較其余模型小一個量級，表現出了最高的效率和最佳的輕量性。進一步觀察，Alter-Second Model利用深度可分離卷積在很大程度上減小了模型的規格，使其計算量和參數量在與本文模型保持在同一量級。AlexNet-Camelia雖然擁有較少的計算量，但因采用微調的遷移學習方式凍結了大量的權重，使模型仍有龐大的參數量和計算冗余。本文模型與ResNet18-Green tea比較，計算量、參數量及內存占用量3項指標僅為后者的10.17%、16.4%、16.73%，由此得知本文模型在輕量化上遠優于ResNet18-Green tea。

綜上所述，本文提出的茶葉嫩芽分級方法，能在確保模型輕量化的前提下從感知野和網絡深度兩個層面對茶葉圖像的多尺度特征信息進行聚焦，并借助結合雙遷移學習和知識蒸餾的方法，使模型在處理小規模茶葉嫩芽數據集時，不僅擁有極強的抗過擬合能力和分類性能，還同時擁有所有對比模型中最少的計算冗余參數和最小的模型規格，對設備要求最低，更易部署到茶葉嫩芽分級實際場景中，展現出了所有對比模型中最佳的綜合能力。

4 結束語

本文將MCBAM嵌入ShuffleNet-V2 0.5x網絡的基本單元中，并通過引入MDS構建了一種輕量型茶葉嫩芽分級模型ShuffleNet-V2 0.5x-SMAU。在利用本文方法進行茶葉嫩芽分級時，將在兩個大型數據集上預訓練后的ResNet32、ShuffleNet-V2 0.5x-SMAU依次作為教師模型與學生模型。在自建茶葉嫩芽圖像數據集HF-G3上設置了2組實驗：在消融實驗中，依次對SMAU、本文方法的合理性及優越性進行了充分論證；在與FI-DenseNet、Alter-Second Model、ResNet18-Green tea、AlexNet-Camelia對比實驗中，本文方法對單芽、一芽一葉、一芽兩葉的分級準確率為100%、92.70%、89.89%，平均準確率為94.19%，計算量為2.219 6×109，參數量為1.826 9×106，模型內存占用量為7.13 MB，在抗過擬合能力、分級能力、分級均衡性、模型規格上表現出了最優異的綜合性能。