YOLOv7训练自己的数据集
上一个博客中讲了怎么安装,并利用COCO数据集训练好的权重进行检测。这个博客讲怎么利用train.py文件来训练自己的数据集。
准备数据集
数据集的准备已经有很多教程讲过,随便一搜就有,对于标注文件来说比较重要的是标注格式,YOLOv7的会将边框进行归一化,标注文件格式如下:
种类 中心横坐标比 中心纵坐标比 宽度比 高度比
看起来就像下图中所示,第一列的3、4、5代表种类,后面的用来表示位置。
所以在训练之前需要将格式更改为对应YOLOv7的格式,转格式同样也有很多教程,是比较简单的数学计算和txt文件写入。因为不同数据集的标注格式不同,所以推荐大家直接搜索数据集名怎么转YOLO格式。
数据集存放结构一般如下图,我用的是VisDrone数据集。cache文件是在训练时生成的临时文件(比如在这一轮要训练这几张,就临时生成方便读取),直接删了没关系。
实际上YOLO支持读取的数据格式还比较灵活,在project中的coco.yaml里可以看到,有三种结构可以用,所以如果使用我所展示的方式3(这种情况是常用于数据集已经分好三类train val和test),只要注意文件名不出错就行。
新建并更改yaml文件
新建yaml文件是为了让YOLO知道你的数据集种类总数,种类名称和数据集存储位置。
其实就和上面所说的存储方式比较灵活,但是注意images和labels名称不要搞错。
nc就是种类总数,names是种类名。
以YOLOv7模型为例,打开文件cfg/trainning
底下的yolov7.yaml,可以直接改也可以新建一个你想要的名字,需要改的只有nc,改为你的数据集种类总数。YOLOv7使用了autoanchor,这使得你不用自己设置anchors。
运行train.py文件
把train.py拉到下面可以看到一系列参数设置。以YOLOv7模型为例:
初始权重–weights是yolov7.pt(YOLOv7在readme中也提供了用于迁移训练的初始权重,但实际上训练轮次多了之后这个初始权重对训练结果并没有什么影响);
–cfg和–data就改为刚刚修改好的对应的yaml文件;
–hpy如果设备支持或者数据量比较小的情况下可以试着改出较好的参数;
–epochs也是试出来的,建议先设置一个大的比如200,然后观察训练曲线,选择比较适合的轮次。
batch-size和img-size也看设备,支持就大点,不支持的话16和640也是比较好的选择。
这些参数可以在文件中直接修改,也可以在命令行中输入,比如
python train.py --device 0 --batch-size 32 --data data/custom.yaml --img 640 640 --cfg cfg/training/yolov7-custom.yaml --weights 'yolov7_training.pt' --name yolov7-custom --hyp data/hyp.scratch.custom.yaml
得到结果
数据集大且轮次多的话会运行很久很久,也许需要一整天。训练结果默认存储在runs/train里。
如果训练次数比较多,建议及时存储得到的权重结果并改名,不然一个一个看图分辨很麻烦(不要问我怎么知道的)。虽然runs里会给一些曲线,但是我感觉使用tensorboard看更方便和全面。
tensorboard的使用方法是直接在命令行中输入
tensorboard --logdir=./runs/train
这样可以可视化训练结果。