Yolo系列 —— 使用自制数据集训练yolo模型

Ultralytics 基于多年计算机视觉和人工智能领域的基础研究，创建了尖端、最先进的(SOTA) YOLO 模型。模型不断更新，以提高性能和灵活性，使其快速、准确且易于使用。模型在目标检测、跟踪、实例分割、图像分类和姿态估计任务中表现出色。本文记录基于自己的数据集 使用yolov11 训练自己的yolo模型的学习过程，如有错误还请指正。

Yolo系列 —— Ubuntu 安装和使用标注工具 labelImg
Yolo系列 —— 使用自制数据集训练yolo模型

一、环境准备

本次使用自制的数据集 通过yolo实现 robocup 足球场的定位点识别(课程链接)，文件目录准备如下：

.
├── predict_data
│   ├── vedio
│   │   └── person perspective.mp4
├── predict.py
├── README.md
├── results
├── runs
├── train.py
├── yolo11n.pt
├── yolo11.yaml
├── yolo_dataset
│   ├── train
│   │   ├── images
│   │   ├── labels
│   │   └── labels.cache
│   └── val
│       ├── images
│       ├── labels
│       └── labels.cache
└── yolov11_goalpost_ball.yaml

文件目录说明：

• predict_data ：验证视频，用于验证yolo视频物体识别效果；
• predict.py ：测试脚本，用于验证yolo视频物体识别效果；
• results ：测试结果保存目录；
• runs ：训练过程的过程权重文件、
• train.py ：训练脚本；
• yolo11n.pt ：yolo模型文件，有n(nano)、s(small)、m(medium)、x(extra-large)多个版本，模型官网下载；
• yolo11.yaml ：创建 conda虚拟环境 yaml文件；
• yolo_dataset ：训练集、测试集；
• yolov11_goalpost_ball.yaml ：yolo配置

二、训练与验证

1. 创建 conda 环境

创建conda环境相关操作，参考此链接

# 方式一：已有 yolo11.yaml
conda env create -f yolo11.yaml # 方式二： 确保安装了必要的依赖
pip install ultralytics
pip install torch torchvision  # 根据你的CUDA版本选择合适版本

2. 训练与测试

当图片使用labelImg 标注完后如上图作为数据集，触发 train.py 训练脚本 进行自定义数据集训练， train.py 脚本 内容如下：

from ultralytics import YOLO
import yaml
import os# 创建数据集配置文件时使用绝对路径
dataset_config = {'path': os.path.abspath('./yolo_dataset'),                  # 相对路径 转 绝对路径'train': 'train/images',                                    # 相对于 path 的 训练集的图像目录路径'val': 'val/images',                                        # 相对于 path 的 验证集的图像目录路径'names': {0: 'Ball', 1: 'Post', 2: "L", 3: "T", 4: "X"}     # 数据集的类别名称字典
}config_path = './yolov11_goalpost_ball.yaml'
with open(config_path, 'w') as f:yaml.dump(dataset_config, f)# 初始化模型（选择不同规模的模型）
model = YOLO('yolo11n.pt')  # 官方预训练模型
# 可选模型：yolov11n/s/m/l/x.pt （从轻量到大型）# 训练参数配置
train_params = {'data': config_path,            # 指定之前创建的数据集配置文件路径'epochs': 1050,                 # 整个训练数据集将被遍历的次数'batch': 16,                    # 批大小（Batch Size），即每次梯度更新使用的图像数量'imgsz': 320 * 2,               # 输入图像的尺寸（宽度和高度，单位像素）。图像在训练前会被缩放到此尺寸'device': '0',                  # GPU设备ID，CPU使用'cpu'， 多卡可用 '0,1,2''workers': 8,                   # 用于数据加载的子进程数'optimizer': 'auto',            # 可选 SGD/Adam/AdamW等'project': './runs/train',      # 训练结果保存目录'name': 'football_detection',   # 本次实验的名称'exist_ok': True,               # 允许覆盖已有训练结果 (False 时，会自动创建新文件夹)'augment': True,                # 启用数据增强 (通过随机缩放、裁剪、翻转、调整色调等操作来增加数据的多样性，是防止过拟合、提升模型泛化能力的关键手段)'lr0': 0.005,                   # 初始学习率  'lrf': 0.2,                     # 最终学习率 = lr0 * lrf'save_period': 10,              # 每10个epoch保存一次模型'amp': True,                    # 启用自动混合精度训练 'patience': 100                 # 如果连续 100 个 epoch 验证指标（如 mAP）没有提升，则提前终止训练
} # 开始训练
results = model.train(**train_params)# 验证最佳模型
best_model = YOLO(os.path.join(train_params['project'], train_params['name'], 'weights/best.pt'))   # 加载训练过程中在验证集上表现最好的模型权重
metrics = best_model.val(data = config_path, # 使用同样的数据集配置batch = 32,         # 验证时的批大小。通常可以比训练时设得大一些，因为验证不需要计算梯度，显存占用更小conf = 0.001,       # 验证置信度阈值, 仅考虑置信度高于此值的预测框。（0.001阈值很低，目的是评估模型输出的所有可能的目标，避免因阈值过高而漏掉一些真实预测，从而得到最全面的评估结果）iou = 0.6,          # NMS(非极大值抑制) IoU阈值。 用于合并重叠的预测框IoU高于此值的两个框会被认为是预测的同一个物体，只保留置信度最高的那个。exist_ok = True     # 允许覆盖已有训练结果 (False 时，会自动创建新文件夹)
)

