如何训练自己的模型

  1. 打开github，克隆yolov5 7.0

image.png

  1. 了解一下整个目录结构

image.png

  • data：主要是存放一些超参数的配置文件，用来配置训练集和测试集还有验证集的路径的
  • models：里面主要是一些网络构建的配置文件和函数
  • utils：存放的是工具类的函数
  • weights：放置训练好的权重参数。
  • detect.py：利用训练好的权重参数进行目标检测
  • train.py：训练自己的数据集的函数。
  • test.py：测试训练的结果的函数。
  • requirements.txt：yolov5项目的环境依赖包（pip install -r requirements.txt）
  1. 数据标注好，划分训练集和验证集，创建目录，将其放到对应图像和标签目录里面 数据如何标注请看5.2怎么打标签 标签转化请看5.3标注文件的相互转化 划分训练集和验证集请看5.4 划分数据集和验证集

image.png

  1. 获取预训练权重 进入此网站https://github.com/ultralytics/yolov5/releases，放到yolov5

image.png

  1. 修改配置文件，准备训练

image.png

复制一份VOC.yaml，命名一个跟自己模型相关的名字

image.png

写好之前创建好的目录，然后写上要训练的类别

image.png

假设你要训练的是yolov5s的模型，在models文件夹的yolov5s.yaml文件里面修改类别的数量，改为5

image.png

在train.py文件里面，修改模型的pt文件（默认在yolov5-7.0主目录下），yolov5s的yaml文件，数据集的yaml文件，训练轮数可以设置为1000，训练到不能提升会自动停下来，batch_size可以根据自己的GPU显存进行调整