Windows + PyTorch + YOLOv8：小白也能玩轉(zhuǎn)AI目標(biāo)檢測

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

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

系統(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默認(rèn)工程圖

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

圖3 ? 終端使用圖

（ GPU版本）

PyTorch的模型訓(xùn)練可以選擇CPU訓(xùn)練或者GPU加速，GPU加速的情況下訓(xùn)練速度會比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)絡(luò)問題可以多安裝幾次。

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

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

查詢pytorch并且是否支持GPU

圖6 ? ?pytorch驗(yàn)證

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

（1-6）安裝YOLO

pip?install ultralytics

YOLO庫代碼安裝

圖7 ? ?Yolo驗(yàn)證

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

2、訓(xùn)練集準(zhǔn)備

圖8 ? ?文件存放架構(gòu)圖

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

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

圖9 ? ?Yolo標(biāo)注格式

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

圖10 Labelimg使用實(shí)拍

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

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

dataset.yaml配置文件用于等會模型訓(xùn)練

圖11 ? ?yaml配置文件

? ? （2-3）模型訓(xùn)練代碼

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

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

模型訓(xùn)練Demo

圖12 ? ?訓(xùn)練時命令行截圖

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

圖13 ? ?結(jié)果圖

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

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

（2-4）模型使用

from ultralytics import YOLOimport cv2# 1. 加載訓(xùn)練好的模型（可以是 yolov8n.pt 或你的 best.pt）model = YOLO("best.pt")?# 2. 讀取一張圖片img = cv2.imread("test.jpg")# 3. 進(jìn)行推理（返回結(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、測試效果

標(biāo)注過程是真的痛苦啊，感覺像是電子打螺絲