模型选择:

• yolo11n.pt: nano，轻量级模型，速度快；
• yolo11s.pt: small，小型模型，平衡速度和精度；
• yolo11m.pt: medium，中型模型；
• yolo11l.pt: large，大型模型；
• yolo11x.pt: extra-large，超大型模型，精度最高；

yolov11_goalpost_ball.yaml 内容主要包含：标签类型说明、数据基础路径和基于数据基础路径的训练集和测试集路径，具体内容如下：

names:0: Ball1: Post2: L3: T4: X
path: $(env HOME)/project/code_learn/train_yolo11/yolo_dataset
train: train/images
val: val/images

当前模型支持5个类别：

• 0: Ball (球)
• 1: Post (门柱)
• 2: L (L形标记)
• 3: T (T形标记)
• 4: X (X形标记)

在conda虚拟环境执行如下命令进行训练

python3 train.py 

下面是训练过程一次 epoch日志打印：

Epoch    GPU_mem   box_loss   cls_loss   dfl_loss  Instances       Size
1/1050      2.41G      2.286       4.72      1.464        159        640: 100%|██████████| 7/7 [00:01<00:00,  6.95it/s]Class     Images  Instances      Box(P          R      mAP50  mAP50-95): 100%|██████████| 1/1 [00:00<00:00, >11.41it/s]all         25        206   0.000256     0.0516   0.000253   6.73e-05

其具体参数分析下：

① 训练阶段参数：

• Epoch：当前训练轮次/总训练轮次
• GPU_mem：当前GPU内存使用量
• Loss Functions：损失函数，损失值越小越好，表示模型预测越准确；
  • box_loss：边界框回归损失，衡量预测框与真实框之间的差异
  • cls_loss：分类损失，衡量类别预测的准确性
  • dfl_loss：分布焦点损失，YOLOv8中引入的改进损失函数
• Instances：实例数量，当前批次中检测到的目标实例数量，159 表示当前批次中有159个目标对象；
• Size：图像尺寸，YOLO 通常使用固定的输入尺寸进行训练；
   
   

② 验证/评估阶段参数：

