使用YOLOv5训练自己的数据集模型
训练YOLOv5模型
- 下载YOLOv5源码
- 训练YOLOv5模型
- 训练自己的数据集
-
- 创建DataSet.yaml
- 组织目录
- 开始训练
- 预测
本篇主要讲述如何训练自己的数据集。
下载YOLOv5源码
在githubYOLOv5源码上选择tag V5.0下载即可。
训练YOLOv5模型
打开 requirements.txt 文件，文件内容如下：
# pip install -r requirements.txt
# base ----------------------------------------
matplotlib>=3.2.2
numpy>=1.18.5
opencv-python>=4.1.2
Pillow
PyYAML>=5.3.1
scipy>=1.4.1
torch>=1.7.0
torchvision>=0.8.1
tqdm>=4.41.0
# logging -------------------------------------
tensorboard>=2.4.1
# wandb
# plotting ------------------------------------
seaborn>=0.11.0
pandas
# export --------------------------------------
# coremltools>=4.1
# onnx>=1.8.1
# scikit-learn==0.19.2 # for coreml quantization
# extras --------------------------------------
thop # FLOPS computation
pycocotools>=2.0 # COCO mAP
运行首行 pip install -r requirements.txt 命令，可能遇到的问题在之前的博客：深度学习笔记：windows安装pycocotools中进行了描述。安装完成后运行train.py文件，可能遇到的问题在之前的博客：深度学习笔记：YOLOv5tagV5.0调试过程中的遇到的问题中进行了描述。
训练完成后运行detect.py进行预测：
训练自己的数据集
创建DataSet.yaml
官方的train.py文件中，读取的数据集使用的是coco128.yaml：
parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data.yaml path')
coco128.yaml如下:
# COCO 2017 dataset http://cocodataset.org - first 128 training images
# Train command: python train.py --data coco128.yaml
# Default dataset location is next to /yolov5:
# /parent_folder
# /coco128
# /yolov5
# download command/URL (optional)
download: https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128.zip
# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
train: ../coco128/images/train2017/ # 128 images
val: ../coco128/images/train2017/ # 128 images
# number of classes
nc: 80
# class names
names: [ 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
'hair drier', 'toothbrush' ]
主要包含：
1）train and val data
2）number of classes，
3）class names：
所以我们的DataSet.yaml也应该包含着这些参数：
# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list:
train: 自己的训练集地址
val: .自己的验证集地址
# number of classes
nc: 自定义类别数
# class names
names:自定义类别名
例如：
组织目录
目录格式如下（数据集划分代码可查看博客深度学习笔记：制作YOLOv5数据集）
│├── DataSet
││├──images
│││├──test 测试集图片
│││├──train 训练集图片
│││├──val 验证集图片
││├──labels
│││├──test 测试集标签
│││├──train 训练集标签
│││├──val 验证集标签
开始训练
将数据集路径改成自己的DataSet.yaml文件所在路径
parser.add_argument('--data', type=str, default='data/DataSet.yaml', help='data.yaml path')
预测
