YOLO模型YOLO（You Only Look Once）是一種高效的目標檢測模型，它將圖像劃分為網(wǎng)格，并在每個網(wǎng)格中同時預測邊界框和類別，從而實現(xiàn)端到端的實時目標檢測。與傳統(tǒng)方法相比，YOLO具有速度快、精度高、適用于實時應用的優(yōu)點，廣泛應用于視頻監(jiān)控、自動駕駛等領域。

博主最近也在學習神經(jīng)網(wǎng)絡相關知識，本期記錄一下Windows系統(tǒng)訓練和部署YOLO神經(jīng)網(wǎng)絡模型的過程。

系統(tǒng)環(huán)境：Python 3.9、PyCharm 2022

1、環(huán)境安裝

Python 3.9和PyCharm的環(huán)境安裝不做過多的描寫，我們主要介紹搭建一個Pytorch和YOLO環(huán)境搭建。

（1-1）首先搭建一個Pycharm的基本工程，Python釋器版本選擇3.9版本。

圖2 PyCharm默認工程圖

（1-2）打開終端我們安裝需要對應PyTorch和YOLO的庫文件。

圖3 ? 終端使用圖

（ GPU版本）

PyTorch的模型訓練可以選擇CPU訓練或者GPU加速，GPU加速的情況下訓練速度會比CPU快很多很多很多?。。∫虼宋覀冃枰樵兒妥约篏PU適配的PyTorch版本（本文內(nèi)容適用于N卡）。

（1-3可選）查詢CUDA版本

?nvidia-smi

查詢CUDA 版本

圖4 ? CUDA版本查看

可以看到我的CUDA版本為12.6。

（1-4）接下來根據(jù)CUDA版本我們使用pip命令來安裝PyTorch,大家可以根據(jù)官網(wǎng)推薦的版本來安裝。

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

安裝cuda12.1版本的pytorch （12.6兼容12.1）

圖5 ? ?Pytorch安裝

庫文件比較大，大概有2G+，如果遇到網(wǎng)絡問題可以多安裝幾次。

（1-5）查詢pytorch是否安裝成功

python?-c?"import torch; print(torch.cuda.is_available())"

查詢pytorch并且是否支持GPU

圖6 ? ?pytorch驗證

返回True即可代表我們可以正常使用GPU

（1-6）安裝YOLO

pip?install ultralytics

YOLO庫代碼安裝

圖7 ? ?Yolo驗證

? ? 安裝成功后輸入yolo會返回很多內(nèi)容，接下來實現(xiàn)模型訓練。

2、訓練集準備

圖8 ? ?文件存放架構圖

（2-1）按照這個路徑我們存放訓練集和驗證集，images和labels分別存儲圖像數(shù)據(jù)和標簽。train文件夾存放用于訓練的數(shù)據(jù)，val存放用于驗證的數(shù)據(jù)。

labels是YOLO格式的txt文件如圖所示：

圖9 ? ?Yolo標注格式

? ? 它表達了我們要識別的內(nèi)容在圖像中的位置信息，可以使用標注工具來制作（可以查詢Labelimg的教學）

圖10 Labelimg使用實拍

? ? （2-2）制作配置文件dataset.yaml（文件名不必相同）

train: ./data/images/train?#訓練集地址val: ./data/images/val ? ?#驗證集地址nc: 4?# 類別數(shù)量names: ['Chongdianbao',?'Red Tea',?'Pen',?'Upan'] ?# 類別名稱

dataset.yaml配置文件用于等會模型訓練

圖11 ? ?yaml配置文件

? ? （2-3）模型訓練代碼

import?torchfrom?ultralytics?import?YOLOimport?multiprocessing
def?run():? ? model = YOLO('yolov8n.pt') ?# 直接使用預訓練模型微調(diào)? ? model.train(? ? ? ? data =?'./dataset.yaml',#模型配置文件? ? ? ? epochs =?800,#總訓練輪數(shù)800輪? ? ? ? imgsz = ?480,#圖片縮放為480*480? ? ? ? batch =?20, ?#一批20張? ? ? ? name =?'yolov8_custom',#模型名字? ? ? ? device =?0, ?#使用GPU進行計算? ? ? ? val =?True? ?#訓練過程中執(zhí)行驗證 ?? ? )? ??# 評估模型? ? results = model.val()? ??print(results)

if?__name__ ==?'__main__':? ? multiprocessing.freeze_support()? ? run()

模型訓練Demo

圖12 ? ?訓練時命令行截圖

? ? 接下來命令行會實時顯示訓練過程，等待訓練結(jié)束后，可以在runs文件夾中找到我們的模型。

圖13 ? ?結(jié)果圖

圖14 ? ?模型結(jié)果

? ? weights中存放著模型，其中best.pt為最優(yōu)模型，last.pt為上一次模型。

（2-4）模型使用

from ultralytics import YOLOimport cv2# 1. 加載訓練好的模型（可以是 yolov8n.pt 或你的 best.pt）model = YOLO("best.pt")?# 2. 讀取一張圖片img = cv2.imread("test.jpg")# 3. 進行推理（返回結(jié)果列表）results = model.predict(source=img, save=False)# 4. 將檢測框畫在圖上并顯示annotated = results[0].plot()cv2.imshow("YOLO Detection", annotated)cv2.waitKey(0)cv2.destroyAllWindows()

模型使用Demo

3、測試效果

標注過程是真的痛苦啊，感覺像是電子打螺絲