【YOLO v5】訓練VisDrone數據集

參考：https://blog.csdn.net/qq_42932308/article/details/115217935

 

1. 數據集標簽文件的格式轉化

visdrone數據集默認標簽格式為.txt，轉化為.xml文件后，修改參數，就可在使用yolov5項目直接訓練。

說明： 創建數據集標簽轉換腳本visdrone2yolo.py，修改root_dir，將train、val、test文件的labels文件轉換得到annotations_voc文件，並將該文件復制到訓練數據集目錄下。

 

2.修改配置文件

修改 data/VisDrone.yaml

train: ../datasets/images/train  # train images (relative to 'path')  6471 images
val: ../datasets/images/val  # val images (relative to 'path')  548 images
test: ../datasets/images/test  # test images (optional)  1610 images

# Classes
nc: 10  # number of classes
names: ['pedestrian', 'people', 'bicycle', 'car', 'van', 'truck', 'tricycle', 'awning-tricycle', 'bus', 'motor']

修改 models/yolo5l.yaml

nc: 10  # number of classes
depth_multiple: 1.0  # model depth multiple
width_multiple: 1.0  # layer channel multiple

3. 下載預訓練模型

Releases · ultralytics/yolov5 (github.com)

python detect.py --weights yolov5s.pt  # P5 models
                           yolov5m.pt
                           yolov5l.pt
                           yolov5x.pt

將下載得到的yolov5l.pt復制到weights文件夾下。

開始訓練

python train.py --data data/VisDrone.yaml --cfg models/yolov5s.yaml --weights weights/yolov5s.pt --batch-size 1 --epochs 30


python train.py --data data/coco128.yaml --cfg models/yolov5s.yaml --weights weights/yolov5s.pt --batch-size 1  --epochs 10

 

4.訓練過程

配置：GTX1050ti 4G

數據集：訓練集6471，驗證集548

模型：yolov5s

batch-size = 1

結果：30min/epoch

問題1 ：RuntimeError: [enforce fail at ..\c10\core\CPUAllocator.cpp:79] data. DefaultCPUAllocator: not enough memory: you tried to allocate 105436960 bytes

驗證時，顯示內存不足。

解決：嘗試減少訓練集：使用4000張圖片訓練（刪除9999966_00000_d_0000034   ~  9999999_00887_d_0000407）

 

 

5. 檢測測試集

python detect.py --source D:\WorkSpace\deeplearning\yolov5\datasets\VisDrone\images\test\ --weights D:\WorkSpace\deeplearning\yolov5\yolov5\runs\train\exp\weights\best.pt