返回 导航

Python / AI

hangge.com

YOLO - 模型训练入门教程3(模型训练与使用)

作者:hangge | 2024-04-25 08:35
    在前两篇文章中,我们已经搭建了 YOLO 模型的训练环境(点击查看)并准备好了数据集(点击查看)。本文我将介绍如何进行模型训练,并使用训练好的模型进行目标检测。

三、模型训练

1,修改配置文件

(1)首先在项目的 data 路径下找到 coco128.yaml 文件,复制到 own_datas 路径下,改名为 animal.yaml

(2)打开 animal.yaml 文件,首先需要修改如下三行内容:
  • 第一行要注释掉
  • 第二行和第三行的存放有训练图片和标签的目录相对路径
#path: ../datasets/coco128 # dataset root dir
train: own_datas/images/train # train images (relative to 'path') 128 images
val: own_datas/images/train # val images (relative to 'path') 128 images

(3)然后根据情况修改如下类别内容:
names:
  0: cat
  1: doge

(4)最后还可以把下面这一行给注释掉,因为我们不使用 coco128 数据集。当然不注释也行,没什么影响。
#download: https://ultralytics.com/assets/coco128.zip

(5)接着在项目 models 路径下找到 yolov5s.yaml 文件,同样复制到 own_datas 路径下:

(6)打开 yolov5s.yaml 文件,根据情况修改 nc 类别数:

2,开始训练

(1)打开项目的 train.py 文件,根据情况修改如下三行内容:
parser.add_argument("--weights", type=str, default=ROOT / "yolov5s.pt", help="initial weights path")
parser.add_argument("--cfg", type=str, default=ROOT / "own_datas/yolov5s.yaml", help="model.yaml path")
parser.add_argument("--data", type=str, default=ROOT / "own_datas/animal.yaml", help="dataset.yaml path")

(2)还有默认的 epoch 训练迭代次数是 100,大家可以根据情况适当增加或者减少试试效果:
parser.add_argument("--epochs", type=int, default=100, help="total training epochs")

(3)然后运行 train.py 开始训练,训练结束后控制台会输出训练结果信息:

(4)而在项目的 runs\train\exp 文件夹中可以查看所有训练相关的信息:

(5)其中 weights 文件夹里面的 best.pt 文件是训练好的模型:

四、模型测试使用

(1)我们将需要检测的图片(或者视频文件)放到 own_datas\images\test 文件夹下:

(2)然后打开 detect.py 文件,根据情况修改如下两行代码:
parser.add_argument("--weights", nargs="+", type=str, default=ROOT / "runs/train/exp/weights/best.pt", help="model path or triton URL")
parser.add_argument("--source", type=str, default=ROOT / "own_datas/images/test", help="file/dir/URL/glob/screen/0(webcam)")

(3)最后运行 detect.py,结束后可以在 runs\detect\exp 文件夹里查看结果:

(4)随便打开其中一张,可以看到目标已经被成功识别并标注出来了:
评论

全部评论(0)

回到顶部