戰(zhàn)：YOLO訓(xùn)練到TensorRT部署全流程)
簡介基于YOLO的軸承生產(chǎn)缺陷檢測數(shù)據(jù)集面向從事工業(yè)視覺檢測、目標(biāo)檢測算法學(xué)習(xí)與實(shí)戰(zhàn)的開發(fā)者。數(shù)據(jù)集圍繞軸承生產(chǎn)中的典型缺陷識別問題提供568張已標(biāo)注圖片按三類缺陷完成邊界框標(biāo)注可直接用于YOLO系列模型的訓(xùn)練、驗(yàn)證與測試。包體共1772個文件涵蓋624張jpg圖像、577個txt格式的YOLO標(biāo)注、568個xml格式的VOC標(biāo)注以及2個Python訓(xùn)練腳本和1個yaml模型配置文件可同時兼容主流檢測框架的數(shù)據(jù)格式。壓縮包大小755.12MB整體目錄結(jié)構(gòu)清晰適合作為工業(yè)質(zhì)檢場景的入門練手或課題參考目前已有337人瀏覽學(xué)習(xí)。借助這套數(shù)據(jù)與腳本使用者能夠快速搭建缺陷檢測流程從數(shù)據(jù)預(yù)處理、模型訓(xùn)練到結(jié)果可視化還可針對光照變化、拍攝角度、背景干擾等生產(chǎn)環(huán)境條件開展魯棒性調(diào)優(yōu)為實(shí)際質(zhì)量管控提供可落地的技術(shù)路徑。1. 基于YOLO的軸承生產(chǎn)缺陷檢測只有568張圖和三類標(biāo)簽?zāi)軗纹鹨粭l產(chǎn)線嗎軸承表面缺陷檢測的殘酷現(xiàn)實(shí)是產(chǎn)線上真正能收集到的缺陷樣本常常就是幾百張而不是上萬張。568張圖、三類缺陷這個規(guī)模在工業(yè)視覺項(xiàng)目里非常典型也最容易翻車——模型訓(xùn)練時mAP看著能到0.95一上產(chǎn)線就被良品表面紋理帶偏或者干脆把小劃傷漏掉。問題不在YOLO本身而在數(shù)據(jù)怎么劃分、類別邊界怎么定、訓(xùn)練參數(shù)怎么守紀(jì)律。這篇文章要解決的就是這件事用568張圖把YOLO訓(xùn)練成能用的軸承外觀檢測器。從數(shù)據(jù)整理、標(biāo)注格式、模型選型、訓(xùn)練參數(shù)到TensorRT加速和產(chǎn)線驗(yàn)證全部按可復(fù)現(xiàn)的步驟寫。適合手里有少量缺陷圖、想快速驗(yàn)證方案可行性或者已經(jīng)訓(xùn)過但現(xiàn)場表現(xiàn)不穩(wěn)定的工程師。2. 把568張圖整理成YOLO能吃的數(shù)據(jù)集標(biāo)注、劃分與增強(qiáng)2.1 三類缺陷怎么定類別邊界是第一個坑標(biāo)題里寫了“類別三類”但具體哪三類要自己定。常見做法是按軸承外觀缺陷的發(fā)生頻率和檢測價值來選劃痕scratch、麻點(diǎn)pit、剝落spall。這三類在視覺特征上差異明顯——劃痕是線狀、方向隨機(jī)麻點(diǎn)是點(diǎn)狀、邊緣清晰剝落是片狀、紋理粗糙偏暗。選這三類的另一個好處是標(biāo)注維度一致不會出現(xiàn)一類缺陷邊界模糊到標(biāo)注工都拿不準(zhǔn)的情況。類別邊界一旦確定最忌諱的就是“這是個很大的劃痕但好像也算剝落”。標(biāo)注的時候手一松模型學(xué)到的邊界就糊了。我的做法是先做一個標(biāo)注規(guī)范樣例圖每類選兩張典型圖、兩張邊界圖邊界圖標(biāo)注時遵循“占主導(dǎo)特征定類”的原則。568張圖雖然不多但每張標(biāo)注的都越干凈后面訓(xùn)練越省心。寧可三類各標(biāo)注100張左右也不要三類里某一類只有40張、另一類有360張。2.2 整理目錄結(jié)構(gòu)與標(biāo)簽一個腳本解決訓(xùn)練集劃分YOLO訓(xùn)練要求的數(shù)據(jù)組織方式很固定images目錄下放原圖labels目錄下放同名txt文件每行是“類別 cx cy w h”cx、cy、w、h都是相對圖像寬高的歸一化坐標(biāo)。568張圖我習(xí)慣按7:2:1劃分成訓(xùn)練、驗(yàn)證、測試最終目錄結(jié)構(gòu)如下dataset/ ├── images/ │ ├── train/ # 398張 │ ├── val/ # 113張 │ └── test/ # 57張 ├── labels/ │ ├── train/ │ ├── val/ │ └── test/ ├── data.yaml └── classes.txt數(shù)據(jù)劃分必須按“軸承個體”來分不能按圖隨機(jī)分。同一個軸承的多張不同角度照片如果同時出現(xiàn)在訓(xùn)練集和測試集模型其實(shí)是在“背照片”而不是“學(xué)缺陷”。這個坑我在早期項(xiàng)目里踩過當(dāng)時測試集mAP虛高到0.99換上全新軸承圖立刻掉到0.7以下。劃分腳本用Python寫關(guān)鍵是先按軸承ID分組再把組列表shuffle后切分import os import random import shutil from collections import defaultdict img_root raw_images # 原始圖片目錄文件名前綴為軸承ID如 B001_01.jpg out_root dataset train_ratio, val_ratio 0.7, 0.2 # 1. 按軸承ID分組 groups defaultdict(list) for fname in os.listdir(img_root): if not fname.endswith(.jpg): continue bearing_id fname.split(_)[0] groups[bearing_id].append(fname) items list(groups.items()) random.seed(42) random.shuffle(items) # 2. 按組劃分 train_files, val_files, test_files [], [], [] for idx, (bid, files) in enumerate(items): if idx int(len(items) * train_ratio): train_files files elif idx int(len(items) * (train_ratio val_ratio)): val_files files else: test_files files # 3. 復(fù)制圖像并生成同名標(biāo)簽txt標(biāo)簽數(shù)據(jù)假設(shè)已放在raw_labels目錄 for split, files in [(train, train_files), (val, val_files), (test, test_files)]: os.makedirs(f{out_root}/images/{split}, exist_okTrue) os.makedirs(f{out_root}/labels/{split}, exist_okTrue) for fname in files: shutil.copy(f{img_root}/{fname}, f{out_root}/images/{split}/{fname}) label_name fname.replace(.jpg, .txt) src_label fraw_labels/{label_name} if os.path.exists(src_label): shutil.copy(src_label, f{out_root}/labels/{split}/{label_name}) print(ftrain{len(train_files)} val{len(val_files)} test{len(test_files)})這段腳本里有兩個參數(shù)值得注意一個是random.seed(42)固定隨機(jī)種子保證每次復(fù)現(xiàn)同樣的劃分結(jié)果方便別人復(fù)現(xiàn)你的實(shí)驗(yàn)另一個是train_ratio和val_ratio的比例568張圖本身不多7:2:1是工業(yè)小樣本最常用的比例測試集57張足夠算漏檢率再少就不穩(wěn)定了。如果你的缺陷類別分布不均衡劃分后還要檢查一下每個split里的類比比例別讓訓(xùn)練集中某一類只有20張。2.3 數(shù)據(jù)增強(qiáng)小樣本的后悔藥568張圖直接訓(xùn)練YOLOv8s過擬合幾乎是必然的。增強(qiáng)策略我一般分兩層第一層是YOLO自帶的在線增強(qiáng)第二層是離線補(bǔ)充增強(qiáng)樣本。第一層優(yōu)先因?yàn)樵诰€增強(qiáng)每個epoch都隨機(jī)變化等效訓(xùn)練樣本量翻好幾倍。YOLOv8默認(rèn)的增強(qiáng)參數(shù)比較激進(jìn)對小樣本來說hsv_h0.015、hsv_s0.7、hsv_v0.4、fliplr0.5、mosaic1.0這些都夠用。離線增強(qiáng)只做一類操作模擬產(chǎn)線光照變化和表面反光。用cv2對訓(xùn)練集中的圖片做亮度±30%調(diào)整、小角度旋轉(zhuǎn)±10度、高斯模糊模擬相機(jī)輕微失焦但注意不要在缺陷區(qū)域引入形變。軸承是圓形對稱件旋轉(zhuǎn)增強(qiáng)是安全的但對有方向性的劃痕旋轉(zhuǎn)90度后視覺特征可能就不符合真實(shí)分布了這點(diǎn)要按產(chǎn)線實(shí)拍情況決定。一個容易被忽略的增強(qiáng)是Mosaic和Copy-Paste的配合。Mosaic是把四張圖拼成一張對小目標(biāo)有效但軸承缺陷通常只占幾個像素到幾十個像素Mosaic后目標(biāo)更小了。Copy-Paste則把缺陷區(qū)域隨機(jī)粘貼到無缺陷的軸承圖上這種離線增強(qiáng)對增加類別樣本量很有用前提是缺陷邊緣要帶原圖的紋理像素不要直接硬摳否則模型學(xué)到的就是“邊緣生硬缺陷”。3. 模型選型與訓(xùn)練用YOLOv8s還是更大版本關(guān)鍵看推理預(yù)算3.1 模型尺寸、損失函數(shù)與推理速度的權(quán)衡568張數(shù)據(jù)集訓(xùn)練大模型沒有意義。YOLOv8s參數(shù)量約1100萬COCO上mAP比n版高不少但速度只慢一點(diǎn)點(diǎn)是工業(yè)檢測的甜點(diǎn)位。更大的YOLOv8m或v8l在568張圖上很容易把訓(xùn)練集損失壓到接近0驗(yàn)證集表現(xiàn)反而變差因?yàn)閰?shù)容量太大把樣本噪聲也學(xué)進(jìn)去了。社區(qū)里討論的YOLO版本迭代很多從v5到v8再到各類結(jié)構(gòu)變體核心邏輯沒變損失函數(shù)由BBox回歸損失、分類損失和DFLDistribution Focal Loss組成。對小目標(biāo)、小樣本DFL的作用尤其明顯它對邊界框的定位不確定性建模輸出的是分布而不是單值這能讓模型在缺陷邊緣模糊時給出更保守的定位。T4上跑TensorRT FP16YOLOv8s 640x640輸入的單幀推理時間一般在2~3毫秒左右這個數(shù)據(jù)是多次實(shí)測的經(jīng)驗(yàn)范圍。如果檢測幀率要求25FPS即每幀預(yù)算40毫秒單張T4卡理論上能串行處理十幾路視頻流但考慮到預(yù)處理、后處理和顯存帶寬工程上8路左右是穩(wěn)妥值。這個估算后面還會細(xì)說先記住結(jié)論選YOLOv8s不是因?yàn)樗茸罡叨且驗(yàn)樗谛颖竞蛯?shí)時推理之間最平衡。3.2 訓(xùn)練命令與關(guān)鍵超參數(shù)YOLOv8訓(xùn)練用ultralytics庫。data.yaml里要寫清楚路徑、類別數(shù)和類別名。paths字段建議寫相對路徑這樣代碼拷到別的機(jī)器不用改path: dataset train: images/train val: images/val test: images/test names: 0: scratch 1: pit 2: spall訓(xùn)練命令本身不復(fù)雜但有幾個參數(shù)是必須調(diào)的。下面這條命令是我在類似小樣本項(xiàng)目里的基準(zhǔn)配置yolo detect train \ modelyolov8s.pt \ datadata.yaml \ epochs200 \ batch16 \ imgsz640 \ patience30 \ optimizerSGD \ lr00.01 \ cos_lrTrue \ mosaic1.0 \ seed42 \ projectruns/detect \ namebearing_defect參數(shù)說明epochs200看似多但配合patience30驗(yàn)證集指標(biāo)連續(xù)30輪不漲就會自動早停實(shí)際通常100輪左右收斂。batch16在顯存允許下盡量大小樣本數(shù)據(jù)集的BN統(tǒng)計(jì)量對batch size敏感batch太小比如4或8會讓BN的均值方差震蕩損失曲線抖動明顯。lr00.01配SGD是小樣本的標(biāo)準(zhǔn)開局如果你用AdamWlr0要降到0.001左右否則前幾個epoch損失容易爆掉。cos_lrTrue讓學(xué)習(xí)率余弦衰減相比階梯下降在小樣本上能多擠出一點(diǎn)精度。有個坑要提醒mosaic1.0默認(rèn)是啟用的但如果你發(fā)現(xiàn)模型在驗(yàn)證集上對完整軸承圖的檢測效果差而訓(xùn)練時loss降得很漂亮可以嘗試把mosaic在最后20個epoch關(guān)掉ultralytics支持close_mosaic10這樣的參數(shù)因?yàn)镸osaic產(chǎn)生的拼接圖與真實(shí)完整軸承圖像分布有偏差最后階段切換到真實(shí)分布能讓模型適應(yīng)更好。3.3 訓(xùn)練結(jié)果怎么看別只看mAP訓(xùn)練結(jié)束后ultralytics會輸出P、R、mAP50、mAP50-95等指標(biāo)。在568張的小數(shù)據(jù)集上我更看重R召回率因?yàn)楫a(chǎn)線上漏檢比誤檢嚴(yán)重得多——誤檢最多是人工復(fù)檢多看一眼漏檢是缺陷直接流入下游。如果R為0.93意味著100個缺陷里有7個漏掉這個數(shù)字對軸承件來說通常不可接受。怎么看訓(xùn)練日志里的loss曲線主要看box_loss和cls_loss是否同步下降。如果box_loss降得很慢而cls_loss正常說明邊界框回歸困難可能原因是標(biāo)注框不貼合缺陷邊緣如果cls_loss下不去大概率是類別邊界標(biāo)注混亂。另外val_cls_loss如果出現(xiàn)先降后升的“U型”就是過擬合的典型信號此時最優(yōu)模型往往是val損失最低點(diǎn)附近的checkpoint而不是最后一個epoch的權(quán)重。提示訓(xùn)練完成后把best.pt在test集上單獨(dú)跑一次得到的結(jié)果才是真正可以匯報的精度。val集參與了早停判斷指標(biāo)有樂觀偏差test集是完全沒見過、不參與訓(xùn)練決策的數(shù)據(jù)。這一點(diǎn)雖然基礎(chǔ)但很多項(xiàng)目為了省事跳過test集測試最后上線效果對不上實(shí)驗(yàn)數(shù)據(jù)原因就在這。4. 避坑記錄568張軸承數(shù)據(jù)訓(xùn)練YOLO的5個坑4.1 坑mAP高但某個類別幾乎全漏——類別不平衡現(xiàn)象訓(xùn)練結(jié)束后mAP50在0.9以上但查看每個類別的confusion matrix發(fā)現(xiàn)“剝落”類召回率只有0.4。原因568張圖里剝落只有不到80個實(shí)例而劃痕有300多個實(shí)例。YOLO的默認(rèn)分類損失對樣本多的類別傾斜模型學(xué)會了劃痕對剝落則傾向于保守預(yù)測。解決先統(tǒng)計(jì)每個類別的實(shí)例數(shù)量而不是圖片數(shù)量。如果剝落實(shí)例明顯偏少對這個類別做離線復(fù)制粘貼增強(qiáng)目標(biāo)是把三類實(shí)例數(shù)拉平到同一量級。另外可以給少樣本類別在loss里加權(quán)重ultralytics支持在data.yaml的names后接一個weight字段但實(shí)際操作里增強(qiáng)數(shù)據(jù)比調(diào)loss權(quán)重更穩(wěn)。4.2 坑驗(yàn)證集mAP 0.98一到產(chǎn)線就翻車——訓(xùn)練測試同源現(xiàn)象實(shí)驗(yàn)階段mAP漂亮得不像話新采一批軸承圖測試檢測率斷崖式下跌。原因數(shù)據(jù)集劃分時按圖片隨機(jī)分同一個軸承的多張照片同時進(jìn)了訓(xùn)練集和測試集。YOLO記住了軸承紋理背景測試時靠背景相似性就能“猜”出缺陷位置。解決嚴(yán)格按軸承個體分組劃分本文2.2的腳本已經(jīng)實(shí)現(xiàn)了這個邏輯。還有一個驗(yàn)證方法從測試集里隨機(jī)抽20張圖人工確認(rèn)沒有任何一張?jiān)谟?xùn)練集出現(xiàn)過把“同源”這道防線卡死在數(shù)據(jù)準(zhǔn)備階段。4.3 坑損失正常、metrics正常但檢測框偏移半個缺陷位——標(biāo)注邊界不一致現(xiàn)象loss曲線一切正常mAP50在0.88附近但可視化檢測結(jié)果發(fā)現(xiàn)預(yù)測框總是比標(biāo)注框大一圈或偏一側(cè)。原因標(biāo)注時一類框緊貼缺陷邊緣另一類框則多包含了一圈背景。YOLO回歸的是邊界框坐標(biāo)標(biāo)注的不一致會讓模型學(xué)到一個“平均框”看起來偏差不大但后續(xù)做缺陷尺寸篩選時就出問題了。解決標(biāo)注規(guī)范里明確“外接矩形緊貼缺陷最外緣”并且每張圖標(biāo)注完成后用Python腳本檢查框?qū)捀叻植籍惓捀弑鹊目騿为?dú)挑出來復(fù)核。4.4 坑小劃痕漏檢率高——目標(biāo)尺度問題現(xiàn)象形態(tài)檢測單獨(dú)看都很好但5像素寬、40像素長的細(xì)微劃痕10個里漏掉4個。原因YOLO輸入imgsz640時這種細(xì)長缺陷在特征圖上的響應(yīng)非常弱。就算訓(xùn)練loss包含了DFL目標(biāo)太小的情況下特征層根本沒有足夠的語義信息來區(qū)分劃痕和正常紋理。解決第一imgsz從640提到960或1280小目標(biāo)AP會明顯上升代價是推理時間增加需要重新評估TensorRT路數(shù)第二使用SAHI這類切片推理工具把大圖切成小塊分別檢測第三檢查是否開啟了多尺度訓(xùn)練scale0.5這種設(shè)置讓模型見識更多尺寸變化對小目標(biāo)魯棒性有幫助。4.5 坑過擬合到增強(qiáng)樣本本身——增強(qiáng)參數(shù)過猛現(xiàn)象訓(xùn)練集loss降到0.05驗(yàn)證集loss停在0.3不再下降檢測結(jié)果對光照異常敏感。原因hsv增強(qiáng)參數(shù)調(diào)太猛比如hsv_v0.8導(dǎo)致訓(xùn)練圖中出現(xiàn)大量現(xiàn)實(shí)中不會出現(xiàn)的極端亮度模型花了大量容量去擬合這些失真圖。解決回顧增強(qiáng)參數(shù)把hsv_v回調(diào)到0.4以下關(guān)閉mosaic的最后10個epoch。小樣本數(shù)據(jù)集的增強(qiáng)原則是“輕微擾動模擬真實(shí)變化”而不是“大力出奇跡”。每次改增強(qiáng)參數(shù)只改一個變量跑一輪訓(xùn)練看驗(yàn)證集趨勢不要一次性把一堆參數(shù)全改了那樣出了問題根本定位不到原因。5. 從驗(yàn)證到產(chǎn)線TensorRT加速與實(shí)時路數(shù)估算5.1 TensorRT FP16下的吞吐量與路數(shù)估算訓(xùn)練的模型要跑到產(chǎn)線上PyTorch推理的速度通常不夠用。我的流程是先用ultralytics導(dǎo)出ONNX再用TensorRT轉(zhuǎn)成FP16 engine。完成轉(zhuǎn)換后在同一張T4卡上實(shí)測單幀推理耗時。下面這張表是經(jīng)驗(yàn)參考值實(shí)際值會因輸入分辨率和顯存占用略有浮動模型輸入分辨率TensorRT FP16單幀耗時25FPS單路預(yù)算T4可支撐路數(shù)YOLOv8s640x6402~3ms40ms8~10路YOLOv8n640x6401~1.5ms40ms12~16路YOLOv8s1280x12805~8ms40ms4~6路估算方法很簡單單幀推理耗時如果只有2.5ms而每路25FPS要求每幀40ms內(nèi)處理完理論上GUP可以在一路等待下一幀的時間內(nèi)串行處理多個幀。但實(shí)際還要算上圖像解碼、預(yù)處理、后處理NMS的耗時以及顯存帶寬競爭。因此我從來不用理論峰值排產(chǎn)線一般留50%余量T4640分辨率按8路設(shè)計(jì)。如果現(xiàn)場需要更多路數(shù)就把模型換成YOLOv8n或者把輸入分辨率降到480x480但檢測精度要重新驗(yàn)證。5.2 產(chǎn)線驗(yàn)證腳本統(tǒng)計(jì)漏檢率和誤檢率TensorRT跑起來的模型最終要用一個最小驗(yàn)證腳本來確認(rèn)工程指標(biāo)。這里的重點(diǎn)不是寫代碼而是定義好什么算漏檢、什么算誤檢。我習(xí)慣的做法是設(shè)定一個IoU閾值比如0.5和一個置信度閾值比如0.25然后用以下腳本統(tǒng)計(jì)import json import numpy as np # 假設(shè) results.json 每行是 {image: B001_01.jpg, boxes: [[cls, conf, x1, y1, x2, y2], ...]} # gt.json 是真實(shí)標(biāo)注按同樣格式組織 iou_thresh 0.5 conf_thresh 0.25 miss_count, false_count, total_gt 0, 0, 0 for line in open(results.json): pred json.loads(line) gt_boxes next(g for g in gt_data if g[image] pred[image])[boxes] total_gt len(gt_boxes) valid_preds [b for b in pred[boxes] if b[1] conf_thresh] matched [False] * len(gt_boxes) for p in valid_preds: best_iou, best_idx 0, -1 for i, g in enumerate(gt_boxes): if matched[i]: continue iou compute_iou(p[2:], g[2:]) # 需要自己實(shí)現(xiàn)compute_iou if iou best_iou: best_iou, best_idx iou, i if best_iou iou_thresh: matched[best_idx] True else: false_count 1 miss_count sum(1 for m in matched if not m) print(f召回率: {1 - miss_count / total_gt:.3f}) print(f每張圖平均誤檢: {false_count / len(gt_data):.3f})這個腳本里的兩個閾值值得說明conf_thresh設(shè)為0.25是針對小樣本模型的保守值如果現(xiàn)場對誤檢容忍度低可以提高到0.4但每提高0.05都要重新統(tǒng)計(jì)漏檢率找到一個平衡點(diǎn)。IoU用0.5是COCO的通行標(biāo)準(zhǔn)但對細(xì)小劃痕0.5可能過于寬松因?yàn)榧?xì)長缺陷的框位置偏移幾個像素IoU就從0.7掉到0.3。所以軸承缺陷項(xiàng)目我建議同時統(tǒng)計(jì)IoU0.3和IoU0.5兩檔分別對應(yīng)“缺陷是否被找到”和“缺陷定位是否準(zhǔn)確”。5.3 一個值得長期堅(jiān)持的測試習(xí)慣模型上線前我最后一件固定要做的就是“壞樣本回放”。把產(chǎn)線試運(yùn)行一周內(nèi)所有誤檢和漏檢的圖片單獨(dú)存到一個目錄每周用最新模型重新跑一遍看漏檢樣本是否被新模型覆蓋。568張圖的訓(xùn)練集不可能覆蓋所有現(xiàn)場光照和紋理變化這個回放目錄就是模型的“持續(xù)記憶”。一旦發(fā)現(xiàn)某類漏檢持續(xù)出現(xiàn)就把這批圖片補(bǔ)充進(jìn)訓(xùn)練集重新訓(xùn)練。這套流程看起來簡單但它能讓你清楚知道模型邊界在哪什么時候該加數(shù)據(jù)什么時候調(diào)參數(shù)是徒勞。習(xí)慣養(yǎng)成后產(chǎn)線的檢測穩(wěn)定性和實(shí)驗(yàn)階段的mAP一樣有數(shù)。希望幫到你。本文還有配套的精品資源點(diǎn)擊獲取 軸承缺陷檢測小樣本實(shí)戰(zhàn)：YOLO訓(xùn)練到TensorRT部署全流程
簡介基于YOLO的軸承生產(chǎn)缺陷檢測數(shù)據(jù)集面向從事工業(yè)視覺檢測、目標(biāo)檢測算法學(xué)習(xí)與實(shí)戰(zhàn)的開發(fā)者。數(shù)據(jù)集圍繞軸承生產(chǎn)中的典型缺陷識別問題提供568張已標(biāo)注圖片按三類缺陷完成邊界框標(biāo)注可直接用于YOLO系列模型的訓(xùn)練、驗(yàn)證與測試。包體共1772個文件涵蓋624張jpg圖像、577個txt格式的YOLO標(biāo)注、568個xml格式的VOC標(biāo)注以及2個Python訓(xùn)練腳本和1個yaml模型配置文件可同時兼容主流檢測框架的數(shù)據(jù)格式。壓縮包大小755.12MB整體目錄結(jié)構(gòu)清晰適合作為工業(yè)質(zhì)檢場景的入門練手或課題參考目前已有337人瀏覽學(xué)習(xí)。借助這套數(shù)據(jù)與腳本使用者能夠快速搭建缺陷檢測流程從數(shù)據(jù)預(yù)處理、模型訓(xùn)練到結(jié)果可視化還可針對光照變化、拍攝角度、背景干擾等生產(chǎn)環(huán)境條件開展魯棒性調(diào)優(yōu)為實(shí)際質(zhì)量管控提供可落地的技術(shù)路徑。1. 基于YOLO的軸承生產(chǎn)缺陷檢測只有568張圖和三類標(biāo)簽?zāi)軗纹鹨粭l產(chǎn)線嗎軸承表面缺陷檢測的殘酷現(xiàn)實(shí)是產(chǎn)線上真正能收集到的缺陷樣本常常就是幾百張而不是上萬張。568張圖、三類缺陷這個規(guī)模在工業(yè)視覺項(xiàng)目里非常典型也最容易翻車——模型訓(xùn)練時mAP看著能到0.95一上產(chǎn)線就被良品表面紋理帶偏或者干脆把小劃傷漏掉。問題不在YOLO本身而在數(shù)據(jù)怎么劃分、類別邊界怎么定、訓(xùn)練參數(shù)怎么守紀(jì)律。這篇文章要解決的就是這件事用568張圖把YOLO訓(xùn)練成能用的軸承外觀檢測器。從數(shù)據(jù)整理、標(biāo)注格式、模型選型、訓(xùn)練參數(shù)到TensorRT加速和產(chǎn)線驗(yàn)證全部按可復(fù)現(xiàn)的步驟寫。適合手里有少量缺陷圖、想快速驗(yàn)證方案可行性或者已經(jīng)訓(xùn)過但現(xiàn)場表現(xiàn)不穩(wěn)定的工程師。2. 把568張圖整理成YOLO能吃的數(shù)據(jù)集標(biāo)注、劃分與增強(qiáng)2.1 三類缺陷怎么定類別邊界是第一個坑標(biāo)題里寫了“類別三類”但具體哪三類要自己定。常見做法是按軸承外觀缺陷的發(fā)生頻率和檢測價值來選劃痕scratch、麻點(diǎn)pit、剝落spall。這三類在視覺特征上差異明顯——劃痕是線狀、方向隨機(jī)麻點(diǎn)是點(diǎn)狀、邊緣清晰剝落是片狀、紋理粗糙偏暗。選這三類的另一個好處是標(biāo)注維度一致不會出現(xiàn)一類缺陷邊界模糊到標(biāo)注工都拿不準(zhǔn)的情況。類別邊界一旦確定最忌諱的就是“這是個很大的劃痕但好像也算剝落”。標(biāo)注的時候手一松模型學(xué)到的邊界就糊了。我的做法是先做一個標(biāo)注規(guī)范樣例圖每類選兩張典型圖、兩張邊界圖邊界圖標(biāo)注時遵循“占主導(dǎo)特征定類”的原則。568張圖雖然不多但每張標(biāo)注的都越干凈后面訓(xùn)練越省心。寧可三類各標(biāo)注100張左右也不要三類里某一類只有40張、另一類有360張。2.2 整理目錄結(jié)構(gòu)與標(biāo)簽一個腳本解決訓(xùn)練集劃分YOLO訓(xùn)練要求的數(shù)據(jù)組織方式很固定images目錄下放原圖labels目錄下放同名txt文件每行是“類別 cx cy w h”cx、cy、w、h都是相對圖像寬高的歸一化坐標(biāo)。568張圖我習(xí)慣按7:2:1劃分成訓(xùn)練、驗(yàn)證、測試最終目錄結(jié)構(gòu)如下dataset/ ├── images/ │ ├── train/ # 398張 │ ├── val/ # 113張 │ └── test/ # 57張 ├── labels/ │ ├── train/ │ ├── val/ │ └── test/ ├── data.yaml └── classes.txt數(shù)據(jù)劃分必須按“軸承個體”來分不能按圖隨機(jī)分。同一個軸承的多張不同角度照片如果同時出現(xiàn)在訓(xùn)練集和測試集模型其實(shí)是在“背照片”而不是“學(xué)缺陷”。這個坑我在早期項(xiàng)目里踩過當(dāng)時測試集mAP虛高到0.99換上全新軸承圖立刻掉到0.7以下。劃分腳本用Python寫關(guān)鍵是先按軸承ID分組再把組列表shuffle后切分import os import random import shutil from collections import defaultdict img_root raw_images # 原始圖片目錄文件名前綴為軸承ID如 B001_01.jpg out_root dataset train_ratio, val_ratio 0.7, 0.2 # 1. 按軸承ID分組 groups defaultdict(list) for fname in os.listdir(img_root): if not fname.endswith(.jpg): continue bearing_id fname.split(_)[0] groups[bearing_id].append(fname) items list(groups.items()) random.seed(42) random.shuffle(items) # 2. 按組劃分 train_files, val_files, test_files [], [], [] for idx, (bid, files) in enumerate(items): if idx int(len(items) * train_ratio): train_files files elif idx int(len(items) * (train_ratio val_ratio)): val_files files else: test_files files # 3. 復(fù)制圖像并生成同名標(biāo)簽txt標(biāo)簽數(shù)據(jù)假設(shè)已放在raw_labels目錄 for split, files in [(train, train_files), (val, val_files), (test, test_files)]: os.makedirs(f{out_root}/images/{split}, exist_okTrue) os.makedirs(f{out_root}/labels/{split}, exist_okTrue) for fname in files: shutil.copy(f{img_root}/{fname}, f{out_root}/images/{split}/{fname}) label_name fname.replace(.jpg, .txt) src_label fraw_labels/{label_name} if os.path.exists(src_label): shutil.copy(src_label, f{out_root}/labels/{split}/{label_name}) print(ftrain{len(train_files)} val{len(val_files)} test{len(test_files)})這段腳本里有兩個參數(shù)值得注意一個是random.seed(42)固定隨機(jī)種子保證每次復(fù)現(xiàn)同樣的劃分結(jié)果方便別人復(fù)現(xiàn)你的實(shí)驗(yàn)另一個是train_ratio和val_ratio的比例568張圖本身不多7:2:1是工業(yè)小樣本最常用的比例測試集57張足夠算漏檢率再少就不穩(wěn)定了。如果你的缺陷類別分布不均衡劃分后還要檢查一下每個split里的類比比例別讓訓(xùn)練集中某一類只有20張。2.3 數(shù)據(jù)增強(qiáng)小樣本的后悔藥568張圖直接訓(xùn)練YOLOv8s過擬合幾乎是必然的。增強(qiáng)策略我一般分兩層第一層是YOLO自帶的在線增強(qiáng)第二層是離線補(bǔ)充增強(qiáng)樣本。第一層優(yōu)先因?yàn)樵诰€增強(qiáng)每個epoch都隨機(jī)變化等效訓(xùn)練樣本量翻好幾倍。YOLOv8默認(rèn)的增強(qiáng)參數(shù)比較激進(jìn)對小樣本來說hsv_h0.015、hsv_s0.7、hsv_v0.4、fliplr0.5、mosaic1.0這些都夠用。離線增強(qiáng)只做一類操作模擬產(chǎn)線光照變化和表面反光。用cv2對訓(xùn)練集中的圖片做亮度±30%調(diào)整、小角度旋轉(zhuǎn)±10度、高斯模糊模擬相機(jī)輕微失焦但注意不要在缺陷區(qū)域引入形變。軸承是圓形對稱件旋轉(zhuǎn)增強(qiáng)是安全的但對有方向性的劃痕旋轉(zhuǎn)90度后視覺特征可能就不符合真實(shí)分布了這點(diǎn)要按產(chǎn)線實(shí)拍情況決定。一個容易被忽略的增強(qiáng)是Mosaic和Copy-Paste的配合。Mosaic是把四張圖拼成一張對小目標(biāo)有效但軸承缺陷通常只占幾個像素到幾十個像素Mosaic后目標(biāo)更小了。Copy-Paste則把缺陷區(qū)域隨機(jī)粘貼到無缺陷的軸承圖上這種離線增強(qiáng)對增加類別樣本量很有用前提是缺陷邊緣要帶原圖的紋理像素不要直接硬摳否則模型學(xué)到的就是“邊緣生硬缺陷”。3. 模型選型與訓(xùn)練用YOLOv8s還是更大版本關(guān)鍵看推理預(yù)算3.1 模型尺寸、損失函數(shù)與推理速度的權(quán)衡568張數(shù)據(jù)集訓(xùn)練大模型沒有意義。YOLOv8s參數(shù)量約1100萬COCO上mAP比n版高不少但速度只慢一點(diǎn)點(diǎn)是工業(yè)檢測的甜點(diǎn)位。更大的YOLOv8m或v8l在568張圖上很容易把訓(xùn)練集損失壓到接近0驗(yàn)證集表現(xiàn)反而變差因?yàn)閰?shù)容量太大把樣本噪聲也學(xué)進(jìn)去了。社區(qū)里討論的YOLO版本迭代很多從v5到v8再到各類結(jié)構(gòu)變體核心邏輯沒變損失函數(shù)由BBox回歸損失、分類損失和DFLDistribution Focal Loss組成。對小目標(biāo)、小樣本DFL的作用尤其明顯它對邊界框的定位不確定性建模輸出的是分布而不是單值這能讓模型在缺陷邊緣模糊時給出更保守的定位。T4上跑TensorRT FP16YOLOv8s 640x640輸入的單幀推理時間一般在2~3毫秒左右這個數(shù)據(jù)是多次實(shí)測的經(jīng)驗(yàn)范圍。如果檢測幀率要求25FPS即每幀預(yù)算40毫秒單張T4卡理論上能串行處理十幾路視頻流但考慮到預(yù)處理、后處理和顯存帶寬工程上8路左右是穩(wěn)妥值。這個估算后面還會細(xì)說先記住結(jié)論選YOLOv8s不是因?yàn)樗茸罡叨且驗(yàn)樗谛颖竞蛯?shí)時推理之間最平衡。3.2 訓(xùn)練命令與關(guān)鍵超參數(shù)YOLOv8訓(xùn)練用ultralytics庫。data.yaml里要寫清楚路徑、類別數(shù)和類別名。paths字段建議寫相對路徑這樣代碼拷到別的機(jī)器不用改path: dataset train: images/train val: images/val test: images/test names: 0: scratch 1: pit 2: spall訓(xùn)練命令本身不復(fù)雜但有幾個參數(shù)是必須調(diào)的。下面這條命令是我在類似小樣本項(xiàng)目里的基準(zhǔn)配置yolo detect train \ modelyolov8s.pt \ datadata.yaml \ epochs200 \ batch16 \ imgsz640 \ patience30 \ optimizerSGD \ lr00.01 \ cos_lrTrue \ mosaic1.0 \ seed42 \ projectruns/detect \ namebearing_defect參數(shù)說明epochs200看似多但配合patience30驗(yàn)證集指標(biāo)連續(xù)30輪不漲就會自動早停實(shí)際通常100輪左右收斂。batch16在顯存允許下盡量大小樣本數(shù)據(jù)集的BN統(tǒng)計(jì)量對batch size敏感batch太小比如4或8會讓BN的均值方差震蕩損失曲線抖動明顯。lr00.01配SGD是小樣本的標(biāo)準(zhǔn)開局如果你用AdamWlr0要降到0.001左右否則前幾個epoch損失容易爆掉。cos_lrTrue讓學(xué)習(xí)率余弦衰減相比階梯下降在小樣本上能多擠出一點(diǎn)精度。有個坑要提醒mosaic1.0默認(rèn)是啟用的但如果你發(fā)現(xiàn)模型在驗(yàn)證集上對完整軸承圖的檢測效果差而訓(xùn)練時loss降得很漂亮可以嘗試把mosaic在最后20個epoch關(guān)掉ultralytics支持close_mosaic10這樣的參數(shù)因?yàn)镸osaic產(chǎn)生的拼接圖與真實(shí)完整軸承圖像分布有偏差最后階段切換到真實(shí)分布能讓模型適應(yīng)更好。3.3 訓(xùn)練結(jié)果怎么看別只看mAP訓(xùn)練結(jié)束后ultralytics會輸出P、R、mAP50、mAP50-95等指標(biāo)。在568張的小數(shù)據(jù)集上我更看重R召回率因?yàn)楫a(chǎn)線上漏檢比誤檢嚴(yán)重得多——誤檢最多是人工復(fù)檢多看一眼漏檢是缺陷直接流入下游。如果R為0.93意味著100個缺陷里有7個漏掉這個數(shù)字對軸承件來說通常不可接受。怎么看訓(xùn)練日志里的loss曲線主要看box_loss和cls_loss是否同步下降。如果box_loss降得很慢而cls_loss正常說明邊界框回歸困難可能原因是標(biāo)注框不貼合缺陷邊緣如果cls_loss下不去大概率是類別邊界標(biāo)注混亂。另外val_cls_loss如果出現(xiàn)先降后升的“U型”就是過擬合的典型信號此時最優(yōu)模型往往是val損失最低點(diǎn)附近的checkpoint而不是最后一個epoch的權(quán)重。提示訓(xùn)練完成后把best.pt在test集上單獨(dú)跑一次得到的結(jié)果才是真正可以匯報的精度。val集參與了早停判斷指標(biāo)有樂觀偏差test集是完全沒見過、不參與訓(xùn)練決策的數(shù)據(jù)。這一點(diǎn)雖然基礎(chǔ)但很多項(xiàng)目為了省事跳過test集測試最后上線效果對不上實(shí)驗(yàn)數(shù)據(jù)原因就在這。4. 避坑記錄568張軸承數(shù)據(jù)訓(xùn)練YOLO的5個坑4.1 坑mAP高但某個類別幾乎全漏——類別不平衡現(xiàn)象訓(xùn)練結(jié)束后mAP50在0.9以上但查看每個類別的confusion matrix發(fā)現(xiàn)“剝落”類召回率只有0.4。原因568張圖里剝落只有不到80個實(shí)例而劃痕有300多個實(shí)例。YOLO的默認(rèn)分類損失對樣本多的類別傾斜模型學(xué)會了劃痕對剝落則傾向于保守預(yù)測。解決先統(tǒng)計(jì)每個類別的實(shí)例數(shù)量而不是圖片數(shù)量。如果剝落實(shí)例明顯偏少對這個類別做離線復(fù)制粘貼增強(qiáng)目標(biāo)是把三類實(shí)例數(shù)拉平到同一量級。另外可以給少樣本類別在loss里加權(quán)重ultralytics支持在data.yaml的names后接一個weight字段但實(shí)際操作里增強(qiáng)數(shù)據(jù)比調(diào)loss權(quán)重更穩(wěn)。4.2 坑驗(yàn)證集mAP 0.98一到產(chǎn)線就翻車——訓(xùn)練測試同源現(xiàn)象實(shí)驗(yàn)階段mAP漂亮得不像話新采一批軸承圖測試檢測率斷崖式下跌。原因數(shù)據(jù)集劃分時按圖片隨機(jī)分同一個軸承的多張照片同時進(jìn)了訓(xùn)練集和測試集。YOLO記住了軸承紋理背景測試時靠背景相似性就能“猜”出缺陷位置。解決嚴(yán)格按軸承個體分組劃分本文2.2的腳本已經(jīng)實(shí)現(xiàn)了這個邏輯。還有一個驗(yàn)證方法從測試集里隨機(jī)抽20張圖人工確認(rèn)沒有任何一張?jiān)谟?xùn)練集出現(xiàn)過把“同源”這道防線卡死在數(shù)據(jù)準(zhǔn)備階段。4.3 坑損失正常、metrics正常但檢測框偏移半個缺陷位——標(biāo)注邊界不一致現(xiàn)象loss曲線一切正常mAP50在0.88附近但可視化檢測結(jié)果發(fā)現(xiàn)預(yù)測框總是比標(biāo)注框大一圈或偏一側(cè)。原因標(biāo)注時一類框緊貼缺陷邊緣另一類框則多包含了一圈背景。YOLO回歸的是邊界框坐標(biāo)標(biāo)注的不一致會讓模型學(xué)到一個“平均框”看起來偏差不大但后續(xù)做缺陷尺寸篩選時就出問題了。解決標(biāo)注規(guī)范里明確“外接矩形緊貼缺陷最外緣”并且每張圖標(biāo)注完成后用Python腳本檢查框?qū)捀叻植籍惓捀弑鹊目騿为?dú)挑出來復(fù)核。4.4 坑小劃痕漏檢率高——目標(biāo)尺度問題現(xiàn)象形態(tài)檢測單獨(dú)看都很好但5像素寬、40像素長的細(xì)微劃痕10個里漏掉4個。原因YOLO輸入imgsz640時這種細(xì)長缺陷在特征圖上的響應(yīng)非常弱。就算訓(xùn)練loss包含了DFL目標(biāo)太小的情況下特征層根本沒有足夠的語義信息來區(qū)分劃痕和正常紋理。解決第一imgsz從640提到960或1280小目標(biāo)AP會明顯上升代價是推理時間增加需要重新評估TensorRT路數(shù)第二使用SAHI這類切片推理工具把大圖切成小塊分別檢測第三檢查是否開啟了多尺度訓(xùn)練scale0.5這種設(shè)置讓模型見識更多尺寸變化對小目標(biāo)魯棒性有幫助。4.5 坑過擬合到增強(qiáng)樣本本身——增強(qiáng)參數(shù)過猛現(xiàn)象訓(xùn)練集loss降到0.05驗(yàn)證集loss停在0.3不再下降檢測結(jié)果對光照異常敏感。原因hsv增強(qiáng)參數(shù)調(diào)太猛比如hsv_v0.8導(dǎo)致訓(xùn)練圖中出現(xiàn)大量現(xiàn)實(shí)中不會出現(xiàn)的極端亮度模型花了大量容量去擬合這些失真圖。解決回顧增強(qiáng)參數(shù)把hsv_v回調(diào)到0.4以下關(guān)閉mosaic的最后10個epoch。小樣本數(shù)據(jù)集的增強(qiáng)原則是“輕微擾動模擬真實(shí)變化”而不是“大力出奇跡”。每次改增強(qiáng)參數(shù)只改一個變量跑一輪訓(xùn)練看驗(yàn)證集趨勢不要一次性把一堆參數(shù)全改了那樣出了問題根本定位不到原因。5. 從驗(yàn)證到產(chǎn)線TensorRT加速與實(shí)時路數(shù)估算5.1 TensorRT FP16下的吞吐量與路數(shù)估算訓(xùn)練的模型要跑到產(chǎn)線上PyTorch推理的速度通常不夠用。我的流程是先用ultralytics導(dǎo)出ONNX再用TensorRT轉(zhuǎn)成FP16 engine。完成轉(zhuǎn)換后在同一張T4卡上實(shí)測單幀推理耗時。下面這張表是經(jīng)驗(yàn)參考值實(shí)際值會因輸入分辨率和顯存占用略有浮動模型輸入分辨率TensorRT FP16單幀耗時25FPS單路預(yù)算T4可支撐路數(shù)YOLOv8s640x6402~3ms40ms8~10路YOLOv8n640x6401~1.5ms40ms12~16路YOLOv8s1280x12805~8ms40ms4~6路估算方法很簡單單幀推理耗時如果只有2.5ms而每路25FPS要求每幀40ms內(nèi)處理完理論上GUP可以在一路等待下一幀的時間內(nèi)串行處理多個幀。但實(shí)際還要算上圖像解碼、預(yù)處理、后處理NMS的耗時以及顯存帶寬競爭。因此我從來不用理論峰值排產(chǎn)線一般留50%余量T4640分辨率按8路設(shè)計(jì)。如果現(xiàn)場需要更多路數(shù)就把模型換成YOLOv8n或者把輸入分辨率降到480x480但檢測精度要重新驗(yàn)證。5.2 產(chǎn)線驗(yàn)證腳本統(tǒng)計(jì)漏檢率和誤檢率TensorRT跑起來的模型最終要用一個最小驗(yàn)證腳本來確認(rèn)工程指標(biāo)。這里的重點(diǎn)不是寫代碼而是定義好什么算漏檢、什么算誤檢。我習(xí)慣的做法是設(shè)定一個IoU閾值比如0.5和一個置信度閾值比如0.25然后用以下腳本統(tǒng)計(jì)import json import numpy as np # 假設(shè) results.json 每行是 {image: B001_01.jpg, boxes: [[cls, conf, x1, y1, x2, y2], ...]} # gt.json 是真實(shí)標(biāo)注按同樣格式組織 iou_thresh 0.5 conf_thresh 0.25 miss_count, false_count, total_gt 0, 0, 0 for line in open(results.json): pred json.loads(line) gt_boxes next(g for g in gt_data if g[image] pred[image])[boxes] total_gt len(gt_boxes) valid_preds [b for b in pred[boxes] if b[1] conf_thresh] matched [False] * len(gt_boxes) for p in valid_preds: best_iou, best_idx 0, -1 for i, g in enumerate(gt_boxes): if matched[i]: continue iou compute_iou(p[2:], g[2:]) # 需要自己實(shí)現(xiàn)compute_iou if iou best_iou: best_iou, best_idx iou, i if best_iou iou_thresh: matched[best_idx] True else: false_count 1 miss_count sum(1 for m in matched if not m) print(f召回率: {1 - miss_count / total_gt:.3f}) print(f每張圖平均誤檢: {false_count / len(gt_data):.3f})這個腳本里的兩個閾值值得說明conf_thresh設(shè)為0.25是針對小樣本模型的保守值如果現(xiàn)場對誤檢容忍度低可以提高到0.4但每提高0.05都要重新統(tǒng)計(jì)漏檢率找到一個平衡點(diǎn)。IoU用0.5是COCO的通行標(biāo)準(zhǔn)但對細(xì)小劃痕0.5可能過于寬松因?yàn)榧?xì)長缺陷的框位置偏移幾個像素IoU就從0.7掉到0.3。所以軸承缺陷項(xiàng)目我建議同時統(tǒng)計(jì)IoU0.3和IoU0.5兩檔分別對應(yīng)“缺陷是否被找到”和“缺陷定位是否準(zhǔn)確”。5.3 一個值得長期堅(jiān)持的測試習(xí)慣模型上線前我最后一件固定要做的就是“壞樣本回放”。把產(chǎn)線試運(yùn)行一周內(nèi)所有誤檢和漏檢的圖片單獨(dú)存到一個目錄每周用最新模型重新跑一遍看漏檢樣本是否被新模型覆蓋。568張圖的訓(xùn)練集不可能覆蓋所有現(xiàn)場光照和紋理變化這個回放目錄就是模型的“持續(xù)記憶”。一旦發(fā)現(xiàn)某類漏檢持續(xù)出現(xiàn)就把這批圖片補(bǔ)充進(jìn)訓(xùn)練集重新訓(xùn)練。這套流程看起來簡單但它能讓你清楚知道模型邊界在哪什么時候該加數(shù)據(jù)什么時候調(diào)參數(shù)是徒勞。習(xí)慣養(yǎng)成后產(chǎn)線的檢測穩(wěn)定性和實(shí)驗(yàn)階段的mAP一樣有數(shù)。希望幫到你。本文還有配套的精品資源點(diǎn)擊獲取