Gold-YOLO的出色性能是对Noah's Ark Lab团队的奉献和专业知识的证明。它不仅超越了其前身YOLOv8,还为实时目标检测设定了新标准。凭借其闪电般快速的处理能力和出色的准确性,Gold-YOLO承诺革命化一系列应用,从自动驾驶车辆到监视系统等等。
我们将了解如何从roboflow训练自定义数据集的Gold-YOLO。在本文中,我们将能够在最后检查评估,然后将权重(.pt)导出到ONNX格式。
数据集来自Roblflow:https://universe.roboflow.com/roboflow-universe-projects/construction-site-safety
官方代码链接:https://github.com/huawei-noah/Efficient-Computing/tree/master/Detection/Gold-YOLO
第1步. 将运行时更改为GPU
点击RUNTIME,然后滚动更改运行时,选择T4 GPU,这是当前免费GPU,并点击保存。
第2步. 通过点击驱动符号,挂载你的Google Drive
标记成黄色
第3步. 克隆Gold YOLO的官方代码库
!git clone https://github.com/huawei-noah/Efficient-Computing.git
第4步. 更改目录
%cd '/content/drive/MyDrive/goldyolo/Efficient-Computing/Detection/Gold-YOLO'
第5步. 安装所需库
!pip install -r requirements.txt
第6步. 从Roboflow导入数据集
在这里,我以yolov6格式导入数据集。Gold Yolo实现是基于YOLOv6的,而一些实现是从Topformer借来的。
!pip install roboflow
from roboflow import Roboflow
rf = Roboflow(api_key="85eFUaSSbJwbBvNmOGlC")
project = rf.workspace("roboflow-100").project("construction-safety-gsnvb")
dataset = project.version(2).download("mt-yolov6")
第7步. 在data.yaml文件中进行更改
在Gold YOLO中,yaml文件的结构有点不同。
train: /content/drive/MyDrive/goldyolo/Efficient-Computing/Detection/Gold-YOLO/construction-safety-2/images/train # train images
val: /content/drive/MyDrive/goldyolo/Efficient-Computing/Detection/Gold-YOLO/construction-safety-2/images/valid
test: /content/drive/MyDrive/goldyolo/Efficient-Computing/Detection/Gold-YOLO/construction-safety-2/images/test
# whether it is coco dataset, only coco dataset should be set to True.
is_coco: False
# Classes
nc: 5 # number of classes
names: ['helmet','no-helmet', 'no-vest', 'person', 'vest'] # class names
如果你看到它写到is_coco: False。因为我们的数据集基于Yolo格式,确保将其设置为False。对于Coco数据集,你需要设置为True。正确设置训练、验证和测试图像路径。
第8步. 下载Gold YOLO的预训练模型并开始训练我们自定义数据集。在这一步,我们需要提供data.yaml文件的路径和批处理大小。
# Train yolo model
!wget -O Gold_n_dist.pt "https://drive.google.com/uc?export=download&id=1nJ9x8Q1Ic-DnCs0ruBh-Zy3t9-IyO9LR"
sed -i "s/pretrained=None/pretrained='.\/Gold_n_dist.pt'/g" ./configs/gold_yolo-n.py
!python tools/train.py --batch 16 --conf configs/gold_yolo-n.py --data "/content/drive/MyDrive/goldyolo/Efficient-Computing/Detection/Gold-YOLO/construction-safety-2/data.yaml" --epoch 50 --fuse_ab --device 0
我将批处理设置为16,时代设置为50。你可以根据需求更改时代。
Results saved to runs/train/exp2
Epoch: 48 | mAP@0.5: 0.9059034791884045 | mAP@0.50:0.95: 0.5014349282982402
Epoch iou_loss dfl_loss cls_loss
49/49 0.727 0 1.204: 100% 63/63 [00:35<00:00, 1.76it/s]
Training completed in 1.069 hours.
一旦训练完成,我们将在runs文件夹中获得三个权重文件(.pt)。
权重文件
第9步. 如果你想在我们的数据集上评估自定义模型
!python tools/eval.py --data "/content/drive/MyDrive/goldyolo/Efficient-Computing/Detection/Gold-YOLO/construction-safety-2/data.yaml" --batch 16 --weights /content/drive/MyDrive/goldyolo/Efficient-Computing/Detection/Gold-YOLO/runs/train/exp2/weights/best_ckpt.pt --task val --reproduce_640_eval
确保你正确给出了权重文件和yaml路径。
第10步. 用于视频和图像推理
!sed -i "s/inferer = Inferer(source, weights, device, yaml, img_size, half)/inferer = Inferer(source, False, None, weights, device, yaml, img_size, half)/g" ./tools/infer.py
!python tools/infer.py --weights /content/drive/MyDrive/goldyolo/Efficient-Computing/Detection/Gold-YOLO/runs/train/exp2/weights/best_ckpt.pt --yaml "/content/drive/MyDrive/goldyolo/Efficient-Computing/Detection/Gold-YOLO/construction-safety-2/data.yaml" --source "/content/drive/MyDrive/goldyolo/Efficient-Computing/Detection/Gold-YOLO/ppe.mp4" --save-dir 'OutPutResults'
确保正确给出权重、yaml和图像/视频路径。
结果展示
第11步. 本文的最后一步是将best.pt转换为ONNX格式
!python deploy/ONNX/export_onnx.py --weights /content/drive/MyDrive/goldyolo/Efficient-Computing/Detection/Gold-YOLO/runs/train/exp2/weights/best_ckpt.pt --device 0 --simplify --batch 32
在这里,我们提供了权重文件路径。确保一切都在正确的目录中。希望这一步可以轻松将.pt文件转换为onnx格式。
结论
这篇教程提供了如何在Roboflow平台上训练Gold-YOLO自定义数据集的全面指南,以及运行各种任务,包括训练、评估、图像和视频推断,以及将最佳模型转换为ONNX格式。它涵盖了设置GPU环境、导入数据集、配置data.yaml文件、训练模型、评估性能以及部署实际用途的关键步骤。
· END ·
HAPPY LIFE
本文仅供学习交流使用,如有侵权请联系作者删除