参考博客:
目标检测—教你利用yolov5训练自己的目标检测模型
YOLOv5跑通VOC2007数据集
下载yolov5源代码
Python 3.8 or later with all requirements.txt dependencies installed, including torch>=1.7. To install run:
根据官网的要求,在pycharm的控制台安装依赖文件。
pip install -r requirements.txt
【数据集的制作】VOC2007数据集格式的转换(voc2yolo)与划分
这里我们使用上次制作的VOC2007数据集。
├── data:如果是训练自己的数据集的话,那么就需要修改其中的yaml文件。但是自己的数据集不建议放在这个路径下面,而是建议把数据集放到yolov5项目的同级目录下面。
根据前面介绍的,我们把数据集放到yolov5项目的同级目录下面。
从官网下载预训练权重,我们总不可能从头开始训练我们所有的参数。根据你任务的要求,来选择是使用大模型还是小模型。这里选择了yolov5s.pt
在源代码同级目录下新建weights文件夹,放入我们下载好的yolov5s.pt。
有了预先训练的模型和数据集,就可以开始训练YOLOV5对象检测模型了。这需要修改两个YAML文件中的参数。
一个是data目录中对应的yaml文件,一个是model目录文件中对应的yaml文件。
修改数据集中的类
复制一份yaml文件,并命名为自己的文件,
train: P:/Project_Python/yolov5-voc/VOCdevkit/VOC2007/train.txt
val: P:/Project_Python/yolov5-voc/VOCdevkit/VOC2007/val.txt
test: P:/Project_Python/yolov5-voc/VOCdevkit/VOC2007/test.txt
# number of classes
nc: 20
# class names
names: [ 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus',
'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse',
'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor' ]
修改模型中的类别数目,voc2007数据集一共有20个类。
主要修改下列参数:
parser.add_argument('--weights', type=str, default=ROOT / 'weights/yolov5s.pt', help='initial weights path')
parser.add_argument('--cfg', type=str, default=ROOT / 'models/yolov5s_voc.yaml', help='model.yaml path')
parser.add_argument('--data', type=str, default=ROOT / 'data/my_VOC.yaml', help='dataset.yaml path')
--weights:初始化的权重文件的路径地址
--cfg:模型yaml文件的路径地址
--data:数据yaml文件的路径地址
修改我们的训练轮数和喂入批次文件的多少
可以看到,这里我训练了50轮。
parser.add_argument('--epochs', type=int, default=50, help='total training epochs')
parser.add_argument('--batch-size', type=int, default=8, help='total batch size for all GPUs, -1 for autobatch')
如果显存不够的话,可以修改下面的缺省值。
parser.add_argument('--workers', type=int, default=8, help='maximum number of dataloader workers')
在pycharm终端输入:
tensorboard --logdir=runs/train
训练后生成一个run文件夹,在train中会有网络学习到的参数。
打开detect.py文件,修改下列参数:
parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'weights/yolov5s.pt', help='model path(s)')
parser.add_argument('--source', type=str, default=ROOT / 'VOCdevkit/VOC2007/images/000032.jpg', help='file/dir/URL/glob, 0 for webcam')
parser.add_argument('--data', type=str, default=ROOT / 'data/my_VOC.yaml', help='(optional) dataset.yaml path')
修改参数后运行,会在run文件夹中生成一个detect文件夹,里面有我们保存的图片。
parser.add_argument('--source', type=str, default=ROOT / 'VOCdevkit/VOC2007/images/000000.jpg', help='file/dir/URL/glob, 0 for webcam')
parser.add_argument('--source', type=str, default=ROOT / '0', help='file/dir/URL/glob, 0 for webcam')
感兴趣自己试一下吧。