改進YOLOv4的蠶豆苗檢測算法及TensorRT加速

楊 肖,袁銳波,李兆旭,曹志鵬,林紅剛,朱 正

(昆明理工大學 機電工程學院,昆明 650504)

農田除草是農業種植中為了保證產量的一個必不可少的環節。目前,智能化農業機器人的研究中基于機器視覺的感知技術是熱點之一[1]。智能化除草機器人的要求是能夠沿作物行自主導航并避障,在行走的同時識別并定位作物,然后控制末端執行機構去除雜草。

除草機器人的檢測作業不僅僅是要求算法能夠準確識別雜草與作物,還對算法的實時性和魯棒性有較高要求。很顯然,以傳統圖像處理算法為基礎的目標檢測算法在檢測速度和抗干擾方面很難滿足以上要求[2]。

基于卷積神經網絡的目標檢測算法主要包括:以R-CNN[3]系列為代表的兩階段檢測算法和以YOLO[4]系列和SSD[5]系列為代表的單階段檢測算法。單階段檢測算法速度一般比雙階段檢測算法算法更快,但精度有所損失,在工業上得到廣泛應用[6]。雙階段的檢測算法通常精度較高,但是檢測速度較慢。修春波等[7]使用Faster R-CNN、YOLO v3、YOLO v4對馬鈴薯葉芽進行檢測,平均精度分別達到82.72%、87.22%、92.75%,檢測時間分別為1.739、0.429、0.519 s,由于模型參數太多,檢測速度還是達不到實際應用需求。YOLOv4[8]相較于其他算法更有優勢,這讓它成為一種非常高效和強大的目標檢測器[9]。

卷積神經網絡中不同的網絡結構和卷積操作對網絡參數量有很大影響。一般常用的主干網絡有VGG[10]、ResNet[11],以上網絡層數多、參數量大,適合在PC平臺下使用GPU進行訓練。疏雅麗等[12]使用VGG16,Resnet50,Resnet101對雜草進行識別,準確率分別達到85.49%、90.12%、90.83%,模型參數分別為YOLOv4的2倍、0.35倍、0.7倍。近年來,人工智能算法也開始向輕量化發展,比較常用的有MobileNet[13]、GhostNet[14],這些網絡以檢測精度的小幅下降為代價,使得網絡參數縮減。

為了解決嵌入式平臺計算力受限,檢測速度較慢,難以滿足實際應用需求的問題,本文提出一種輕量化網絡結構模型MobileNet-YOLOv4,使用MobileNet網絡代替YOLOv4原主干網CSPDarknet53,并用深度可分離卷積替換骨干網絡之后加強特征提取網絡和預測層的普通卷積,減少網絡參數量;然后利用NVIDIA的推理加速器對網絡結構進行重構和優化,提高GPU運行效率,在嵌入式平臺上實現模型的推理與加速。

1 方法與材料

1.1 分析方法

考慮到檢測模型的大小和部署設備的硬件條件,本文實驗方案是在PC端對改進YOLOv4網絡進行訓練,然后將訓練完成模型部署到嵌入式平臺Jetson Nano,利用TensorRT加速,在嵌入式平臺完成預測。

1.1.1YOLOv4

YOLOv4主體包括:一個主干,2個附加模塊和檢測頭的結構,其中主干是CSPDarknet53,2個附加模塊分別為SPP和PANet,檢測頭是YOLO Head。經過主干網絡提取后的特征會分別進入2個附加模塊進行處理,得到的數據會進行內部的合并最后在送入檢測頭。YOLOv4結構如圖1所示。

圖1 YOLOv4檢測算法結構框圖

1.1.2輕量級卷積神經網絡

Mobilenet網絡是Google提出的一種用于移動和嵌入式設備的輕量級神經網絡。特點是引入了深度可分離卷積,即用深度卷積(Depthwise)和1×1卷積(逐點卷積Pointwise)來替代標準卷積。普通卷積如圖2所示,標準的卷積層將輸入一個DF×DF×M的特征圖,并產生一個DF×DF×N的特征圖,當特征圖通過DK×DK×M×N的卷積核時,整個標準卷積的計算量如式(1):

(1)

圖2 普通卷積

深度可分離卷積由兩層組成:深度卷積和點向卷積。深度可分離卷積過程如圖3所示,使用深度卷積對每個輸入通道應用單個濾波器,計算量為M×DK2×DF2,然后還需要一個1×1的卷積層計算深度卷積輸出的線性組合,計算量為N×M×DF2。因此,總的計算量為兩步之和,表示為:

