关于YOLO训练时的回调函数
编程侧
关于YOLO训练时的回调函数
编程侧

关于YOLO训练时的回调函数

多少事，从来急；天地转，光阴迫。一万年太久，只争朝夕。

进行单个YOLO训练时，或许你不需要特别在意YOLO的训练进度，换句话说，你只要瞄一眼现在训练的轮数，和每一轮的时间，就能估计一个大概的时间。但是我们如果想通过代码来获取训练的进度，就无法使用瞄一眼的方法。如果你需要获取YOLO的训练进度，可以参考YOLO内置的回调函数来获取当期的进度，这些回调函数在训练过程的关键阶段作为入口点——这是他们的原话。

在官网中，他提供了所以的回调函数，你可以根据自己的需求，在每个你想要介入的位置，调用对应的回调函数来实现你想要的功能。

训练器回调

回调函数描述
on_pretrain_routine_start在预训练程序开始时触发。
on_pretrain_routine_end在预训练程序结束时触发。
on_train_start在训练开始时触发。
on_train_epoch_start在每个训练 epoch 开始时触发。
on_train_batch_start在每个训练批次开始时触发。
optimizer_step在优化器步骤中触发。
on_before_zero_grad在梯度归零之前触发。
on_train_batch_end在每个训练批次结束时触发。
on_train_epoch_end在每个训练 epoch 结束时触发。
on_fit_epoch_end在每个拟合 epoch 结束时触发。
on_model_save在模型保存时触发。
on_train_end在训练过程结束时触发。
on_params_update在模型参数更新时触发。
teardown在清理训练过程时触发。

验证器回调

回调函数描述
on_val_start在验证开始时触发。
on_val_batch_start在每个验证批次开始时触发。
on_val_batch_end在每个验证批次结束时触发。
on_val_end在验证结束时触发。

预测器回调

回调函数描述
on_predict_start在预测过程开始时触发。
on_predict_batch_start在每个预测批次开始时触发。
on_predict_postprocess_end在预测后处理结束时触发。
on_predict_batch_end在每个预测批次结束时触发。
on_predict_end在预测过程结束时触发。

导出器回调

回调函数描述
on_export_start在导出过程开始时触发。
on_export_end在导出过程结束时触发。

举一个简单的例子，如果我想得知当前的训练进度，那么我可以调用on_train_epoch_end，在每一轮epoch训练结束时，记录一次，那么只要知道训练的轮数，便可得知进度。若你需要更加细腻的进度，你完全可以调用on_train_batch_end来实现，只需要通过简单的计算得知详细的训练进度。

使用epochs计算进度

from ultralytics import YOLO

total_epochs = 100  # 你训练时设置的 epochs


def on_train_epoch_end(trainer):
    """在每个 epoch 结束时计算进度"""

    current_epoch = trainer.epoch + 1  # 从0开始
    progress = current_epoch / total_epochs

    print(f"[Epoch] {current_epoch}/{total_epochs} -> {progress:.2%}")


model = YOLO("yolo11n.pt")

model.add_callback("on_train_epoch_end", on_train_epoch_end)

model.train(
    data="coco8.yaml",
    epochs=total_epochs
)

使用batch计算进度

from ultralytics import YOLO

total_epochs = 100


def on_train_start(trainer):
    """训练开始时初始化总步数"""
    trainer.total_steps = total_epochs * len(trainer.train_loader)
    trainer.current_step = 0


def on_train_batch_end(trainer):
    """每个 batch 结束时更新进度"""

    trainer.current_step += 1

    progress = trainer.current_step / trainer.total_steps

    print(
        f"[Step] {trainer.current_step}/{trainer.total_steps} "
        f"-> {progress:.2%}"
    )


model = YOLO("yolo11n.pt")

model.add_callback("on_train_start", on_train_start)
model.add_callback("on_train_batch_end", on_train_batch_end)

model.train(
    data="coco8.yaml",
    epochs=total_epochs
)

更高阶的回调自定义 Ultralytics 训练流程我还不会，等我再探。