[极简]pytorch版yolov5训练自己的数据集

一、准备数据集

git clone -b v5.0 https://github.com/ultralytics/yolov5.git
cd yolov5-v5.0
mkdir VOCdevkit
cd VOCdevkit
mkdir images
cd images
mkdir train
mkdir val
cd ..
mkdir labels
cd labels
mkdir train
mkdir val
cd ..
mkdir VOC2007
cd VOC2007
mkdir Annotations
mkdir JPEGImages
mkdir YOLOLabels

1、把所有图片放到JPEGImages文件夹内，所有标注文件放到Annotations文件夹内。
2、修改脚本prepare_data.py中的classes并运行该脚本

二、修改配置文件

1、打开models目录下的yolov5s.yaml，修改nc个数
2、打开data目录下的voc.yaml
2.1、注释掉这句download: bash data/scripts/get_voc.sh
2.2、把

train: ../VOC/images/train/  # 16551 images
val: ../VOC/images/val/  # 4952 images

换成

train: ./VOCdevkit/images/train/  # 16551 images
val: ./VOCdevkit/images/val/  # 4952 images

2.3、修改后面的nc、names

三、开始训练

python train.py --data data/voc.yaml --cfg models/yolov5s.yaml --weights weights/yolov5s.pt --batch-size 16 --epochs 300 --workers 4 

batch-size指批处理数，epochs指迭代次数

四、训练过程可视化

tensorboard --logdir=./runs

五、模型测试

python detect.py --source ./测试图片.jpg --weights runs/train/exp/best.pt
python detect.py --source ./测试视频.mp4 --weights runs/train/exp/best.pt

测试结果会保存在runs/detect目录下
命令后可加上目标的置信度阈值，如–conf-thres 0.4

可用以下命令统计性能

python test.py --data data/voc.yaml --weights runs/train/exp/best.pt --batch-size 16