[Object detection]重頭開始教你如何用YOLOv8訓練自己的資料集

泓宇
7 min readJan 15, 2024

--

YOLOv8於2023年由Ultralytics釋出(YOLOv5作者),相較於以前的版本速度及準確度都有明顯提升,以下簡單說明基本環境與訓練、預測的教學。

環境部屬

1. 首先下載相對應的CUDA,以及cuDNN

我們可以從NVIDIA官網下載,選定你的作業系統以及想要下載的版本,須注意顯卡所支援的版本

接著下載對應的cuDNN,另外提醒,cuDNN需註冊會員才能下載

把 C:\Users<username>\Downloads\cuda\bin 資料夾內檔案複製到
C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.1\bin

把 C:\Users<username>\Downloads\cuda\include 資料夾內檔案複製到
C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.1\include

把 C:\Users<username>\Downloads\cuda\lib\x64 資料夾內檔案複製到
C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.1\lib\x64

2. 安裝Python,作者使用的是Anaconda來建置虛擬環境,好處為可以在同一台電腦上安裝不同的python版本,環境之間也不會汙染,較容易維護

3. 安裝pytorch,到pytorch的官網上找到對應的版本進行安裝

安裝完後可以在CMD中打上torch.cuda.is_available(),確認是否安裝成功以及是否有啟動cuda

4. 到Github上下載YOLOv8,可以直接download整個zip檔,也可以使用git clone來下載

到目前為止,基本環境已建置完成,接著我們要開始部屬我們的資料

YOLOv8訓練

1. 資料集配置

將我們的資料切分成訓練,驗證以及測試三個部分,結構如下:

images為存放圖片的資料夾,labels則是相對應的圖片的標註檔,標註的資料須依照YOLO的格式,這邊可用LabelImg來進行。

2. yaml參數檔配置

將train,val和test的路徑改為剛剛配置好的資料夾路徑,Classes則依照自己的label所定義名稱

3. 模型選定

接著需選定一個YOLOv模型來進行訓練,官網上提供了一些基本的預訓練模型可供下載來使用

模型下載網址:https://github.com/ultralytics/ultralytics

或是也可以train from scratch,在YOLOv8的資料夾中作者也附上了各個Model的yaml檔供使用

4. 開始訓練

準備好所有資料後就可以開始訓練,訓練的方法有兩種,第一種可以參考YOLOv8官網提供的documents,直接在CMD下執行

yolo detect train data=data.yaml model=yolov8n.pt epochs=100 imgsz=640 device=0

也可以過python腳本指令進行訓練

from ultralytics import YOLO

model = YOLO('yolov8n.pt')

results = model.train(data='path yo your data.yaml',
epochs=100,
batch=64,
imgsz=640,
device=0,
)

其中device參數若未指定則會使用GPU進行訓練,device=0即使用第一張GPU卡,也可device=0,1使用2張卡來進行運算,而device=cpu即使用CPU來運算

接著直接執行train.py,訓練完的資料會儲存在runs/detect/train底下,訓練好的模型會放在weights底下,接下來我們要教如何使訓練好的模型來進行測試

5. 測試模型訓練結果

跟訓練一樣,有兩種方式進行測試,一種為YOLOv8官網提供的documents,直接在CMD下執行

yolo detect predict model=best.pt source=test/images save=true

另一種為python執行檔

from ultralytics import YOLO

model = YOLO("best.pt")

source="test/images"

result = model.predict(source, imgsz=640, stream=True, save=True)

save = True代表會將預測到的BBOX畫在影像上並儲存下來,還有其他參數可以設定,可自行到YOLOv8官網上查詢

以上即是YOLOv8的訓練以及測試教學,感謝您的觀看 😍

--

--

泓宇
泓宇

Written by 泓宇

Specialist in machine learning, image processing, computer vision, and object recognition.

Responses (1)