檢測實(shí)戰(zhàn):YOLO工程與PyTorch訓(xùn)練推理全流程解析)
簡介面向水下生物目標(biāo)檢測場景這份基于Python與PyTorch的深度學(xué)習(xí)資源包整合了YOLO模型訓(xùn)練與推理所需的數(shù)據(jù)集、腳本及預(yù)訓(xùn)練權(quán)重適合有一定深度學(xué)習(xí)基礎(chǔ)、希望快速上手目標(biāo)檢測項(xiàng)目的開發(fā)者。資源共1830個(gè)文件壓縮包112.1MB其中910張JPG水下生物圖像、448個(gè)XML標(biāo)注和453個(gè)TXT標(biāo)簽構(gòu)成可直接使用的YOLO格式數(shù)據(jù)集3個(gè)Python腳本覆蓋數(shù)據(jù)劃分、模型訓(xùn)練與PyQt可視化識(shí)別另含訓(xùn)練好的pt模型權(quán)重與運(yùn)行日志。已有111人學(xué)習(xí)。按01、02、03順序運(yùn)行腳本即可完成從數(shù)據(jù)集格式轉(zhuǎn)換、模型訓(xùn)練到界面加載圖片檢測的完整流程若僅需調(diào)用已訓(xùn)練模型直接運(yùn)行03pyqt.py即可彈出可視化界面進(jìn)行實(shí)時(shí)識(shí)別。對于想深入理解YOLO訓(xùn)練流程或快速搭建水下檢測演示系統(tǒng)的讀者是一份可操作、結(jié)構(gòu)清晰的參考方案。1. 水下生物目標(biāo)檢測這份 YOLO 工程能直接跑通訓(xùn)練和可視化識(shí)別做水下生物檢測最煩的不是模型選型而是數(shù)據(jù)集和自己標(biāo)注格式對不上。網(wǎng)上公開的水下數(shù)據(jù)集不少但大多是 VOC 或 COCO 格式要落到 YOLO 訓(xùn)練流程里得先寫轉(zhuǎn)換腳本一個(gè)字段對不上就全盤翻車。這份基于 Python 深度學(xué)習(xí)的工程把路鋪好了——數(shù)據(jù)集已經(jīng)整理好、標(biāo)注是 YOLO 格式環(huán)境裝完按順序跑三個(gè) py 文件就能完成從數(shù)據(jù)劃分到可視化檢測的完整流程。它用的是 PyTorch 安裝環(huán)境檢測部分帶 PyQt5 界面既適合剛?cè)腴T目標(biāo)檢測的學(xué)生跑通一個(gè)完整項(xiàng)目也適合需要快速出演示效果的從業(yè)者。整個(gè)工程不是黑匣子三個(gè)腳本分別負(fù)責(zé)數(shù)據(jù)劃分、模型訓(xùn)練、界面推理每一段都能單獨(dú)調(diào)試這也意味著你可以只替換數(shù)據(jù)集把整個(gè)流程復(fù)用到其他檢測場景。2. 環(huán)境配置與目錄結(jié)構(gòu)先看懂每個(gè)文件是干什么的2.1 requirement.txt 之外的硬性環(huán)境要求項(xiàng)目摘要里說得很直接代碼基于 Python PyTorch 環(huán)境環(huán)境需要自己配。常見做法是先裝 PyTorch 再裝其他依賴因?yàn)?requirement.txt 里如果直接寫 torch 的版本號(hào)默認(rèn)會(huì)裝 CPU 版訓(xùn)練速度會(huì)讓你懷疑人生。建議先到 PyTorch 官網(wǎng)按 CUDA 版本生成安裝命令再執(zhí)行 requirement.txt 的剩余依賴# 先確認(rèn) CUDA 版本再裝 PyTorch以 CUDA 11.8 為例 nvcc --version pip install torch2.0.1 torchvision0.15.2 --index-url https://download.pytorch.org/whl/cu118 # 裝剩余依賴 pip install -r requirement.txt代碼里涉及圖像讀取和 PyQt 界面顯示opencv-python 和 PyQt5 這兩個(gè)包版本沖突是常見問題。如果運(yùn)行 03pyqt.py 時(shí)報(bào) Qt 相關(guān)錯(cuò)誤直接強(qiáng)制重裝 PyQt5 的匹配版本即可。Python 版本建議 3.8 到 3.10太新的版本部分依賴包還沒有預(yù)編譯輪子會(huì)現(xiàn)場編譯 C 擴(kuò)展耗時(shí)長且容易失敗。2.2 工程目錄里的產(chǎn)物文件各是什么拿到壓縮包解壓后除了三個(gè) py 文件和數(shù)據(jù)集文件夾還能看到一串帶 events、cache 后綴的文件。這些都是訓(xùn)練過程中自動(dòng)生成的看懂它們能幫你判斷訓(xùn)練是否正常。labels.cache 是數(shù)據(jù)加載時(shí)生成的標(biāo)注緩存下次啟動(dòng)訓(xùn)練時(shí)不需要重新解析所有 txt 標(biāo)注results.csv 記錄每個(gè) epoch 的 loss、mAP、precision 等指標(biāo)val_batch0_labels.jpg 到 val_batch2_pred.jpg 是驗(yàn)證集的可視化結(jié)果左邊是標(biāo)注框右邊是預(yù)測框打開看一眼就知道模型學(xué)到什么程度了。如果你打算重新訓(xùn)練保險(xiǎn)做法是刪掉 labels.cache因?yàn)閿?shù)據(jù)集一旦改動(dòng)緩存里的信息就和實(shí)際標(biāo)注不一致訓(xùn)練時(shí)會(huì)報(bào) shape mismatch 之類的問題。type 命令在命令行依次運(yùn)行 01、02、03 三個(gè)文件即可如果只做推理驗(yàn)證直接跑 03pyqt.py 就行它默認(rèn)加載訓(xùn)練好的模型權(quán)重。3. 數(shù)據(jù)劃分與 YOLO 格式轉(zhuǎn)換01 腳本的完整邏輯拆解3.1 從原始標(biāo)注到 YOLO txt 的轉(zhuǎn)換細(xì)節(jié)運(yùn)行python 01劃分?jǐn)?shù)據(jù)集.py后腳本做三件事把標(biāo)注文件轉(zhuǎn)成 YOLO 格式的 txt、生成 train.txt 和 val.txt、生成 data.yaml。YOLO 格式的標(biāo)注是歸一化后的坐標(biāo)每行內(nèi)容為類別id x_center y_center width height四個(gè)坐標(biāo)值都在 0 到 1 之間。原始標(biāo)注如果來自 LabelImg 的 VOC XML 格式需要從 XML 里提取 object 的 name 和 bndbox 坐標(biāo)再除以圖片寬高做歸一化import xml.etree.ElementTree as ET def voc_to_yolo(xml_file, img_w, img_h, class_list): tree ET.parse(xml_file) root tree.getroot() yolo_lines [] for obj in root.iter(object): name obj.find(name).text cls_id class_list.index(name) # 類別不在列表里會(huì)直接報(bào)錯(cuò) box obj.find(bndbox) xmin float(box.find(xmin).text) ymin float(box.find(ymin).text) xmax float(box.find(xmax).text) ymax float(box.find(ymax).text) x_center (xmin xmax) / 2 / img_w y_center (ymin ymax) / 2 / img_h w (xmax - xmin) / img_w h (ymax - ymin) / img_h yolo_lines.append(f{cls_id} {x_center:.6f} {y_center:.6f} {w:.6f} {h:.6f}) return yolo_lines這段轉(zhuǎn)換邏輯里最容易出錯(cuò)的是class_list.index(name)如果圖片里有類別沒在 class_list 里定義程序會(huì)拋 ValueError 直接中斷。我一般會(huì)在腳本開頭加上一個(gè)類別檢查先把整個(gè)數(shù)據(jù)集里出現(xiàn)的所有 name 打印出來再手動(dòng)對齊 class_list避免訓(xùn)到一半才發(fā)現(xiàn)漏了類別。另外注意坐標(biāo)歸一化用的是圖片的寬高拿錯(cuò)寬高會(huì)導(dǎo)致標(biāo)注框偏移訓(xùn)練時(shí) loss 怎么調(diào)都降不下去。3.2 train.txt、val.txt 和 data.yaml 的生成邏輯數(shù)據(jù)劃分的比例沒有固定標(biāo)準(zhǔn)常見做法是 8 : 2 或 9 : 1。腳本里會(huì)隨機(jī)打亂圖片路徑列表按比例切分后分別寫入 train.txt 和 val.txt每行是一個(gè)圖片的絕對路徑。data.yaml 內(nèi)容是 YOLO 訓(xùn)練時(shí)的數(shù)據(jù)配置入口包含 train/val 路徑、類別數(shù)和類別名train: /path/to/your/dataset/train.txt val: /path/to/your/dataset/val.txt nc: 5 names: [holothurian, echinus, scallop, starfish, water_weeds]這里 nc 的值必須和 names 列表長度一致否則訓(xùn)練會(huì)報(bào)錯(cuò)。如果數(shù)據(jù)集里的類別在 names 里找不到屬于哪個(gè)類別的目標(biāo)都會(huì)被忽略或當(dāng)成背景處理這部分得仔細(xì)核對。data.yaml 的路徑建議寫絕對路徑相對路徑在 Windows 和 Linux 下的解析方式不同換個(gè)環(huán)境就要改一次。3.3 劃分后如何驗(yàn)證數(shù)據(jù)沒問題訓(xùn)練前先做一次數(shù)據(jù)校驗(yàn)?zāi)苁〔簧贂r(shí)間。我一般會(huì)拿 01 腳本生成的訓(xùn)練集路徑隨機(jī)讀取十幾張圖片和對應(yīng)的 txt 標(biāo)注把框畫在原圖上保存肉眼確認(rèn)標(biāo)注框貼不貼目標(biāo)、方向有沒有反、類別對不對。這個(gè)習(xí)慣救過我很多次之前導(dǎo)出的標(biāo)注里 xmin 和 ymin 全部是 0看了可視化才反應(yīng)過來是 XML 解析時(shí)字段名寫錯(cuò)了。4. 訓(xùn)練與可視化推理從 02train.py 到 03pyqt.py 的完整鏈路4.1 02train.py 訓(xùn)練參數(shù)怎么調(diào)下載的工程里 02train.py 的默認(rèn)參數(shù)一般對應(yīng)腳本里的默認(rèn)值核心參數(shù)是 epochs、batch-size、img-size。batch-size 的第一優(yōu)先級(jí)是顯存大小6GB 顯存跑 640x640 輸入batch-size 設(shè)在 8 到 16 之間比較穩(wěn)超出顯存直接 OOM 報(bào)錯(cuò)。img-size 決定輸入分辨率640 是速度和精度的平衡點(diǎn)水下生物有不少小目標(biāo)分辨率太低會(huì)把海參、海膽這類目標(biāo)直接糊掉預(yù)算夠就上 640 以上。訓(xùn)練啟動(dòng)后觀察終端輸出重點(diǎn)關(guān)注 loss 下降趨勢和 mAP 曲線。如果 loss 在最初的幾個(gè) epoch 沒有明顯下降先檢查學(xué)習(xí)率YOLO 系列默認(rèn)學(xué)習(xí)率是 0.01數(shù)據(jù)量小或 batch-size 小的時(shí)候可以調(diào)到 0.001。如果 GIOU loss 在訓(xùn)練后期持續(xù)震蕩說明學(xué)習(xí)率偏大或數(shù)據(jù)里有大量錯(cuò)標(biāo)框此時(shí)優(yōu)先檢查數(shù)據(jù)而不是調(diào)參數(shù)。訓(xùn)練正常結(jié)束后weights 目錄下會(huì)有 best.pt 和 last.ptbest.pt 是驗(yàn)證集上 mAP 最高的權(quán)重界面推理用的就是它。results.csv 里逐行記錄每個(gè) epoch 的指標(biāo)打開看一眼 precision 和 recall 的變化曲線能判斷模型有沒有過擬合或欠擬合。4.2 03pyqt.py 界面推理的工作過程03pyqt.py 做的事情本質(zhì)上是加載模型權(quán)重、讀入圖片、前向推理、畫框顯示。PyQt 界面上的加載圖片按鈕調(diào)用 QFileDialog 選擇圖片文件檢測按鈕把圖片傳給模型推理再把結(jié)果轉(zhuǎn)換回 QImage 顯示在窗口上。對應(yīng)核心邏輯如下from PyQt5.QtWidgets import QFileDialog, QLabel from PyQt5.QtGui import QImage, QPixmap import cv2 import torch class DetectWindow: def __init__(self, model_path): self.model torch.hub.load(ultralytics/yolov5, custom, pathmodel_path) self.model.conf 0.25 # 置信度閾值低于這個(gè)值的檢測框會(huì)被過濾 self.model.iou 0.45 # NMS 的 IoU 閾值重疊框去重 def load_image(self): fname, _ QFileDialog.getOpenFileName(None, 選擇圖片, , 圖片文件 (*.jpg *.png)) self.img cv2.imread(fname) self.img cv2.cvtColor(self.img, cv2.COLOR_BGR2RGB) def detect(self): results self.model(self.img) rendered results.render()[0] h, w, ch rendered.shape bytes_per_line ch * w qimg QImage(rendered.data, w, h, bytes_per_line, QImage.Format_RGB888) self.label.setPixmap(QPixmap.fromImage(qimg))置信度閾值 conf 的取值直接影響識(shí)別結(jié)果的誤檢和漏檢。閾值設(shè)太低0.1 以下會(huì)框出一堆背景雜物設(shè)太高0.5 以上則小目標(biāo)和水下模糊目標(biāo)容易被漏掉。水下生物數(shù)據(jù)集里背景復(fù)雜度高我會(huì)先把閾值設(shè)成 0.3 跑一遍看結(jié)果再往兩個(gè)方向微調(diào)。PyQt 界面里如果要加一個(gè)置信度滑條改動(dòng)也不大核心就一行self.model.conf slider.value() / 100。4.3 訓(xùn)練產(chǎn)物和界面推理的關(guān)系跑通整個(gè)流程后你會(huì)發(fā)現(xiàn)02train.py 訓(xùn)練出的 best.pt 就是 03pyqt.py 的輸入。如果訓(xùn)練中斷過或沒跑完weights 目錄下是空的界面推理會(huì)報(bào)模型文件不存在的錯(cuò)誤。摘要里提到直接運(yùn)行 03pyqt.py 即可調(diào)用已訓(xùn)練好的模型說明壓縮包自帶權(quán)重如果想重新訓(xùn)練又不破壞原有權(quán)重記得先把 weights 目錄備份一份。5. 避坑與常見問題排查四個(gè)高頻翻車現(xiàn)場5.1 訓(xùn)練到一半顯存爆了現(xiàn)象02train.py 運(yùn)行后終端報(bào)CUDA out of memory程序直接退出顯卡看著沒跑多少 epoch 就撐不住了。 原因batch-size 設(shè)太大輸入分辨率高導(dǎo)致中間特征圖占滿顯存或者同時(shí)跑著其他程序占用顯存。 解決按顯存大小設(shè)置 batch-size6GB 顯存從 batch-size 8 起步往下調(diào)直至不報(bào)錯(cuò)。也可以打開 YOLOv5 源碼中的--cache參數(shù)把數(shù)據(jù)緩存到內(nèi)存減少顯存占用或使用梯度累積功能讓顯存占用折半。5.2 labels.cache 文件導(dǎo)致的報(bào)錯(cuò)現(xiàn)象改動(dòng)數(shù)據(jù)集后直接運(yùn)行 02train.py報(bào)assertion error: shape mismatch或數(shù)據(jù)加載卡住不動(dòng)。 原因工程里殘留了之前訓(xùn)練生成的 labels.cache 緩存新數(shù)據(jù)集標(biāo)注內(nèi)容和緩存不一致加載邏輯讀取到舊的標(biāo)注信息。 解決刪掉數(shù)據(jù)集目錄下的 labels.cache 文件再重新訓(xùn)練。這是個(gè)很玄學(xué)的問題因?yàn)閳?bào)錯(cuò)信息不會(huì)直說緩存過期如果你排查了很久找不到原因多半是這里。5.3 中文路徑導(dǎo)致圖片讀取失敗現(xiàn)象03pyqt.py 加載圖片時(shí)界面無反應(yīng)控制臺(tái)報(bào)Could not find a writer for the specified extension或 cv2.imread 返回 None。 原因OpenCV 的 imread 函數(shù)對中文路徑支持不好Windows 系統(tǒng)下尤其明顯文件路徑里有中文就會(huì)讀取失敗。 解決把整個(gè)工程路徑和數(shù)據(jù)路徑都改成純英文圖片文件名也不要帶中文。我習(xí)慣把工程放在D:\projects\underwater_detect\這類純英文目錄下省去一堆潛在問題。5.4 PyQt 界面啟動(dòng)卡死或閃退現(xiàn)象運(yùn)行 03pyqt.py 后界面彈不出來或點(diǎn)擊檢測按鈕時(shí)界面卡住變成白屏無響應(yīng)。 原因檢測邏輯是同步執(zhí)行的圖片推理期間主線程被阻塞UI 無法刷新或者 PyQt5 和 OpenCV 的 DLL 沖突導(dǎo)致加載失敗。 解決推理邏輯放到 QThread 線程里運(yùn)行時(shí)用 QTimer 定時(shí)刷新界面狀態(tài)。如果是 DLL 沖突重新安裝匹配版本的 PyQt5或者把 cv2 的顯示邏輯改成 PyQt 的 QImage 轉(zhuǎn)換方式。6. 進(jìn)階用法把單張圖片檢測擴(kuò)展成批量推理和實(shí)時(shí)檢測跑通 PyQt 界面后整個(gè)工程的能力邊界就清楚了——它解決的是單張圖片的交互式檢測。但實(shí)際項(xiàng)目里遇到的往往是批量和視頻流場景比如水下機(jī)器人傳回一段錄像需要逐幀分析或者攝像頭實(shí)時(shí)畫面需要持續(xù)檢測。把 03pyqt.py 里的模型推理部分抽出來稍微改造就能擴(kuò)展成批量推理腳本。核心邏輯是把讀圖片的循環(huán)改成遍歷目錄下的所有文件檢測結(jié)果統(tǒng)一輸出到指定文件夾import torch import cv2 import glob import os model torch.hub.load(ultralytics/yolov5, custom, pathweights/best.pt) model.conf 0.3 img_dir test_images/ out_dir test_results/ os.makedirs(out_dir, exist_okTrue) for img_path in glob.glob(os.path.join(img_dir, *.jpg)): img cv2.imread(img_path) results model(img) # results.xyxy[0] 是每個(gè)檢測框的坐標(biāo)、置信度和類別id rendered results.render()[0] fname os.path.basename(img_path) cv2.imwrite(os.path.join(out_dir, fname), rendered) print(fdone: {fname}, detections: {len(results.xyxy[0])})模型推理的結(jié)果保存在results.xyxy[0]它是一個(gè) N 行 6 列的 Tensor每行對應(yīng)一個(gè)檢測框前四列是 x1、y1、x2、y2 坐標(biāo)第五列是置信度第六列是類別索引。批量跑完后統(tǒng)計(jì)所有圖片的檢測框數(shù)量可以大致判斷模型在不同光照、不同水質(zhì)條件下的穩(wěn)定性。這個(gè)腳本在工程里可以直接加入新的 py 文件不影響原有的三個(gè)文件。驗(yàn)證模型效果時(shí)還有一個(gè)習(xí)慣值得養(yǎng)成把檢測結(jié)果按置信度分桶統(tǒng)計(jì)。比如設(shè)定 0.5、0.7、0.9 三檔分別統(tǒng)計(jì)每一檔的檢測框數(shù)量和分布。如果大部分檢測框落在低置信度區(qū)間說明模型的區(qū)分度不夠需要補(bǔ)充訓(xùn)練數(shù)據(jù)或調(diào)高 conf 閾值如果高置信度區(qū)間存在大量誤檢則說明數(shù)據(jù)里存在相似干擾物。這個(gè)分析能幫你判斷模型能不能上生產(chǎn)環(huán)境而不只是「看著差不多」。我那之后每次拿到新的檢測工程都會(huì)先跑一遍批量推理腳本把所有結(jié)果圖拼成一張大圖快速翻一遍再?zèng)Q定是調(diào)參數(shù)還是補(bǔ)數(shù)據(jù)基本不會(huì)直接進(jìn)訓(xùn)練環(huán)節(jié)。這個(gè)流程走下來翻車概率小了很多希望幫到你。本文還有配套的精品資源點(diǎn)擊獲取