YOLOv8模型训练方法及训练停止threshold设置问题咨询

1. 如何使用YOLOv8训练模型？ #

• 准备合规数据集：采用YOLO自定义格式，每个图片对应同名的txt标注文件（每行格式为类别id x_center y_center width height，坐标为归一化后的值），同时创建data.yaml配置文件，示例如下：

train: ./train/images  # 训练集图片路径
val: ./val/images      # 验证集图片路径
nc: 2                  # 数据集类别总数
names: ['cat', 'dog']  # 类别名称列表

• 安装依赖：通过pip安装ultralytics库即可

pip install ultralytics

• 启动训练：支持命令行或Python脚本两种方式
  命令行示例：

yolo task=detect mode=train model=yolov8n.pt data=data.yaml epochs=100 imgsz=640

Python脚本示例：

from ultralytics import YOLO

# 加载预训练的YOLOv8n模型
model = YOLO('yolov8n.pt')
# 执行训练，指定数据集配置、训练轮数、输入图片尺寸等参数
results = model.train(data='data.yaml', epochs=100, imgsz=640)

2. 设置验证误差阈值触发训练停止 #

YOLOv8官方确实没有提供直接控制该逻辑的参数，但可以通过自定义回调函数实现：

• 编写一个继承自BaseCallback的回调类，在每个epoch结束后检查验证误差，当低于设定阈值时终止训练，示例代码如下：

from ultralytics import YOLO
from ultralytics.engine.callbacks import BaseCallback

class StopOnValLoss(BaseCallback):
    def __init__(self, threshold):
        self.threshold = threshold  # 设定的验证损失阈值

    def on_epoch_end(self, trainer):
        # 获取当前epoch的验证集总损失
        current_val_loss = trainer.metrics['val/loss']
        if current_val_loss < self.threshold:
            print(f"验证损失 {current_val_loss:.4f} 低于阈值 {self.threshold}，停止训练")
            trainer.stop = True  # 触发训练停止

# 加载模型
model = YOLO('yolov8n.pt')
# 添加自定义回调，这里设置阈值为0.5
model.add_callback(StopOnValLoss(threshold=0.5))
# 启动训练
results = model.train(data='data.yaml', epochs=100, imgsz=640)

你也可以根据需求替换监控的指标，比如换成val/box_loss（边界框损失）、val/cls_loss（分类损失）等，具体可参考训练日志里输出的指标名称。

内容的提问来源于stack exchange，提问作者user24929150