(2)

由此可以得到深度可分離卷積與標準卷積的計算量之比如式(3):

(3)

一般DK=3,且經過層層卷積計算后輸出的通道數N會非常大,這時它的計算量約為標準卷積的1/8到1/9倍,只是精度略有降低。

1.1.3改進YOLOv4網絡

本文提出的MobileNet-YOLOv4算法,使用MobileNet網絡代替YOLOv4原主干網絡CSPDarknet53,網絡特征層位置如圖4所示,提取的3個特征層尺寸分別為64×64、32×32和16×16,同時,將骨干網絡之后的加強特征提取網絡和預測層的普通3×3卷積和5×5卷積用深度可分離卷積替換,改進后的網絡檢測精度有所下降,但是網絡參數減少,檢測速度也有提升。改進的網絡整體結構如圖5所示,其中SeparableConv代表深度可分離卷積代替普通卷積的操作。

圖3 深度可分離卷積

圖4 網絡特征層位置框圖

從圖5可看出,檢測圖片輸入輸入網絡;主干網絡MobileNet提取圖像特征;后續加強特征提取網絡;最后分類檢測頭Head對提取的特征進行分類然后經過篩選后輸出預測結果。

機器學習算法是通過對該目標函數的優化求解。在分類問題中,一般使用損失函數作為目標函數。損失函數衡量的是模型的預測值和實際值不同的程度,損失函數越好,通常模型的性能越好。

圖5 改進的網絡整體結構框圖

目標檢測中的損失函數分為3個部分,即預測框回歸損失、置信度損失和分類損失。回歸一般基于IoU,即預測框與真實框的交并比,計算方式如式(4)所示:

(4)

式中:B和G分別代表預測框與真實框。但IoU只是一個比值,無法反應預測框與真實框的形狀、位置誤差,也沒法優化沒有重疊的部分。針對這一問題,不斷有人提出DIoU、GIoU、CIoU等指標。其中CIoU不僅僅使得目標框的回歸更加穩定,還考慮到了預測框與真實寬之間的距離、尺度、長寬比的擬合程度。MobileNet-YOLOv4即采用了CIoULoss,其公式如式(5)所示:

CIoU=IoU-ρ2(B,G)/c2-αυ

(5)

式中:ρ2(B,G)表示B與G中心點之間的距離。c表示同時包含預測框與真實框所有區域的最小矩形框對角線長度。v表示B與G長寬比之間的距離;α表示權重系數。計算方式分別如式(6)、(7):

(6)

(7)

然后用1-CIoU即可得到預測框回歸損失LCIoU。置信度損失和分類損失都使用二元交叉熵來計算,具體的損失函數如式(8):

(8)

1.1.4TensorRT加速

TensorRT是NVIDIA開發的用于深度學習GPU加速的推理優化器,能提升深度學習算法在jetson系列邊緣計算設備上的推理速度。TensorRT 加速原理是通過網絡層及張量合并,低精度計算及精度校準,多流執行,動態張量內存等方式來生成推理引擎[15]。其可以看作不會改變原有基本結構的一個單向傳播神經網絡,后續直接使用該引擎進行預測即可。TensorRT是將輸入串流到文件中,形成一個優化的engine,在執行時可調取它做推斷。TensorRT加速流程如圖6所示。

圖6 TensorRT加速流程框圖

1.2 實驗環境

實驗基于谷歌的深度學習開源框架tensorflow,并基于keras內置的函數庫來搭建模型代碼。算法最終目標是在小型移動機器人平臺上部署,但受制于機器人端處理器運算能力的不足,網絡的訓練仍然要在PC端進行,然后將訓練好的網絡部署到嵌入式平臺進行預測。

實驗軟件環境配置(基于Windows平臺):tensorflow-gpu1.13.2、keras2.1.5、Python3.6。實驗硬件平臺處理器為八核AMD Ryzen 7- 4800H,顯卡為NVIDIA GeForce GTX 1650 ,該實驗平臺用來訓練更改主干網絡后的YOLOv4。

YOLOv4網絡另使用臺式主機訓練,GPU為NVIDIA GeForce GTX 3050 。訓練使用cuda架構進行加速,充分利用CPU和GPU各自的優點合理分配運算資源。對應版本選擇:cuda 10.1,cudnn 7.6.5。

嵌入式平臺為Jetson Nano B01,GPU為NVIDIA Maxwell 128cores,CPU為4核ARM A57 1.43 GHz。下載官方鏡像文件Jetpack4.6.1,燒錄系統后配置推理環境。顯卡驅動仍選擇cuda10.0來保持訓練與推理環境的一致性,TensorRT版本選擇6.0.1。

