垃圾分揀項目找上門的時候我心里其實是有點打鼓的——場景看著簡單不就是幾個垃圾桶里的瓶瓶罐罐嘛真做起來才知道要在傳送帶上一秒抓一個目標(biāo)還要兼顧準(zhǔn)確率和部署成本光靠通用檢測模型是行不通的。當(dāng)時團隊里有同事提了一嘴用YOLOv5 4.0版本來做目標(biāo)檢測我抱著試一試的心態(tài)跑通了全套流程從環(huán)境配置到數(shù)據(jù)集制作再到模型訓(xùn)練最后落地成一條可用的分揀業(yè)務(wù)流。這篇博客就把這整條鏈路攤開講環(huán)境怎么配、數(shù)據(jù)怎么標(biāo)、模型怎么調(diào)、坑在哪里全部按實操順序來適合剛接手檢測項目、尤其是想在工業(yè)場景里做視覺分揀的同學(xué)參考。我用的硬件是普通工作站加一塊中端顯卡沒有上集群也沒有買高配服務(wù)器YOLOv5 4.0這套東西對硬件的要求其實比想象中低但環(huán)境搭不對照樣跑不起來。先別急著下載代碼下面從環(huán)境配置開始把每一步拆清楚。1. 環(huán)境配置裸機到能跑通的完整清單以及為什么4.0版本這么選1.1 Python環(huán)境與系統(tǒng)版本選擇YOLOv5 4.0版本發(fā)布于2021年中段它的依賴鎖得不算死但也不是隨便拿個Python版本就能跑。我實測比較穩(wěn)的組合是Python 3.8.10或3.9.7這兩個版本在Windows和Linux下都能正常編譯torch擴展。Python版本別選太高3.10及以上在部分機器上裝torch1.9會碰到二進制兼容問題得不償失。Windows下我建議用Anaconda創(chuàng)建獨立虛擬環(huán)境Linux下用conda或venv都行。我的習(xí)慣是創(chuàng)建虛擬環(huán)境conda create -n yolov5 python3.8激活環(huán)境conda activate yolov5確認pip版本pip install --upgrade pip這里有個很多人會忽略的細節(jié)YOLOv5 4.0的requirements.txt里torch和torchvision的版本是帶“”的比如torch1.7.0。如果直接無腦安裝最新版torch碰上CUDA版本不匹配的概率很高。我當(dāng)時踩過一個坑——裝完最新torch后模型訓(xùn)練時loss直接nan查了半天是顯卡驅(qū)動和CUDA庫不匹配。后來老老實實按官網(wǎng)推薦的匹配表來一次性跑通。1.2 CUDA、cudnn與torch版本匹配這一步是環(huán)境配置里最核心、也最容易炸的一環(huán)。YOLOv5 4.0本身不做CUDA強綁定但GPU訓(xùn)練必須用正確的驅(qū)動。我的經(jīng)驗是顯卡驅(qū)動版本決定你能支持到哪一版CUDA Toolkit命令nvidia-smi看右上角的CUDA Version這個數(shù)字是驅(qū)動支持的最高版本不代表你的環(huán)境已經(jīng)裝了對應(yīng)CUDA。實際訓(xùn)練時用的CUDA是torch自帶的不需要單獨安裝完整的CUDA Toolkit但你需要確保驅(qū)動夠新。cudnn也建議裝一下雖然torch可以自動加載部分算子但我遇到過推理速度忽高忽低的情況后來發(fā)現(xiàn)是cudnn未配置導(dǎo)致卷積算法未被優(yōu)化。我這邊的穩(wěn)定組合是驅(qū)動470.xx CUDA 11.1 cudnn 8.0.5 torch 1.9.0 torchvision 0.10.0。裝torch用官方匹配索引pip install torch1.9.0cu111 torchvision0.10.0cu111 -f https://download.pytorch.org/whl/torch_stable.html裝完后驗證一下python -c import torch; print(torch.cuda.is_available())顯示True就說明GPU路線通了。這一步值得多花半小時確認因為后面所有訓(xùn)練時間都依賴這個基礎(chǔ)。1.3 下載YOLOv5 4.0源碼與環(huán)境依賴YOLOv5 4.0的代碼結(jié)構(gòu)相比后來的5.0、6.0版本要簡潔不少特別適合學(xué)習(xí)或者二次改造。官方倉庫的tag切換成v4.0也可以直接下載zip包解壓。解壓后目錄里最顯眼的就是train.py、detect.py、models、data、utils這些模塊。接下來安裝依賴pip install -r requirements.txtrequirements里包含numpy、opencv-python、matplotlib、pyyaml、tqdm這些常用庫還有pycocotools用于COCO指標(biāo)計算。如果是在Windows上pycocotools安裝失敗的概率比較大官方推薦裝pycocotools-windows。我試過兩次裝Windows版可以避免編譯報錯安裝命令是pip install pycocotools-windows依賴裝完后可以立刻用預(yù)訓(xùn)練權(quán)重跑一次demo驗證全鏈路是否通。做法是下載yolov5s.pt放進weights目錄然后執(zhí)行python detect.py --source data/images/bus.jpg --weights weights/yolov5s.pt如果能看到圖像窗口彈出并繪制出檢測框就說明環(huán)境徹底通了。這一步我建議別跳過因為很多問題會在第一次推理時就暴露。1.4 垃圾分揀場景為什么選用4.0版本這個選擇值得多說一句。4.0版本的網(wǎng)絡(luò)結(jié)構(gòu)以CSPDarknet53為主配合PANet特征融合層它的anchor設(shè)置和損失函數(shù)已經(jīng)相當(dāng)成熟訓(xùn)練小目標(biāo)比如礦泉水瓶的瓶蓋效果也能接受。相對于更新的6.0版本4.0的代碼結(jié)構(gòu)更直白修改anchor或自定義dataloader的難度更低對于垃圾分揀這種類別不算多、背景相對固定的場景完全夠用。垃圾分揀業(yè)務(wù)的特殊性在于一是類別少但形態(tài)差異大比如塑料瓶和紙箱的寬高比就完全不同二是場景里會出現(xiàn)密集堆疊的情況目標(biāo)之間互相遮擋對NMS閾值和anchor尺寸特別敏感。4.0版本的自適應(yīng)anchor計算功能和多尺度訓(xùn)練恰好能處理這部分問題。而且4.0版本推理權(quán)重文件更小后續(xù)導(dǎo)出ONNX或TensorRT的兼容性也更好。2. 數(shù)據(jù)集制作垃圾圖片的采集、標(biāo)注與劃分2.1 采集什么樣的圖片才有用訓(xùn)練垃圾分揀模型最忌諱的是圖片太“干凈”。很多團隊圖省事從網(wǎng)上下載了一堆商品圖背景是純白或純灰結(jié)果模型在實驗室里AP高到0.99一到傳送帶上就拉胯。為什么會這樣因為模型的泛化能力來自背景多樣性、光照差異、角度變化和信息噪音。我自己整理數(shù)據(jù)時定了三條硬標(biāo)準(zhǔn)背景必須多樣化覆蓋傳送帶、垃圾桶、地板、戶外、不同光照條件。目標(biāo)必須有多角度每個類別至少要拍到360度不同姿態(tài)同時包含俯視、側(cè)視、傾斜視角。必須有遮擋和堆疊樣本分揀業(yè)務(wù)里瓶子會疊在一起紙箱會半開口如果訓(xùn)練集里全是單目標(biāo)完美擺放模型永遠學(xué)不會處理實際工況。采集方法上我推薦用視頻抽幀而不是連拍照片。拿手機或工業(yè)相機對著垃圾堆橫向掃一遍錄10分鐘視頻然后每隔3到5幀抽一張這樣能極大豐富目標(biāo)姿態(tài)。抽幀腳本用OpenCV寫就行幾十行代碼搞定。2.2 標(biāo)注工具與標(biāo)注規(guī)范YOLOv5 4.0的標(biāo)注格式是每個圖片對應(yīng)一個txt文件每行是class_id x_center y_center width height坐標(biāo)值都需要歸一化到0到1之間。很多同學(xué)第一次做標(biāo)注會用labelimg生成VOC格式的XML然后再轉(zhuǎn)YOLO格式其實labelimg本身就直接支持YOLO格式輸出不用多這一步轉(zhuǎn)換。我用的工具是labelimg操作要點如下打開labelimg點擊“PascalVOC”切換成“YOLO”模式。用Open Dir打開圖片目錄Change Save Dir指定標(biāo)簽保存目錄?？旖萱IW新建框D下一張A上一張CtrlS保存。標(biāo)注規(guī)范這塊垃圾分揀很容易犯“框得太大”的毛病——把瓶子周圍的背景也包進去了。正確的做法是框要緊貼目標(biāo)可見部分的外輪廓如果目標(biāo)被部分遮擋就標(biāo)可見部分不要臆想完整的形狀。置信度閾值在訓(xùn)練時會因為框不準(zhǔn)確而出現(xiàn)嚴重波動寧可框小一點也不要框大。類別定義我建議用英文做標(biāo)簽名比如plastic_bottle、cardboard_box、aluminum_can、other_waste后面在yaml文件里再映射中文顯示名。這個習(xí)慣可以避免很多編碼問題。2.3 數(shù)據(jù)集劃分與目錄結(jié)構(gòu)YOLOv5 4.0期待的目錄結(jié)構(gòu)有兩種一種是官方COCO式的images和labels兩個大目錄然后在yaml里指定train和val路徑另一種是每個數(shù)據(jù)集子集單獨建目錄比如dataset/ ├── images/ │ ├── train/ │ └── val/ └── labels/ ├── train/ └── val/兩種都行我推薦第二種因為后續(xù)擴展test集更清晰。劃分比例我按8:1:1也就是80%訓(xùn)練、10%驗證、10%測試。注意劃分的時候要保證同一個物體的不同圖片不要被割裂得太嚴重但也不要刻意把同一視頻幀的數(shù)據(jù)全放一組否則驗證集會虛高。一個細節(jié)垃圾分揀數(shù)據(jù)集的類別數(shù)量不建議超過10類除非你確實需要區(qū)分很多細分材質(zhì)。類別越多每個類的訓(xùn)練樣本數(shù)就會攤薄模型在傳送帶高速場景下的誤檢率會明顯上升。我們項目初期定了12類訓(xùn)練后mAP只有0.78后來合并成8類mAP直接漲到0.92所以寧可在類別定義上做文章也別硬扛分類細度。2.4 數(shù)據(jù)增強策略與臟數(shù)據(jù)清洗YOLOv5 4.0的train.py里自帶數(shù)據(jù)增強包括mosaic、mixup、隨機透視、HSV擾動、隨機翻轉(zhuǎn)這些。對于垃圾分揀場景最有用的是mosaic和HSV擾動。因為垃圾材質(zhì)的顏色受光照影響特別大塑料瓶在不同燈光下反光效果完全不同HSV擾動會讓模型忽略顏色絕對數(shù)值、更多學(xué)習(xí)形狀和紋理特征。但增強也不是越狠越好。我在訓(xùn)練過程中發(fā)現(xiàn)mixup開啟過強會導(dǎo)致模型對小目標(biāo)的定位框偏移尤其是瓶蓋這種小物件。4.0版本里通過超參數(shù)hsv_h、hsv_s、hsv_v可以調(diào)節(jié)增強幅度我習(xí)慣把hsv_s和hsv_v調(diào)低到0.4左右保留一點顏色信息又不至于讓材質(zhì)失真。臟數(shù)據(jù)清洗是很多人忽略的一步。標(biāo)注完以后我建議跑一次分析腳本查看每類目標(biāo)的數(shù)量、目標(biāo)面積分布、圖片像素寬度比。如果發(fā)現(xiàn)某張圖片的標(biāo)簽特別多比如20多個框或者框面積小于圖片面積的0.1%這種樣本在訓(xùn)練時往往會算進太多背景負樣本必要時可以直接刪除。另外如果兩張圖片內(nèi)容幾乎一樣也會導(dǎo)致過擬合需要做去重。3. 模型訓(xùn)練參數(shù)調(diào)整、訓(xùn)練過程監(jiān)控與過擬合排查3.1 從預(yù)訓(xùn)練權(quán)重開始訓(xùn)練的底層邏輯垃圾分揀數(shù)據(jù)集規(guī)模通常只有幾千張從頭訓(xùn)練一個檢測模型非常虧。YOLOv5 4.0提供了一組在COCO上預(yù)訓(xùn)練好的權(quán)重這些權(quán)重已經(jīng)學(xué)會了“什么是物體邊緣”“什么是紋理”“什么是遮擋”這些通用特征。我們只需要在它的基礎(chǔ)上微調(diào)最后幾層讓它更懂“瓶子”“紙箱”“易拉罐”這幾個具體類別的模樣。我在項目里使用yolov5s.pt作為初始權(quán)重考慮到部署設(shè)備是嵌入式級別的板卡沒有直接上yolov5l或者yolov5x。有個常見誤區(qū)是模型越大效果越好但垃圾分揀的實時性要求很高yolov5s在1080Ti上推理能達到6到8毫秒一幀完全夠用yolov5l雖然mAP高一點但推理時間翻倍工業(yè)現(xiàn)場根本用不了。訓(xùn)練命令python train.py --data garbage.yaml --cfg models/yolov5s.yaml --weights weights/yolov5s.pt --batch-size 16 --epochs 200 --img 640--img 640的含義是訓(xùn)練輸入尺寸640x640垃圾分揀目標(biāo)一般中等偏小這個尺寸比較平衡。如果追求更快的訓(xùn)練能把尺寸降到416但小目標(biāo)的召回率會下降不太推薦。3.2 數(shù)據(jù)配置yaml文件寫法與anchor調(diào)整在寫yaml文件前我建議先按官方數(shù)據(jù)集格式組織目錄。garbage.yaml的寫法如下train: dataset/images/train val: dataset/images/val test: dataset/images/test nc: 8 names: [plastic_bottle, cardboard_box, aluminum_can, glass_bottle, plastic_bag, wood_block, metal_parts, other_waste]這里面最容易出問題的是路徑寫法。YOLOv5 4.0中路徑可以寫相對路徑也可以寫絕對路徑我在Windows上碰到過一次路徑分隔符反斜杠導(dǎo)致的讀取失敗改成正斜杠就正常了。Anchor是YOLOv5一個比較特別的設(shè)計。4.0版本在訓(xùn)練時會自動根據(jù)標(biāo)注數(shù)據(jù)重新計算anchor尺寸所以如果數(shù)據(jù)集里目標(biāo)比COCO普遍小很多anchor會自動變小不需要手工調(diào)整。但自動anchor計算也有個前提——目標(biāo)標(biāo)注必須足夠準(zhǔn)確。如果標(biāo)注框普遍松弛計算出的anchor也會偏大。訓(xùn)練日志里會有這樣一句Autoanchor: 4.6 anchors, 8.2, 9.1, etc, Better than old anchors看到這句話就說明anchor已經(jīng)重算過了。如果顯示Unchanged anchors說明預(yù)設(shè)anchor已經(jīng)夠用不用擔(dān)心。3.3 訓(xùn)練過程中的指標(biāo)監(jiān)控與loss曲線判斷訓(xùn)練時我最關(guān)心的不是總loss而是box_loss、obj_loss和cls_loss這三個分量的變化。YOLOv5 4.0訓(xùn)練會實時打印類似這樣的內(nèi)容Epoch 0: box_loss0.068, obj_loss0.035, cls_loss0.021, precision0.52, recall0.65, mAP0.50.48如果box_loss一直下降說明定位框?qū)W習(xí)正常如果obj_loss降不下去說明目標(biāo)和背景的區(qū)分度不夠這時候需要考慮增加難負樣本比如空垃圾袋、紙屑這類容易被當(dāng)成目標(biāo)的“疑似物”cls_loss抖動劇烈往往和某個類別的樣本過少有關(guān)。我更習(xí)慣用tensorboard實時監(jiān)控曲線YOLOv5 4.0自帶了tensorboard --logdir runs的啟動命令。訓(xùn)練了大概50輪后能明顯看到曲線斜率變緩這是正?，F(xiàn)象。到150輪左右mAP基本穩(wěn)定再往后訓(xùn)練收益很低反而有過擬合風(fēng)險。3.4 過擬合問題排查一個典型的分揀業(yè)務(wù)案例我們項目第一次訓(xùn)練時驗證集mAP達到了0.94很漂亮但現(xiàn)場實測識別率只有0.6。這種現(xiàn)象幾乎可以斷定是過擬合。怎么排查先看訓(xùn)練集和驗證集的loss曲線。如果訓(xùn)練集loss持續(xù)下降而驗證集loss在第120輪后開始反彈說明模型開始死記硬背訓(xùn)練集特征了。我做了三件事關(guān)掉部分數(shù)據(jù)增強將--cache的緩存大圖勾選去掉同時調(diào)低mosaic概率。增加驗證集數(shù)量原驗證集只占總數(shù)據(jù)5%后來重新劃分成15%。在模型最后一層加dropoutYOLOv5沒有簡單的dropout開關(guān)但可以通過降低模型深度系數(shù)、增加正則化來緩解。我把yolov5s.yaml的depth_multiple從0.33降到0.3相當(dāng)于讓模型變窄一點強制它學(xué)更抽象的特征。這三步做完現(xiàn)場識別率恢復(fù)到0.88。經(jīng)驗就是垃圾分揀場景的數(shù)據(jù)集規(guī)模不會太大與其追高訓(xùn)練集指標(biāo)不如把驗證集指標(biāo)壓穩(wěn)。3.5 訓(xùn)練完成后的指標(biāo)評估與選擇best.pt還是last.pt訓(xùn)練結(jié)束會在runs/train/exp下生成兩個權(quán)重文件best.pt和last.pt。名字已經(jīng)說明了一切——best.pt是訓(xùn)練過程中驗證集mAP最高的權(quán)重last.pt是最后一輪的權(quán)重。工業(yè)落地上我強烈建議用best.pt但要注意best.pt的“best”是基于驗證集算出來的驗證集和實際分揀場景可能還有分布差異。穩(wěn)妥的做法是挑三到四個不同訓(xùn)練階段的best.pt用現(xiàn)場錄像截圖做成一個獨立測試集跑val.py對比出真正的“最優(yōu)”。這個額外驗證步驟能過濾掉很多“紙面mAP大師”。另外驗證集mAP0.5和mAP0.5:0.95的區(qū)別也要關(guān)注。垃圾分揀業(yè)務(wù)一般只看mAP0.5因為分揀機械臂抓取時并不嚴格要求框和真實目標(biāo)的高度貼合IOU閾值0.5已經(jīng)足夠。如果mAP0.5:0.95很低但mAP0.5很高說明定位精度一般但對業(yè)務(wù)夠用。4. 部署到分揀業(yè)務(wù)流模型導(dǎo)出、推理加速與現(xiàn)場調(diào)優(yōu)4.1 從PyTorch到ONNX再到TensorRT的導(dǎo)出路徑訓(xùn)練完后的PyTorch權(quán)重沒法直接在生產(chǎn)環(huán)境里跑需要一個導(dǎo)出過程。YOLOv5 4.0自帶了export.py能導(dǎo)出ONNX、TorchScript、CoreML等格式。我這邊最終部署用的是TensorRT因為分揀工位的算力卡非常怕CPU推理而TensorRT能最大程度利用GPU的并行能力。導(dǎo)出ONNX的命令python export.py --weights runs/train/exp/best.pt --include onnx --opset 11導(dǎo)出時會有一個注意點YOLOv5 4.0的檢測頭里有部分動態(tài)shape操作直接用onnxruntime跑會報錯需要在export時加--dynamic參數(shù)。但我更推薦的是一步到位導(dǎo)出TensorRT enginepython export.py --weights best.pt --include engine --device 0TensorRT引擎是序列化過的優(yōu)化計算圖推理速度比ONNX還快。前提是TensorRT版本和你用的CUDA版本兼容我這邊用了TensorRT 8.0配合CUDA 11.1非常順。4.2 推理端代碼改造與輸入預(yù)處理YOLOv5 4.0官方detect.py是面向演示的直接放到生產(chǎn)環(huán)境會有問題一是每幀都打印日志和畫框影響吞吐二是一次只能處理單張圖無法應(yīng)對連續(xù)幀流水。我改造推理端時做了三件事移除可視化部分只返回檢測框數(shù)組。使用torch.cuda.synchronize()控制幀率。把輸入圖像resize為640x640時使用letterbox方式保留長寬比而不是強制拉伸。這個細節(jié)極其重要——垃圾瓶子細長如果拉伸變形寬高比信息丟失檢測框會嚴重偏移。letterbox的算法本身很簡單計算縮放比例補邊為114像素的灰色。YOLOv5 4.0的utils里已經(jīng)有l(wèi)etterbox函數(shù)直接復(fù)用即可。如果你自己寫一定要記得在推理前對圖像執(zhí)行同樣的預(yù)處理、推理后對坐標(biāo)做反向映射。4.3 NMS閾值和置信度閾值的現(xiàn)場調(diào)優(yōu)模型訓(xùn)練時的閾值和部署時的閾值是兩回事。訓(xùn)練階段的conf_thres和iou_thres影響的是mAP計算方式部署階段的這兩個參數(shù)直接決定現(xiàn)場檢測行為。我剛開始部署時用的默認閾值conf_thres0.25, iou_thres0.45結(jié)果是傳送帶上偶爾漏撿。后來把conf_thres降到0.15把iou_thres調(diào)到0.5漏檢率好了不少但誤檢略有增加。于是我在中間加了一個按類別設(shè)定閾值的邏輯對“塑料瓶”這種形態(tài)穩(wěn)定的類別用0.3對“其他垃圾”這種外形不規(guī)則的類別用0.15。這樣整體識別效果更符合分揀業(yè)務(wù)的需求。還有一個坑當(dāng)目標(biāo)密集堆疊時NMS的max_det參數(shù)默認是300如果一幀圖像里檢測超過300個框后面的目標(biāo)會被丟棄導(dǎo)致漏檢。垃圾分揀場景單幀一般不超過50個目標(biāo)所以我把max_det設(shè)為100減少內(nèi)存壓力。4.4 現(xiàn)場現(xiàn)場光照變化與魯棒性驗證分揀傳送帶的現(xiàn)場環(huán)境遠不是實驗室那種穩(wěn)定光源。我在部署時發(fā)現(xiàn)上午9點陽光照進廠房時模型對“鋁罐”的識別率跌了20%。原因是鋁罐高光反射蛋白白色區(qū)域在圖片里和透明塑料袋特別像。解決方案不是重新采集大量視頻數(shù)據(jù)而是先用opencv對接入攝像頭的曝光參數(shù)做了硬約束把曝光時間固定、自動白平衡關(guān)閉。另外在內(nèi)網(wǎng)里部署了夜間紅外補光燈進一步杜絕環(huán)境光波動。這種工程手段比加數(shù)據(jù)更直接。但注意如果換了攝像頭型號曝光參數(shù)需要重新標(biāo)定不能直接套用。4.5 與PLC和機械臂的聯(lián)動邏輯分揀業(yè)務(wù)的最終目標(biāo)是把檢測結(jié)果轉(zhuǎn)成機械臂抓取指令。這里有個檢測速度之外的瓶頸YOLOv5推理返回的是圖像坐標(biāo)系里的框坐標(biāo)而機械臂需要的是相機坐標(biāo)系或世界坐標(biāo)系下的抓取點。我的做法是先用張正友標(biāo)定法對相機做內(nèi)參標(biāo)定再通過手眼標(biāo)定建立相機到機械臂底座的變換矩陣。之后推理出的目標(biāo)中心點坐標(biāo)通過矩陣變換映射到機械臂空間。由于傳送帶一直在動我還加了一個簡單的速度補償根據(jù)傳送帶線速度和檢測到目標(biāo)到機械臂末端的運動時間差預(yù)測目標(biāo)到達抓取點時的新位置。這部分一開始用固定速度仿真效果還行但實際傳輸帶速度波動時會抓偏后來在編碼器上加了一個滑動平均濾波誤差控制在0.5厘米以內(nèi)。這一整套邏輯下來從圖像采集到最后機械臂抓取端到端延遲大概在120毫秒左右滿足低中速傳送帶0.3米/秒以下的分揀需求。5. 訓(xùn)練之外必須知道的事數(shù)據(jù)迭代閉環(huán)和成本控制5.1 建立“臟樣本回流”機制模型上線后絕對不是一勞永逸。分揀現(xiàn)場每周都會出現(xiàn)新的“意外垃圾”比如濕紙巾、果皮、捆扎帶這些在訓(xùn)練集里根本沒有對應(yīng)類別模型會強行把它們預(yù)測成最相近的類別結(jié)果機械臂抓了不該抓的東西。解決方法是建立回流機制現(xiàn)場IPC把置信度低于0.3的檢測框和對應(yīng)截圖自動保存到一個“待標(biāo)注”目錄。每隔兩周標(biāo)注人員只處理這些新增難例補充到訓(xùn)練集重新訓(xùn)練一輪。這個機制第一周就幫我們發(fā)現(xiàn)了70多個誤檢案例第二周模型識別率又漲了4個百分點。對于垃圾分揀這種長尾分布極其明顯的場景這個閉環(huán)比一開始堆數(shù)據(jù)集更經(jīng)濟、更有效。5.2 硬件選型與訓(xùn)練時間成本估算在訓(xùn)練成本上YOLOv5s用單張1080Ti訓(xùn)練200輪、每輪約1000張圖大概需要2.5到3小時。如果是用CPU訓(xùn)練同樣的工作量可能要40個小時以上。所以哪怕只是短期項目我也建議租一塊云GPU來跑訓(xùn)練而不是用本機CPU硬抗。訓(xùn)練完成后把best.pt拉回本地部署即可。數(shù)據(jù)集標(biāo)注成本也要算清楚每張圖如果有2到3個目標(biāo)熟練標(biāo)注員每小時能標(biāo)注80到120張圖。一個3000張圖的垃圾數(shù)據(jù)集純標(biāo)注工作在25小時左右?？雌饋聿欢嗟由腺|(zhì)量檢查和返工實際大概需要2人日。所以標(biāo)注前定嚴格規(guī)范非常有必要否則返工成本會翻倍。5.3 類別不均衡和背景負樣本的取舍垃圾分揀數(shù)據(jù)集里塑料瓶往往占60%以上而金屬零件可能只有5%。直接按原始分布訓(xùn)練模型會對金屬零件類別的召回率特別低。我的做法是對樣本量少的類別做離線復(fù)制增強不是簡單的復(fù)制粘貼而是對該類圖片做小幅旋轉(zhuǎn)、平移、縮放后重新保存為新的樣本。同時在下采樣階段限制每張圖每類的實例數(shù)不超過30防止一個場景里的同類別目標(biāo)刷屏導(dǎo)致訓(xùn)練失衡。背景負樣本是另一個容易被忽略的坑。分揀場景里空無一物的傳送帶圖片在訓(xùn)練集里一張都不能少。這些負樣本能讓模型的obj_loss快速收斂減少誤報。我每批次會混入5%到10%的空背景圖效果很明顯。5.4 模型更新時的權(quán)重管理每個版本迭代出的best.pt我都會在權(quán)重文件名里加上日期和數(shù)據(jù)量后綴比如best_20240601_3200img.pt。這樣在比對現(xiàn)場效果時能快速定位到底是哪套數(shù)據(jù)訓(xùn)出來的模型。有一次現(xiàn)場反饋識別率下降排查了半天發(fā)現(xiàn)是某個同事用了舊版模型文件沒有同步更新。權(quán)重管理聽著啰嗦但在工業(yè)項目里非常重要。另外最好保留訓(xùn)練當(dāng)時的opt.yaml文件YOLOv5 4.0每次訓(xùn)練都會自動保存。這個文件記錄了所有超參數(shù)想要復(fù)現(xiàn)結(jié)果或排查差異時直接載入即可不用猜。6. 一段供直接參考的完整訓(xùn)練流程腳本這里把我自己在項目里常用的訓(xùn)練流程串起來給你一份可以照抄的清單不包含花哨操作只保留穩(wěn)定跑通的部分。# 1. 激活環(huán)境 conda activate yolov5 # 2. 數(shù)據(jù)目錄準(zhǔn)備好后先驗證標(biāo)注格式 python -c import os img_dir dataset/images/train label_dir dataset/labels/train count 0 for img in os.listdir(img_dir): label os.path.join(label_dir, img.replace(.jpg, .txt)) if not os.path.exists(label): print(缺少label:, img) count 1 print(缺失標(biāo)簽數(shù)量:, count) 運行這個驗證腳本能第一時間發(fā)現(xiàn)漏標(biāo)問題比直接開訓(xùn)練再報錯要高效得多。我曾經(jīng)靠這個腳本一次性找出37張漏標(biāo)圖片。# 2. 修改garbage.yaml路徑確認nc和names正確 # 3. 開始訓(xùn)練 python train.py --data garbage.yaml --cfg models/yolov5s.yaml --weights weights/yolov5s.pt --batch-size 16 --epochs 200 --img 640 --project garbage_project --name v1 # 4. 訓(xùn)練完后在runs/train/garbage_project/v1/weights下找到best.pt # 5. 用val.py評估測試集 python val.py --data garbage.yaml --weights runs/train/garbage_project/v1/weights/best.pt --img 640 # 6. 導(dǎo)出TensorRT engine python export.py --weights runs/train/garbage_project/v1/weights/best.pt --include engine --device 0我訓(xùn)練時加過很多額外參數(shù)比如--cos-lr、--label-smoothing但實測下來對垃圾分揀數(shù)據(jù)集提升不大反而增加了訓(xùn)練時間。如果你是第一次跑建議先不加這些花哨參數(shù)跑通整條鏈路再說。關(guān)于批次大小如果顯存只有8G--batch-size 16跑640分辨率可能會OOM可以把batch降到8或者把--img降到544。不要為了湊batch瘋狂調(diào)低分辨率垃圾分揀的小目標(biāo)會因此丟失得不償失。7. 最后再分享一些實際部署中的小心得這個項目前后折騰了將近兩個月最深的體會是YOLOv5 4.0本身只是一個工具真正的難點全都在工具之外。環(huán)境配置看起來繁瑣但按部就班來半天時間絕對能搞定數(shù)據(jù)集制作才是整個項目里最花時間的部分占掉了我六成精力而模型訓(xùn)練反而最簡單因為官方代碼已經(jīng)足夠成熟剩下的事情就是盯曲線、調(diào)閾值。再提一個很多人沒注意到的細節(jié)4.0版本的--cache參數(shù)官方默認表示緩存圖片到內(nèi)存如果內(nèi)存不足會有警告。但很多教程會把--cache當(dāng)成默認照抄下來我建議顯存緊張時不要開緩存否則訓(xùn)練過程中可能因為內(nèi)存交換導(dǎo)致整機卡死。我踩過的另一個坑是Windows平臺上PyTorch的DataLoader多線程問題。如果num_workers設(shè)得過高在Windows下會出現(xiàn)DataLoader worker超時導(dǎo)致訓(xùn)練中斷。建議Windows用戶把--workers設(shè)為0或2Linux下可以設(shè)到8。分揀業(yè)務(wù)流真正落地時視覺檢測只是其中一環(huán)機械臂的軌跡規(guī)劃、傳送帶的啟停邏輯、PLC信號的時序配合每一個環(huán)節(jié)都需要跨團隊溝通。作為算法工程師別只顧著調(diào)模型多去現(xiàn)場看看機械臂抓空或者抓漏的瞬間很多模型上的問題在現(xiàn)場看一次就明白了。代碼和數(shù)據(jù)可以隨時更新但現(xiàn)場的經(jīng)驗不會自動積累。希望這篇內(nèi)容能讓你少踩一些我踩過的坑也歡迎交流你們在垃圾分揀或者類似檢測場景里遇到的具體問題。