YOLO-V5 训练自己的分类模型
- 1、获取官方源码
- 2、测试官方源码
- 2.1、公开数据集测试源码
- 3、源码模块解析
- 3.1、数据读取
- 4、快速开始训练自己的数据
- 4.1、准备自己的数据
- 4.2、配置训练参数
1、获取官方源码
官方代码:https://github.com/ultralytics/yolov5
更新版本:v6.2 添加了分类模块,v7.0 添加分割模块,所以至少要下载 v6.2 之后的版本
环境安装:按照下面的方式安装,
拉取特定版本的代码:参考我的博文 YOLO-V5 算法和代码解析系列(一)—— 快速开始
2、测试官方源码
下图是官方的代码结构(v6.2
),增加了一个文件夹【classify】,存放分类的代码,包括训练代码 train.py
,预测代码 predict.py
,评估代码 val.py
,
2.1、公开数据集测试源码
按照上述代码结构,进入【classify】文件夹,打开训练脚本【train.py
】。根据训练脚本最开始的注释内容,可以使用官方数据集,也可以使用自己的数据集,以及单卡和多卡的训练命令,
官方训练命令如下,
python classify/train.py --model yolov5s-cls.pt --data cifar100 --epochs 5 --img 128
运行上述命令,会自动下载所需的数据集【cifar100】,以及模型文件【yolov5s-cls.pt】。下载完成后,数据的存放结构如下图所示,包含测试集【test
】和训练集【train
】。注意,参考官方的数据存放结构,可以准备自己数据集
每一个文件夹内,存放一个类别的图片,如下图所示
训练过程如下,训练基本正常,
3、源码模块解析
3.1、数据读取
- 数据路径【data_dir/train】下的文件夹数量,用于判别训练的类别数量,具体代码如下:
nc = len([x for x in (data_dir / 'train').glob('*') if x.is_dir()]) # number of classes
yolov5-cls 的网络结构如下图所示,只用了目标检测网络的主干部分。基础模块包括:C3,Conv,其详细的结构,可以参考之前写的博文,YOLO-V5 系列算法和代码解析(四)—— 网络结构
4、快速开始训练自己的数据
按照在官方数据集上的测试过程,要想快速开始训练,只要两方面的准备:(1)准备好自己的数据集;(2)配置部分训练参数。具体的过程在以下章节陈述
4.1、准备自己的数据
根据【2.1】节的陈述,可以确认自己数据集的存放结构,每一个文件夹放一个类别的图片。【mono_classify】就是存放自己的训练数据的目录名称,也是训练时传入的数据标志,
4.2、配置训练参数
为了训练自己的数据集,首先保证官方代码能够在自己的数据集上正常训练,并且达到基本的训练效果即可,之后再进行网络调优。所以,这一章节的目的是让官方源码在自己的数据集上快速开始训练,并且能够得到基本的训练效果。为此,修改训练脚本(train.py
)的部分训练参数,具体如下,
-
模型选择,默认是
yolov5-cls
,也可以是efficient-b0,b1,b2,b3
,resnet18
,parser.add_argument('--model', type=str, default='yolov5s-cls.pt', help='initial weights path')
-
修改训练数据路径
parser.add_argument('--data', type=str, default='mono_classify', help='cifar10, cifar100, mnist, imagenet, etc.')
-
根据自己的需要,修改如下参数
parser.add_argument('--epochs', type=int, default=10) parser.add_argument('--batch-size', type=int, default=64, help='total batch size for all GPUs') parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=128, help='train, val image size (pixels)')
-
根据是否需要预训练,修改参数
parser.add_argument('--pretrained', nargs='?', const=True, default=True, help='start from i.e. --pretrained False')
配置完参数后,开启训练,过程如下,从精度来看,基本正常,达到最基本的效果,