回调#
Ultralytics 框架支持回调,它们可作为 train、val、export 和 predict 模式关键阶段的入口。每个回调都会根据操作类型接收 Trainer、Validator 或 Predictor 对象。这些对象的所有属性都在 BaseTrainer、BaseValidator 和 BasePredictor 参考页面中详细说明。
观看: 如何使用 Ultralytics 回调 | 预测、训练、验证和导出回调 | Ultralytics YOLO26 🚀
示例#
在预测中返回附加信息#
在此示例中,我们演示如何在每个结果对象中同时返回原始帧:
from ultralytics import YOLO
def on_predict_batch_end(predictor):
"""Combine prediction results with corresponding frames."""
_, image, _ = predictor.batch
# 确保 image 是一个列表
image = image if isinstance(image, list) else [image]
# 将预测结果与对应帧组合
predictor.results = zip(predictor.results, image)
# 创建 YOLO 模型实例
model = YOLO("yolo26n.pt")
# 向模型添加自定义回调
model.add_callback("on_predict_batch_end", on_predict_batch_end)
# 遍历结果和帧
for result, frame in model.predict(): # 或 model.track()
pass训练前注册 PyTorch 前向钩子#
model.train() 根据模型配置构建训练模型,并将已加载的权重迁移到该模型中。在调用此方法之前直接在 model.model 上注册的 PyTorch 钩子属于运行时状态,因此重建后的训练模型不会继承这些钩子。请在模型设置完成后,于 on_pretrain_routine_end 中注册训练钩子:
from ultralytics import YOLO
from ultralytics.utils.torch_utils import unwrap_model
def preprocess_input(module, inputs):
"""Center each input channel before the first model layer."""
images = inputs[0]
return (images - images.mean(dim=(-2, -1), keepdim=True),)
def register_forward_pre_hook(trainer):
"""Attach preprocessing to the active training model."""
train_model = unwrap_model(trainer.model)
train_model.model[0].register_forward_pre_hook(preprocess_input)
model = YOLO("yolo26n.pt")
model.add_callback("on_pretrain_routine_end", register_forward_pre_hook)
model.train(data="coco8.yaml", epochs=1)unwrap_model() 同时支持单设备训练和 DistributedDataParallel 训练。不要将本地定义的钩子附加到 trainer.ema.ema,因为训练检查点会序列化 EMA 模型,其他进程加载检查点时可能无法导入该回调。如果训练验证期间也必须运行相同的预处理,请将其实现为可导入的模型组件,而不是运行时钩子。NaN 恢复会根据训练模型重建 EMA,因此以这种方式注册的钩子会被复制到 EMA,并在恢复后保存到检查点;如果检查点必须能在其他进程中加载,请将钩子函数定义在可导入的模块中。
独立的 model.val() 会在每次调用时复制已加载的模型。预测会在首次调用时创建并缓存一个副本,因此请在首次调用 model.predict() 或 model.track() 之前,在 model.model 上注册钩子;之后添加的钩子不会传递到缓存的预测器。在新进程中加载检查点后,请重新注册运行时钩子。
使用 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)所有回调#
以下是所有受支持的回调。有关详情,请参阅回调源代码。
训练器回调#
| 回调 | 说明 |
|---|---|
on_pretrain_routine_start | 在预训练流程开始时触发,此时尚未加载数据和设置模型。 |
on_pretrain_routine_end | 在预训练流程结束时触发,此时数据加载和模型设置均已完成。 |
on_train_start | 训练开始时触发,在第一个轮次开始之前。 |
on_train_epoch_start | 每个训练轮次开始时触发,在批次迭代开始之前。 |
on_train_batch_start | 每个训练批次开始时触发,在前向传播之前。 |
optimizer_step | 优化器步骤期间触发。此回调预留给自定义集成使用;默认训练循环不会调用它。 |
on_before_zero_grad | 梯度清零之前触发。此回调预留给自定义集成使用;默认训练循环不会调用它。 |
on_train_batch_end | 每个训练批次结束时触发,在反向传播之后。由于梯度累积,优化器步骤可能会延后。 |
on_train_epoch_end | 每个训练轮次结束时触发,即所有批次处理完毕后、验证开始之前。此时验证指标和适应度可能尚不可用。 |
on_model_save | 模型检查点保存时触发,在验证之后。 |
on_fit_epoch_end | 每个完整训练轮次(训练 + 验证)结束时触发,即验证和任何检查点保存之后。验证指标此时可用;对于逐轮训练调用,适应度也可用。最终最佳模型评估期间也会调用此回调,但此时不会保存检查点,适应度也可能不存在。 |
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对象。有关这些对象属性的详情,请参阅BaseTrainer、BaseValidator和BasePredictor参考页面。要使用回调,请定义一个函数,并通过
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,让你能够添加自定义指标、处理逻辑或日志记录。下面介绍如何在每个训练轮次结束时记录每个优化器参数组的学习率:
from ultralytics import YOLO def log_learning_rates(trainer): """Print the current learning rates after each training epoch.""" print(f"Epoch {trainer.epoch + 1}: {trainer.lr}") model = YOLO("yolo26n.pt") model.add_callback("on_train_epoch_end", log_learning_rates) model.train(data="coco8.yaml", epochs=3)提示使用
freeze时,无需通过回调来冻结 BatchNorm 统计信息:训练器会在每个轮次开始时自动将冻结层中的 BatchNorm 层设为评估模式。有关如何有效使用训练回调的更多详情,请参阅训练指南。
在 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")有关如何在验证流程中集成回调的更多说明,请参阅验证指南。
要为 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如需更全面地了解用法,请参阅预测指南,其中包含详细说明和更多自定义选项。
Ultralytics YOLO 支持多种实用的回调实现,可增强和自定义训练、验证和预测等不同阶段。以下是一些实用示例:
示例:使用
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查看回调源代码,了解更多选项和示例。