超越YOLOv8？基于Gold YOLO的自定义数据集训练

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

本文仅供学习交流使用，如有侵权请联系作者删除