使用yolov5訓練自己的目標檢測數據集

yolov4出來后不久，又出現了yolov5，沒有論文。雖然作者沒有放上和yolov4的直接測試對比，但在COCO數據集的測試效果還是很可觀的。很多人考慮到YOLOv5的創新性不足，對算法是否能夠進化，因此yolov5的名頭仍有爭議。但是既然github上有如此多的star，說明還是受到大多數人認可的一個工程。

并且yolov5是十分容易上手使用的一個目標檢測開源項目，只需要改幾個文件，配置好自己的數據集格式和目錄結構，即可開始訓練，對新手十分友好。可以先把訓練跑起來，在根據需要去細扣其中的模塊設計。

并且，yolov5項目一直在更新維護，除了日常處理bug，還有更優算法的版本（如激活函數的更新）更迭和新的訓練可視化工具wandb的加入等，這也是十分難得的。

今天，筆者就來介紹一下自己第一次使用yolov5訓練自己的數據集的過程。

工程項目準備

PyTorch等其他依賴的安裝這里就不說了，網上有很多教程。

首先從yolov5 github上克隆下代碼：

git clone https://github.com/ultralytics/yolov5.git

然后從github下載預訓練權重(.pt文件)，沒有百度云版本，這里筆者提供下載好的權重文件放到云盤（密碼: afp3）供不方便的朋友直接下載。下載好后放到yolov5/weights文件夾下。

數據集準備

yolov5項目所需的數據集文件路徑格式如下：

---yolov5
---yolo_data---train---images---labels---val---images---labels---data.yaml

數據集文件夾放在與工程文件同級的目錄下。其中包含兩個子文件夾和一個配置文件data.yaml，兩個子文件夾分別存放訓練樣本和驗證樣本，注意images和labels文件夾的名字不要改。

先說一下data.yaml配置文件內的內容

train: ../yolo_data/train
val: ../yolo_data/valnc: [你的檢測數據集的類別數]
names: [你的檢測數據集的類別名稱列表]

共有四個鍵值對，train，val，nc，names，分別對應如下表：

鍵	值
train	你的訓練數據的文件夾路徑
val	你的驗證數據的文件夾路徑
nc	你的檢測數據集的類別數（如COCO數據集中nc=80）
names	你的檢測數據集的類別名稱列表（如：[‘plane’, ‘car’, ‘dog’, …]）

然后說一下數據集格式的要求：

在train/val目錄下要有images和labels兩個目錄，分別存放圖像和標簽，其中的文件名要相互對應，分別為.jpg格式和.txt格式，即：

---train---imagesaaaa.jpg	...---labelsaaaa.txt...

其中圖像不必多說，放入你的數據集圖像即可，而標簽label目錄下的txt文件要求格式如下：

類別索引           x 	      y				w				h
0 			0.53857421875 0.8779296875 0.158203125 0.06184895833333333

第一列是圖像中每個object的類別索引，注意這里要與data.yaml配置文件中的names相互對應。后面四個列就是該object的檢測框坐標了，這里采用的是經過歸一化的中心長寬xywh格式，如果你的數據集原生不是這種格式（如對角坐標xyxy）的話，請注意調整。

開始訓練

直接運行yolov5/train.py即可，注意配置好參數

python train.py \
--data ../yolo_data/data.yaml \
--cfg models/yolov5s.yaml \
--weights weights/yolov5s.yaml \
--batch-size 8 \
--epochs 100

yolov5最近一次更新增加了wandb（weights and biases，其官網，直接訪問不太穩定，建議科學上網訪問）訓練可視化工具，沒有安裝的話需要先安裝一下，直接pip安裝即可

pip install wandb

第一次使用的話會需要注冊賬號，訪問站點https://wandb.ai/authorize，拿到API Key，復制到終端中即可，首次綁定API Key賬號后再訓練就不在需要輸入API Key了，直接將訓練可視化結果送到你的wandb賬號，可在該站點登錄查看。

不僅有訓練過程可視化，loss曲線等

甚至還有訓練過程中系統的信息的展示：

還是非常簡單實用的一個訓練可視化工具。

訓練完成后再yolov5/runs/train/exp文件夾下可查看一些結果PR曲線，F1曲線等，還有一些樣本可視化的結果。

可以看到效果還是可以的（我這個數據集只有豬仔一個類 ^^可愛小豬仔）。

測試

訓練完成后可以運行yolov5.detect.py文件，指定一個文件夾，測試其中的圖像。

python detect.py \
--weights runs/train/exp/weights/best.pt \
--source ../images/test \
--save-txt 

如果設置打開–save-txt，則會保存結果文件。

測試結果會保存在yolov5/runs/detect內。txt結果文件則會保存在yolov5/runs/detect/labels內。

使用yolov5訓練自己的數據集就先介紹到這里，如果有問題歡迎留言討論。