高级自定义#
Ultralytics YOLO 命令行界面和 Python 界面都是构建在基础引擎执行器之上的高级抽象。本指南重点介绍 Trainer 引擎,并解释如何针对你的具体需求对其进行自定义。
Watch: Mastering Ultralytics YOLO: Advanced Customization
有关常见训练器自定义(自定义指标、类别加权损失、模型保存、主干网络冻结和分层学习率)的实际示例,请参阅自定义训练器指南。
BaseTrainer#
BaseTrainer 类提供了一个通用的训练例程,适用于各种任务。你可以通过重写特定函数或操作来进行自定义,同时遵循所需的格式。例如,通过重写以下函数集成你自己的自定义模型和数据加载器:
get_model(cfg, weights):构建要训练的模型。get_dataloader():构建数据加载器。
有关更多详细信息和源代码,请参阅 BaseTrainer 参考。
DetectionTrainer#
以下是如何使用和自定义 Ultralytics YOLO DetectionTrainer 的方法:
from ultralytics.models.yolo.detect import DetectionTrainer
trainer = DetectionTrainer(overrides={...})
trainer.train()
trained_model = trainer.best # Get the best model自定义 DetectionTrainer#
若要训练不支持的自定义检测模型,请重载现有的 get_model 功能:
from ultralytics.models.yolo.detect import DetectionTrainer
class CustomTrainer(DetectionTrainer):
def get_model(self, cfg=None, weights=None, verbose=True):
"""Loads a custom detection model given configuration and weight files."""
trainer = CustomTrainer(overrides={...})
trainer.train()通过修改损失函数或添加回调来每 10 个轮次将模型上传到 Google Drive,从而进一步自定义训练器。以下是一个示例:
from ultralytics.models.yolo.detect import DetectionTrainer
from ultralytics.nn.tasks import DetectionModel
class MyCustomModel(DetectionModel):
def init_criterion(self):
"""Initializes the loss function and adds a callback for uploading the model to Google Drive every 10 epochs."""
class CustomTrainer(DetectionTrainer):
def get_model(self, cfg=None, weights=None, verbose=True):
"""Returns a customized detection model instance configured with specified config and weights."""
return MyCustomModel(...)
# Callback to upload model weights
def log_model(trainer):
"""Logs the path of the last model weight used by the trainer."""
last_weight_path = trainer.last
print(last_weight_path)
trainer = CustomTrainer(overrides={...})
trainer.add_callback("on_train_epoch_end", log_model) # Adds to existing callbacks
trainer.train()有关回调触发事件和入口点的更多信息,请参阅回调指南。
其他引擎组件#
类似地自定义其他组件,例如 Validators 和 Predictors。有关更多信息,请参考验证器和预测器的文档。
在自定义训练器中使用 YOLO#
YOLO 模型类为训练器类提供了一个高级包装器。你可以利用此架构在机器学习工作流中获得更大的灵活性:
from ultralytics import YOLO
from ultralytics.models.yolo.detect import DetectionTrainer
# Create a custom trainer
class MyCustomTrainer(DetectionTrainer):
def get_model(self, cfg=None, weights=None, verbose=True):
"""Custom code implementation."""
# Initialize YOLO model
model = YOLO("yolo26n.pt")
# Train with custom trainer
results = model.train(trainer=MyCustomTrainer, data="coco8.yaml", epochs=3)这种方法允许你在自定义底层训练过程以适应特定需求的同时,保持 YOLO 接口的简洁性。
常见问题#
通过重写其方法以适应你的自定义模型和数据加载器,针对特定任务自定义
DetectionTrainer。首先继承DetectionTrainer并重新定义诸如get_model之类的方法来实现自定义功能。以下是一个示例:from ultralytics.models.yolo.detect import DetectionTrainer class CustomTrainer(DetectionTrainer): def get_model(self, cfg=None, weights=None, verbose=True): """Loads a custom detection model given configuration and weight files.""" trainer = CustomTrainer(overrides={...}) trainer.train() trained_model = trainer.best # Get the best modelBaseTrainer是训练例程的基础,可以通过重写其通用方法来针对各种任务进行自定义。关键组件包括:get_model(cfg, weights):构建要训练的模型。get_dataloader():构建数据加载器。preprocess_batch():在模型前向传播之前处理批次预处理。set_model_attributes():根据数据集信息设置模型属性。get_validator():返回用于模型评估的验证器。
有关自定义和源代码的更多详细信息,请参阅
BaseTrainer参考。在
DetectionTrainer中添加回调以监视和修改训练过程。以下是如何添加回调以便在每个训练轮次后记录模型权重的示例:from ultralytics.models.yolo.detect import DetectionTrainer # Callback to upload model weights def log_model(trainer): """Logs the path of the last model weight used by the trainer.""" last_weight_path = trainer.last print(last_weight_path) trainer = DetectionTrainer(overrides={...}) trainer.add_callback("on_train_epoch_end", log_model) # Adds to existing callbacks trainer.train()有关回调事件和入口点的更多详细信息,请参考回调指南。
Ultralytics YOLO 在强大的引擎执行器之上提供了高级抽象,使其成为快速开发和自定义的理想选择。主要优势包括:
- 易用性:命令行界面和 Python 界面都简化了复杂的任务。
- 性能:针对实时目标检测和各种视觉 AI 应用进行了优化。
- 自定义:易于为自定义模型、损失函数和数据加载器进行扩展。
- 模块化:可以独立修改组件,而不会影响整个流水线。
- 集成:与 ML 生态系统中的流行框架和工具无缝协作。
通过浏览主要的Ultralytics YOLO页面,了解有关 YOLO 功能的更多信息。
可以,
DetectionTrainer高度灵活,可针对非标准模型进行自定义。继承DetectionTrainer并重载方法以支持你的特定模型需求。以下是一个简单的示例:from ultralytics.models.yolo.detect import DetectionTrainer class CustomDetectionTrainer(DetectionTrainer): def get_model(self, cfg=None, weights=None, verbose=True): """Loads a custom detection model.""" trainer = CustomDetectionTrainer(overrides={...}) trainer.train()有关全面的说明和示例,请查看
DetectionTrainer参考。