• Class：评估的类别，all 表示所有类别的综合评估结果；
• Images：用于评估的图像数量，25 表示使用了25张图像进行评估；
• Instances：实例数量，验证集中所有目标实例的总数，205 表示当前批次中有205个目标对象；
• 评估指标：
  • Box(P) ：精确率/准确率 (Precision)，预测为正样本中真正为正样本的比例( TP / (TP + FP) )，值越高表示误检越少，理想值为1；
  • Box(R)：召回率 (Recall)，真正为正样本中被正确预测为正样本的比例( TP / (TP + FN) )，值越高表示漏检越少，理想值为1；
  • mAP50 ：平均精度 (mean Average Precision at IoU=0.5)，在IoU阈值为0.5时的平均精度，综合衡量检测性能的主要指标，值越高越好；
  • mAP50-95 ： 平均精度 (mean Average Precision at IoU=0.5:0.95)，在IoU阈值从0.5到0.95（步长0.05）的平均精度，更严格的评估指标，考虑了不同IoU阈值下的性能；

3. 训练效果

训练结束后会在runs 路径下生成两个文件夹 train 和 detect ，下面是 train 文件夹下的训练过程数据：

F1_curve	P_curve
R_curve	PR_curve

三、测试训练模型

测试脚本 predict.py 内容如下：注意修改 source='./predict_data/vedio 为自己的视频路径。

from ultralytics import YOLOmodel = YOLO("./runs/train/football_detection/weights/best.pt")# 预测并保存结果
results = model.predict(source='./predict_data/vedio',conf=0.25,           # 设置对象置信度阈值，高于阈值的检测框才会保留和输出，默认通常为0.25:cite[2]:cite[4]save=True,           # 保存带标注的预测图片save_txt=True,       # 保存检测框的 txt 文件（YOLO 格式）save_conf=True,      # 保存置信度到 txt 文件save_crop=False,     # 保存裁剪的检测目标（可选）project='./results', # 指定保存路径（默认是 runs/detect）exist_ok=True        # 允许覆盖已有结果
)

上述脚本执行完后，会在 results 路径下生成 predict 保存测试视频的识别结果，效果如下所示：

四、模型转换及查看

1. 模型转换

现在有了 训练自己数据集的模型 best.pt 格式模型，接下来将其进行转换成其他格式模型文件，如 onnx

# 默认安装依赖库
pip install onnx onnxslim onnxruntime# 要求版本安装依赖库
pip install "onnx>=1.12.0,<1.18.0" onnxslim onnxruntime

参考 官网 创建 pt2onnx.py 文件，内容如下：

from ultralytics import YOLO# Load a model
model = YOLO("yolo11n.pt")  # load an official model
# model = YOLO("path/to/best.pt")  # load a custom trained model
model = YOLO("./runs/train/football_detection/weights/best.pt")  # load a custom trained model# Export the model
model.export(format="onnx")

导出参数说明：
 

参数	类型	默认值	描述
format	str	'torchscript'	导出模型的目标格式，例如 'onnx', 'torchscript', 'engine' （TensorRT）等。每种格式都支持与不同的 部署环境.
imgsz	int 或 tuple	640	模型输入所需的图像大小。可以是正方形图像的整数（例如， 640 对于 640x640）或元组 (height, width) 用于指定特定维度。
keras	bool	False	启用导出为 Keras 格式，用于 TensorFlow SavedModel，提供与 TensorFlow serving 和 API 的兼容性。
optimize	bool	False	导出到 TorchScript 时，应用针对移动设备的优化，可能会减小模型大小并提高 推理 性能。与 NCNN 格式或 CUDA 设备不兼容。
half	bool	False	启用 FP16（半精度）量化，从而减小模型大小并可能加快受支持硬件上的推理速度。与 ONNX 的 INT8 量化或仅 CPU 导出不兼容。 int8 bool False 激活 INT8 量化，进一步压缩模型并加速推理，同时最大限度地减少精度损失，主要用于边缘设备。与 TensorRT 结合使用时，执行训练后量化 (PTQ)。
dynamic	bool	False	允许 ONNX、TensorRT 和 OpenVINO 导出使用动态输入大小，从而提高处理不同图像尺寸的灵活性。自动设置为 True 当将TensorRT与INT8一起使用时。
simplify	bool	True	使用以下方式简化 ONNX 导出的模型图 onnxslim，从而可能提高性能以及与推理引擎的兼容性。
opset	int	None	指定 ONNX opset 版本，以与不同的 ONNX 解析器和运行时兼容。如果未设置，则使用最新支持的版本。
workspace	float 或 None	None	设置最大工作区大小，单位为GiB，用于 TensorRT 优化，平衡内存使用和性能。使用 None 用于 TensorRT 自动分配，最高可达设备最大值。
nms	bool	False	如果支持，则将非极大值抑制 (NMS) 添加到导出的模型（请参阅导出格式），从而提高检测后处理效率。不适用于 end2end 模型。
batch	int	1	指定导出模型批量推理大小，或导出模型将并发处理的最大图像数量，单位为。 predict 模式。对于 Edge TPU 导出，此项会自动设置为 1。
device	str	None	指定导出设备：GPU (device=0），CPU（device=cpu），适用于 Apple 芯片的 MPS（device=mps）或适用于 NVIDIA Jetson 的 DLA（device=dla:0 或 device=dla:1)。TensorRT 导出自动使用 GPU。
data	str	'coco8.yaml'	路径指向 数据集 配置文件（默认： coco8.yaml)，这对于 INT8 量化校准至关重要。如果启用 INT8 但未指定，则将分配默认数据集。
fraction	float	1.0	指定用于 INT8 量化校准的数据集比例。允许在完整数据集的子集上进行校准，这对于实验或资源有限时非常有用。如果未在使用 INT8 启用时指定，则将使用完整数据集。

 
 
 
导出格式:
 

