yolov5训练自己的数据集(缺陷检测分类——看了就懂系列)
Posted 凌风探梅
tags:
篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了yolov5训练自己的数据集(缺陷检测分类——看了就懂系列)相关的知识,希望对你有一定的参考价值。
这里记录一下将yolov5用于训练自己的数据集,这里v5的github地址
Environment
本文是直接 pip install -r requirement.txt直接安装成功,如果不容易成功,就一个个安装即可。如果出现报错,多关注一下torch的版本即可。
Data preparation
label的格式:
<object-class> <x_center> <y_center> <width> <height>
举个🌰:
414_735_0.txt
7 0.503 0.499 0.098 0.451
7 0.465 0.633 0.033 0.053
数据集格式
data
├──coco.yaml
├──coco128.yaml
├──defect_data
├──images
├──train
├──val
├──labels
├──train
├──val
├──…
然后制作自己的defect.yaml配置文件(直接复制coco.yaml修改或者新建一个文件将一下代码copy进去):
train: ./data/defect_data/images/train # absolute_path_train_img
val: ./data/defect_data/images/val # absolute_path_val_img
nc: 8 # 类别数量
names: [‘quexian1’, ‘quexian2’, ... ‘quexian8’]
修改models/xxx.yaml,有s, m, l, x四个模型供选择,需要修改类别nc数目;同时最好把anchors也修改为自己数据对应的anchors(需要自己提前聚类),不改也可train。
Pretrained Model
建议爬梯子下载最新的预训练权重,这个repo经常更新,可能和以前的预训练模型不适应。下载好放在weights下,在train.py里修改相应路径。
Training
配置train.py
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--weights', type=str, default='weights/yolov5s.pt', help='initial weights path')
parser.add_argument('--cfg', type=str, default='models/yolov5s.yaml', help='model.yaml path')
parser.add_argument('--data', type=str, default='data/defect.yaml', help='data.yaml path')
parser.add_argument('--epochs', type=int, default=300)
parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs')
parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='[train, test] image sizes')
parser.add_argument('--device', default='0,1,2,3', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
opt = parser.parse_args()
训练好的模型会被保存在./runs/train/exp3/weights/last.pt和best.pt
Inference
配置detect.py
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--weights', nargs='+', type=str, default='runs/train/exp3/weights/last.pt', help='model.pt path(s)')
parser.add_argument('--source', type=str, default='data/defect_data/images/val', help='source') # file/folder, 0 for webcam
parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
parser.add_argument('--device', default='0,1,2,3', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
opt = parser.parse_args()
如果需要inference 自己数据集原图大小,需要修改utils/dataset.py中
class LoadImages: # for inference
# Padded resize
img = letterbox(img0, new_shape=self.img_size)[0]
修改为:
class LoadImages: # for inference
# Padded resize
img = letterbox(img0, new_shape=(math.ceil(img0.shape[0]/32)*32, math.ceil(img0.shape[1]/32)*32))[0]
Visualization of Training
利用tensorboard进行可视化,训练开始会在主目录生成一个runs文件,利用tensorboard打开:
tensorboard --logdir=./runs
Thanks
欢迎评论,每天回复☼
以上是关于yolov5训练自己的数据集(缺陷检测分类——看了就懂系列)的主要内容,如果未能解决你的问题,请参考以下文章
yolov5训练自己的数据集(缺陷检测分类——看了就懂系列)
yolov5训练自己的数据集(缺陷检测分类——看了就懂系列)
智能生活垃圾检测与分类系统(UI界面+YOLOv5+训练数据集)
笔记2:yolov5训练自己的目标检测模型_创建并划分数据集