1.3 數據來源

實驗中使用的蠶豆植株圖像于2021年11月在湖北省監利市農場采集。在采集蠶豆植株樣本時,圖片要包含足夠多的特征,算法最終會在自然條件下運行并受到光照強度變化的干擾。考慮到上述情況,本次采樣時分3個時間段進行拍攝,分別在早晨、中午和陰天條件下模擬機器人行走路線和攝像頭角度拍攝,得到1 100張、每張包含4～6株作物的大像素圖像(分辨率為4 160×3 120)作為原始樣本。后續對這些圖像進行篩選和裁剪,得到約1 200張分辨率2 080×1 560的蠶豆植株圖片。每張圖片中包含植株2～3株。后續實驗以這些圖片作為數據集進行訓練和測試。

實驗數據集使用labelimg進行人工標注,在原始數據上用矩形框框住目標植株并打上標簽broad bean。由于不識別雜草所以物體標簽只有一類。部分標注后的樣本如圖7所示。實驗從所有數據中按10%的比例抽出180張圖片作為測試集,用來評估檢測算法的泛化能力,剩下的 1 020張圖片用作訓練與驗證數據,以訓練網絡并調整各層的權值參數。以9∶1的比例將這些數據分為訓練集和驗證集。

雖然在采集樣本時已經充分考慮樣本和環境的多樣性,但實際農田中環境復雜多變,會出現許多實驗樣本無法涵蓋的情況。例如,因小車行走顛簸造成的拍攝角度變化、光照角度的變化等。為了進一步提高模型對復雜環境的泛化性與適應能力,實驗在網絡訓練之前先使用隨機的數據增強來擴充數據集。結合本實驗中的蠶豆數據集,背景單一、檢測目標單一且基本無小目標,故實驗中直接采用隨機裁剪、縮放、水平翻轉、色彩抖動、附加噪聲進行數據增強。

圖7 部分標注后的樣本圖像及對應標簽

2 實驗與分析

2.1 模型評價指標

損失函數的值只能表示模型訓練是否已經收斂,并不能直接用來判斷模型檢測效果的好壞。目標檢測模型常用的精度衡量指標主要有準確度、召回率、平均精度(AP)和平均精度均值(mAP),本實驗為了使算法在移動端也能有不錯的效果,采用了輕量級網絡模型,為了突出其整體參數量小、運算速度快的優點,除了精度外必須把檢測速度也作為非常重要的一組評價指標,本研究采用圖片的檢測時間作為評價標準。準確率和召回率的計算公式如式(9)、式(10):

(9)

(10)

式中:P、R分別表示準確度和召回率;TP、TN、FP、FN分別代表真正例、真負例、假正例和假負例。設置數據集的指標confidence(置信度)為0.7,即檢測結果的IoU>0.7時,認定樣本為真正例,反之為假正例。最后統計所有正例,再減去真正例,即得到了假負例。

根據所求P,R可以得到準確率與召回率的對應曲線,即R-P曲線。用積分的方式計算某一類目標的R-P曲線與坐標軸圍成的面積即為該類目標的AP,計算公式如式(11),對所有類別目標的AP求均值,即為整個數據集的mAP。本實驗中只檢測蠶豆植株這一類目標,即mAP與AP值完全相等。

(11)

2.2 改進YOLOv4網絡訓練

1) 主干網絡替換實驗

實驗數據為采集的蠶豆樣本1 200張,其中訓練與驗證圖片1 080和測試圖片120張。總體訓練策略為:① 導入預訓練權值文件;② 凍結主干網絡參數進行訓練(50個epoch);③ 解凍網絡,對整體網絡再進行一定輪次的訓練(100個epoch)。凍結主干網絡是因為主干提取的一般為淺層特征,具有通用性,而剩下的部分(分類器)更多是針對具體訓練目標的特征。所以先凍結主干訓練分類器,解凍后再對整個網絡進行微調。MobileNet-YOLOv4和GhostNet-YOLOv4網絡訓練損失及驗證損失如圖8所示。

圖8 模型訓練和驗證損失曲線