格式	format 参数	模型	元数据	参数
PyTorch	-	yolo11n.pt	✅	-
TorchScript	torchscript	yolo11n.torchscript	✅	imgsz, half, dynamic, optimize, nms, batch, device
ONNX	onnx	yolo11n.onnx	✅	imgsz, half, dynamic, simplify, opset, nms, batch, device
OpenVINO	openvino	yolo11n_openvino_model/	✅	imgsz, half, dynamic, int8, nms, batch, data, fraction, device
TensorRT	engine	yolo11n.engine	✅	imgsz, half, dynamic, simplify, workspace, int8, nms, batch, data, fraction, device
CoreML	coreml	yolo11n.mlpackage	✅	imgsz, half, int8, nms, batch, device
TF SavedModel	saved_model	yolo11n_saved_model/	✅	imgsz, keras, int8, nms, batch, device
TF GraphDef	pb	yolo11n.pb	❌	imgsz, batch, device
TF Lite	tflite	yolo11n.tflite	✅	imgsz, half, int8, nms, batch, data, fraction, device
TF Edge TPU	edgetpu	yolo11n_edgetpu.tflite	✅	imgsz, device
TF.js	tfjs	yolo11n_web_model/	✅	imgsz, half, int8, nms, batch, device
PaddlePaddle	paddle	yolo11n_paddle_model/	✅	imgsz, batch, device
MNN	mnn	yolo11n.mnn	✅	imgsz, batch, int8, half, device
NCNN	ncnn	yolo11n_ncnn_model/	✅	imgsz, half, batch, device
IMX500	imx	yolo11n_imx_model/	✅	imgsz, int8, data, fraction, device
RKNN	rknn	yolo11n_rknn_model/	✅	imgsz, batch, name, device

执行如下命令即可在待转换的 .pt模型文件同目录下获得 .onnx 文件。

python3 pt2onnx.py

2. 模型查看

安装 Netron 软件：Netron-xx.AppImage 、icon ，具体图形化打开软件配置方法 参看 LabelImg 的安装及配置 。

sudo mkdir /usr/local/Netron
sudo cp netron_icon.png /usr/share/icons/
sudo cp Netron-8.5.2.AppImage /usr/local/Netron/sudo chmod +x /usr/local/Netron/Netron-8.5.2.AppImage
sudo vim /usr/share/applications/Netron.desktop

Netron.desktop 文件 内容如下：

[Desktop Entry]
Name=Netron
Comment=Netron
Exec=/usr/local/Netron/Netron-8.5.2.AppImage
Icon=/usr/share/icons/netron_icon.png
Type=Application
StartupNotify=true
Categories=Application;

然后就可以打开 Netron软件，愉快的玩耍了~