回调函数#
Ultralytics 框架支持回调函数,它们用作 train、val、export 和 predict 模式中关键阶段的入口点。每个回调函数接受一个 Trainer、Validator 或 Predictor 对象,具体取决于操作类型。这些对象的所有属性都在文档的 Reference section 中有详细说明。
Watch: How to use Ultralytics Callbacks | Predict, Train, Validate and Export Callbacks | Ultralytics YOLO🚀
示例#
在预测时返回额外信息#
在此示例中,我们将演示如何返回原始帧以及每个结果对象:
from ultralytics import YOLO
def on_predict_batch_end(predictor):
"""Combine prediction results with corresponding frames."""
_, image, _, _ = predictor.batch
# Ensure that image is a list
image = image if isinstance(image, list) else [image]
# Combine the prediction results with the corresponding frames
predictor.results = zip(predictor.results, image)
# Create a YOLO model instance
model = YOLO("yolo26n.pt")
# Add the custom callback to the model
model.add_callback("on_predict_batch_end", on_predict_batch_end)
# Iterate through the results and frames
for result, frame in model.predict(): # or model.track()
pass使用 on_model_save 回调函数访问模型指标#
此示例展示了如何在保存检查点后使用 on_model_save 回调函数检索训练详情,例如 best_fitness 得分、total_loss 以及其他指标。
from ultralytics import YOLO
# Load a YOLO model
model = YOLO("yolo26n.pt")
def print_checkpoint_metrics(trainer):
"""Print trainer metrics and loss details after each checkpoint is saved."""
print(
f"Model details\n"
f"Best fitness: {trainer.best_fitness}, "
f"Loss names: {trainer.loss_names}, " # List of loss names
f"Metrics: {trainer.metrics}, "
f"Total loss: {trainer.tloss}" # Total loss value
)
if __name__ == "__main__":
# Add on_model_save callback.
model.add_callback("on_model_save", print_checkpoint_metrics)
# Run model training on custom dataset.
results = model.train(data="coco8.yaml", epochs=3)所有回调函数#
以下是所有支持的回调函数。更多详情,请参考回调函数 source code。
训练器回调#
| 回调函数 | 描述 |
|---|---|
on_pretrain_routine_start | 在预训练例程开始时触发,位于数据加载和模型设置之前。 |
on_pretrain_routine_end | 在预训练例程结束时触发,位于数据加载和模型设置完成之后。 |
on_train_start | 在训练开始时、第一个 epoch 开始之前触发。 |
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_model_save | 在模型检查点保存时触发,位于验证之后。 |
on_fit_epoch_end | 在每个拟合 epoch (训练 + 验证) 结束时触发,位于验证和任何检查点保存 之后。验证指标可用,且每轮训练调用的适应度也是可用的。此回调函数也会在最终最佳模型评估期间调用,此时不会保存检查点,且适应度可能不存在。 |
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 | 在导出过程结束时触发。 |
常见问题解答#
Ultralytics 回调函数是专用的入口点,会在模型操作的关键阶段(例如训练、验证、导出和预测)被触发。这些回调函数在过程中的特定点启用自定义功能,允许对工作流进行增强和修改。每个回调函数接受一个
Trainer、Validator或Predictor对象,具体取决于操作类型。有关这些对象的详细属性,请参考 Reference section。要使用回调函数,请定义一个函数并使用
model.add_callback()方法将其添加到模型中。以下是一个在预测期间返回附加信息的示例:from ultralytics import YOLO def on_predict_batch_end(predictor): """Handle prediction batch end by combining results with corresponding frames; modifies predictor results.""" _, image, _, _ = predictor.batch image = image if isinstance(image, list) else [image] predictor.results = zip(predictor.results, image) model = YOLO("yolo26n.pt") model.add_callback("on_predict_batch_end", on_predict_batch_end) for result, frame in model.predict(): pass通过在训练过程的特定阶段注入逻辑来定制你的 Ultralytics 训练常规。Ultralytics YOLO 提供了各种训练回调函数,例如
on_train_start、on_train_end和on_train_batch_end,允许你添加自定义指标、处理或日志记录。以下是在使用回调函数冻结层时冻结 BatchNorm 统计信息的方法:
from ultralytics import YOLO # Add a callback to put the frozen layers in eval mode to prevent BN values from changing def put_in_eval_mode(trainer): n_layers = trainer.args.freeze if not isinstance(n_layers, int): return for i, (name, module) in enumerate(trainer.model.named_modules()): if name.endswith("bn") and int(name.split(".")[1]) < n_layers: module.eval() module.track_running_stats = False model = YOLO("yolo26n.pt") model.add_callback("on_train_epoch_start", put_in_eval_mode) model.train(data="coco.yaml", epochs=10)有关有效使用训练回调函数的更多详情,请参见 Training Guide。
在 Ultralytics YOLO 的验证过程中使用回调函数可以通过启用自定义处理、日志记录或指标计算来增强模型评估。像
on_val_start、on_val_batch_end和on_val_end这样的回调函数提供了注入自定义逻辑的入口点,确保详细且全面的验证过程。例如,如果要绘制所有验证批次而不仅仅是前三个:
import inspect from ultralytics import YOLO def plot_samples(validator): frame = inspect.currentframe().f_back.f_back v = frame.f_locals validator.plot_val_samples(v["batch"], v["batch_i"]) validator.plot_predictions(v["batch"], v["preds"], v["batch_i"]) model = YOLO("yolo26n.pt") model.add_callback("on_val_batch_end", plot_samples) model.val(data="coco.yaml")有关将回调函数并入验证过程的更多见解,请参见 Validation Guide。
要在 Ultralytics YOLO 中为预测模式附加自定义回调函数,请定义一个回调函数并将其注册到预测过程中。常见的预测回调函数包括
on_predict_start、on_predict_batch_end和on_predict_end。这些函数允许修改预测输出并集成额外功能,例如数据日志记录或结果转换。这是一个示例,其中自定义回调函数会根据是否存在特定类别的对象来保存预测结果:
from ultralytics import YOLO model = YOLO("yolo26n.pt") class_id = 2 def save_on_object(predictor): r = predictor.results[0] if class_id in r.boxes.cls: predictor.args.save = True else: predictor.args.save = False model.add_callback("on_predict_postprocess_end", save_on_object) results_stream = model("pedestrians.mp4", stream=True, save=True) for result in results_stream: pass有关更全面的用法,请参考 Prediction Guide,其中包含详细说明和附加自定义选项。
Ultralytics YOLO 支持各种实用的回调函数实现,以增强和自定义不同的阶段,如训练、验证和预测。一些实际示例包括:
- Logging Custom Metrics: 在不同阶段记录附加指标,例如在训练或验证 epochs 结束时。
- Data Augmentation: 在预测或训练批次期间实现自定义数据转换或增强。
- 中间结果:保存中间结果(例如预测或帧)以供进一步分析或可视化。
示例:在预测期间使用
on_predict_batch_end将帧与预测结果结合:from ultralytics import YOLO def on_predict_batch_end(predictor): """Combine prediction results with frames.""" _, image, _, _ = predictor.batch image = image if isinstance(image, list) else [image] predictor.results = zip(predictor.results, image) model = YOLO("yolo26n.pt") model.add_callback("on_predict_batch_end", on_predict_batch_end) for result, frame in model.predict(): pass探索 callback source code 以获取更多选项和示例。