YOLO - 模型训练入门教程3(模型训练与使用)
作者:hangge | 2024-04-25 08:35
三、模型训练
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)