前50輪訓練是預訓練,用較高的學習率微調網絡。MobileNet-YOLOv4網絡的損失函數曲線如圖9。可以看出,模型損失在訓練50輪后下降速率減緩,90輪左右接近收斂。在93輪因為連續4輪val loss不下降而提前停止訓練。網絡在第89輪獲得最小的訓練損失3.64,在第83輪獲得最小的驗證損失2.10。GhostNet-YOLOv4網絡在50輪訓練后呈現波動中緩慢下降的趨勢,從70輪開始驗證損失趨于平緩,在90時已連續多輪不下降,根據earlystoping的設置在92輪提前停止網絡的訓練。驗證損失在第89輪得到最小值1.35,訓練損失在第90輪得到最小值2.02。不同的網絡不能直接通過比較損失函數的大小來評估檢測效果的好壞,需要通過其他參數具體評估。

2) 檢測效果對比

一般來說,損失值越小的輪次得到的檢測效果越好,根據損失函數曲線圖,在90輪附近使用不同的權重文件計算mAP(主干網絡為GhostNet時),一共測試116張圖片,目標數量(蠶豆植株數)為221。在第89輪得到了最優的效果,真正例和假正例分別為210和39。網絡的召回率、準確率、mAP參數如圖9所示。可以看出GhostNet-YOLOv4模型的mAP值在89輪訓練時已經達到了93.44%。同樣的,以MobileNet為主干時,在91輪得到最優的mAP值,真正例和假正例分別為206和46,其余參數如圖10所示。

圖9 GhostNet-YOLO模型評估參數

圖10 MobileNet-yolo模型評估參數

為了直觀比較MobileNet-YOLOv4、GhostNet-YOLOv4和YOLOv4三種網絡的檢測效果,以參數量、AP值、FPS為評估指標,具體參數如表1所示。

從表中可以看出YOLOv4擁有最高的AP值95.36%,但網絡體積約是2種改進網絡的6倍,且檢測速度也較低。在除草作業中不僅要保持作物檢測的位置精度,還需要足夠的檢測速度保證導航-識別-除草作業的實時進行。同時網絡體積過大,導致機器人嵌入式處理器有更高的性能要求,對小型機器人來說很困難,因此,不考慮使用該網絡作為檢測算法。GhostNet-YOLOv4對比YOLOv4,體積減小且AP值僅下降了0.92%,檢測速度也約為4倍。而MobileNet-YOLOv4對比GhostNet-YOLOv4,AP值下降了2.22%,但檢測速度有所提高,結合檢測速度和精度,MobileNet-YOLOv4檢測效果相對較好。

表1 3種模型性能對比

2.3 TensorRT加速

YOLOv4網絡精度只是略微高于優化網絡,且檢測速度較慢,部署實驗中使用優化后的2種模型做對比實驗。使用測試集的180張圖片作為測試對象,僅在Jetson nano 上測試算法在不同框架下的推理效果。具體數值如表2所示。

表2 優化算法在不同框架和數據精度下的檢測效果

由于使用INT8精度的數據推理使得精度下降超過預期范圍,這里只給出FP32和FP16兩種精度下的推理效果指標。總體來看:MobileNet-YOLOv4和GhostNet-YOLOv4在檢測速度和精度方面各有優勢,兩種模型在Tensorrt框架下使用單精度數據推理時,對比原keras下的精度僅是略有下降,下降值分別為1.09%和1.27%,但在檢測速度上都有約50%的巨大提升。而數據精度降低為半精度FP16時,MobileNet-YOLOv4模型的mAP值進一步下降為85.85%,但FPS提升為原框架下的2.3倍。GhostNet-YOLOv4的mAP上比原keras框架下下降了3.31%,FPS則提升至原本的2.17倍。

2種模型(圖11中,左圖為GhostNet-YOLOv4,右圖為MobileNet-YOLOv4)在除草機器人端處于不同光照環境下的實際檢測效果,如圖11所示。可以看到,2種模型在不同的光照條件下都能準確檢測到作物,基本滿足機器人端視覺模塊的需求。

圖11 模型在不同光照下的檢測圖像

3 結論

針對嵌入式平臺下深度學習網絡推理速度慢,性能低的問題,提出了一種基于改進YOLOv4的輕量級蠶豆苗檢測方法。利用NVIDIA的推論加速器TensorRT重構優化網絡結構,提高GPU的運行效率,在嵌入平臺上實現深度學習模型的推論加速。實驗結果顯示,改進網絡在Jetson Nano嵌入式平臺上,改進后的網絡模型推理速度達到20.3 FPS。提出的網絡模型能支持深度學習模型在嵌入式平臺的實時應用。從識別結果看,嵌入式平臺與PC端性能差異很大,模型在嵌入式平臺上的表現遠低于PC端。為了驗證模型實際工作性能,針對性能較弱的嵌入式平臺做模型的優化是必要的。