跳转至内容

自定义训练器

Ultralytics 训练管道围绕 BaseTrainer 和任务专用训练器，例如 DetectionTrainer. 这些类开箱即用地处理训练循环、验证、检查点和日志记录。当您需要更多控制时——例如 track 自定义指标、调整损失权重或实现学习率调度——您可以子类化训练器并覆盖特定方法。

本指南将介绍五种常见的自定义方法：

  1. 在每个epoch结束时记录自定义指标（F1分数）
  2. 添加类别权重以处理类别不平衡
  3. 基于不同指标保存最佳模型
  4. 冻结骨干网络，在前 N 个 epoch 后再解冻
  5. 指定每层学习率

准备工作

在阅读本指南之前，请确保您熟悉以下基础知识： 训练 YOLO 模型 和 高级自定义 页面，其中涵盖了 BaseTrainer 架构。

自定义训练器的工作原理

字段 YOLO 模型类接受一个 trainer 参数在 train() 方法。这允许您传入自己的训练器类，以扩展默认行为：

from ultralytics import YOLO
from ultralytics.models.yolo.detect import DetectionTrainer


class CustomTrainer(DetectionTrainer):
    """A custom trainer that extends DetectionTrainer with additional functionality."""

    pass  # Add your customizations here


model = YOLO("yolo26n.pt")
model.train(data="coco8.yaml", epochs=10, trainer=CustomTrainer)

您的自定义训练器继承了所有功能自 DetectionTrainer，因此您只需覆盖您想要自定义的特定方法。

记录自定义指标

字段 验证 步骤计算 精度, 召回率和 mAP. 如果您需要额外的指标，例如按类别划分的 F1 分数，覆盖 validate():

import numpy as np

from ultralytics import YOLO
from ultralytics.models.yolo.detect import DetectionTrainer
from ultralytics.utils import LOGGER


class MetricsTrainer(DetectionTrainer):
    """Custom trainer that computes and logs F1 score at the end of each epoch."""

    def validate(self):
        """Run validation and compute per-class F1 scores."""
        metrics, fitness = super().validate()
        if metrics is None:
            return metrics, fitness

        if hasattr(self.validator, "metrics") and hasattr(self.validator.metrics, "box"):
            box = self.validator.metrics.box
            f1_per_class = box.f1
            class_indices = box.ap_class_index
            names = self.validator.names

            valid_f1 = f1_per_class[f1_per_class > 0]
            mean_f1 = np.mean(valid_f1) if len(valid_f1) > 0 else 0.0

            LOGGER.info(f"Mean F1 Score: {mean_f1:.4f}")
            per_class_str = [
                f"{names[i]}: {f1_per_class[j]:.3f}" for j, i in enumerate(class_indices) if f1_per_class[j] > 0
            ]
            LOGGER.info(f"Per-class F1: {per_class_str}")

        return metrics, fitness


model = YOLO("yolo26n.pt")
model.train(data="coco8.yaml", epochs=5, trainer=MetricsTrainer)

这会在每次验证运行后记录所有类别的平均F1分数和每个类别的详细分解。

可用指标

验证器通过以下方式提供对众多指标的访问 self.validator.metrics.box:

属性描述
f1每类 F1 分数
p每类精度
r各类别召回率
ap50每个类别的 IoU 0.5 AP
ap每个类别的 IoU 0.5:0.95 AP
mp, mr平均精确率和召回率
map50, map平均AP指标

添加类别权重

如果您的数据集存在不平衡类别（例如，制造检测中的罕见缺陷），您可以在损失函数中增加代表性不足类别的权重。这使得模型对稀有类别的错误分类施加更重的惩罚。

要自定义损失函数，请对损失函数类、模型和训练器进行子类化：

import torch
from torch import nn

from ultralytics import YOLO
from ultralytics.models.yolo.detect import DetectionTrainer
from ultralytics.nn.tasks import DetectionModel
from ultralytics.utils import RANK
from ultralytics.utils.loss import E2ELoss, v8DetectionLoss


class WeightedDetectionLoss(v8DetectionLoss):
    """Detection loss with class weights applied to BCE classification loss."""

    def __init__(self, model, class_weights=None, tal_topk=10, tal_topk2=None):
        """Initialize loss with optional per-class weights for BCE."""
        super().__init__(model, tal_topk=tal_topk, tal_topk2=tal_topk2)
        if class_weights is not None:
            self.bce = nn.BCEWithLogitsLoss(
                pos_weight=class_weights.to(self.device),
                reduction="none",
            )


class WeightedE2ELoss(E2ELoss):
    """E2E Loss with class weights for YOLO26."""

    def __init__(self, model, class_weights=None):
        """Initialize E2E loss with weighted detection loss."""

        def weighted_loss_fn(model, tal_topk=10, tal_topk2=None):
            return WeightedDetectionLoss(model, class_weights=class_weights, tal_topk=tal_topk, tal_topk2=tal_topk2)

        super().__init__(model, loss_fn=weighted_loss_fn)


class WeightedDetectionModel(DetectionModel):
    """Detection model that uses class-weighted loss."""

    def init_criterion(self):
        """Initialize weighted loss criterion with per-class weights."""
        class_weights = torch.ones(self.nc)
        class_weights[0] = 2.0  # upweight class 0
        class_weights[1] = 3.0  # upweight rare class 1
        return WeightedE2ELoss(self, class_weights=class_weights)


class WeightedTrainer(DetectionTrainer):
    """Trainer that returns a WeightedDetectionModel."""

    def get_model(self, cfg=None, weights=None, verbose=True):
        """Return a WeightedDetectionModel."""
        model = WeightedDetectionModel(cfg, nc=self.data["nc"], verbose=verbose and RANK == -1)
        if weights:
            model.load(weights)
        return model


model = YOLO("yolo26n.pt")
model.train(data="coco8.yaml", epochs=10, trainer=WeightedTrainer)

从数据集计算权重

您可以根据数据集的标签分布自动计算类别权重。一种常见的方法是逆频率加权：

import numpy as np

# class_counts: number of instances per class
class_counts = np.array([5000, 200, 3000])
# Inverse frequency: rarer classes get higher weight
class_weights = max(class_counts) / class_counts
# Result: [1.0, 25.0, 1.67]

通过自定义指标保存最佳模型

训练器保存 best.pt 基于适应度，默认为 0.9 × mAP@0.5:0.95 + 0.1 × mAP@0.5. 要使用不同的指标（例如 mAP@0.5 或召回），覆盖 validate() 并将您选择的指标作为适应度值返回。内置的 save_model() 将自动使用它：

from ultralytics import YOLO
from ultralytics.models.yolo.detect import DetectionTrainer


class CustomSaveTrainer(DetectionTrainer):
    """Trainer that saves the best model based on mAP@0.5 instead of default fitness."""

    def validate(self):
        """Override fitness to use mAP@0.5 for best model selection."""
        metrics, fitness = super().validate()
        if metrics:
            fitness = metrics.get("metrics/mAP50(B)", fitness)
            if self.best_fitness is None or fitness > self.best_fitness:
                self.best_fitness = fitness
        return metrics, fitness


model = YOLO("yolo26n.pt")
model.train(data="coco8.yaml", epochs=20, trainer=CustomSaveTrainer)

可用指标

可用的常见指标 self.metrics 验证后包括：

键描述
metrics/precision(B)精确度
metrics/recall(B)召回率
metrics/mAP50(B)IoU 0.5时的mAP
metrics/mAP50-95(B)IoU 0.5:0.95时的mAP

冻结与解冻主干网络

迁移学习 工作流程通常受益于在前N个epoch冻结预训练主干网络，从而让detect头在适应之前。 微调 整个网络。Ultralytics 提供一个 freeze 参数在训练开始时冻结层，您可以使用 回调函数 在N个epoch后解冻它们：

from ultralytics import YOLO
from ultralytics.models.yolo.detect import DetectionTrainer
from ultralytics.utils import LOGGER

FREEZE_EPOCHS = 5


def unfreeze_backbone(trainer):
    """Callback to unfreeze all layers after FREEZE_EPOCHS."""
    if trainer.epoch == FREEZE_EPOCHS:
        LOGGER.info(f"Epoch {trainer.epoch}: Unfreezing all layers for fine-tuning")
        for name, param in trainer.model.named_parameters():
            if not param.requires_grad:
                param.requires_grad = True
                LOGGER.info(f"  Unfroze: {name}")
        trainer.freeze_layer_names = [".dfl"]


class FreezingTrainer(DetectionTrainer):
    """Trainer with backbone freezing for first N epochs."""

    def __init__(self, *args, **kwargs):
        """Initialize and register the unfreeze callback."""
        super().__init__(*args, **kwargs)
        self.add_callback("on_train_epoch_start", unfreeze_backbone)


model = YOLO("yolo26n.pt")
model.train(data="coco8.yaml", epochs=20, freeze=10, trainer=FreezingTrainer)

字段 freeze=10 参数在训练开始时冻结前 10 层（骨干网络）。 on_train_epoch_start 回调函数在每个 epoch 开始时触发，并在冻结期完成后解冻所有参数。

选择要冻结的内容

  • freeze=10 冻结前10层（通常是YOLO架构中的主干网络）
  • freeze=[0, 1, 2, 3] 按索引冻结特定层
  • 更高 FREEZE_EPOCHS 这些值让头部在骨干网络改变之前有更多时间进行适应

逐层学习率

网络的不同部分可以从不同的学习率中受益。一种常见策略是为预训练的主干网络使用较低的学习率，以保留已学习的特征，同时允许检测头以更高的速率更快地适应：

import torch

from ultralytics import YOLO
from ultralytics.models.yolo.detect import DetectionTrainer
from ultralytics.utils import LOGGER
from ultralytics.utils.torch_utils import unwrap_model


class PerLayerLRTrainer(DetectionTrainer):
    """Trainer with different learning rates for backbone and head."""

    def build_optimizer(self, model, name="auto", lr=0.001, momentum=0.9, decay=1e-5, iterations=1e5):
        """Build optimizer with separate learning rates for backbone and head."""
        backbone_params = []
        head_params = []

        for k, v in unwrap_model(model).named_parameters():
            if not v.requires_grad:
                continue
            is_backbone = any(k.startswith(f"model.{i}.") for i in range(10))
            if is_backbone:
                backbone_params.append(v)
            else:
                head_params.append(v)

        backbone_lr = lr * 0.1

        optimizer = torch.optim.AdamW(
            [
                {"params": backbone_params, "lr": backbone_lr, "weight_decay": decay},
                {"params": head_params, "lr": lr, "weight_decay": decay},
            ],
        )

        LOGGER.info(
            f"PerLayerLR optimizer: backbone ({len(backbone_params)} params, lr={backbone_lr}) "
            f"| head ({len(head_params)} params, lr={lr})"
        )
        return optimizer


model = YOLO("yolo26n.pt")
model.train(data="coco8.yaml", epochs=20, trainer=PerLayerLRTrainer)

学习率调度器

内置学习率调度器（cosine 或 linear) 仍然适用于每个组的基础学习率之上。主干网络和头部学习率都将遵循相同的衰减计划，在整个训练过程中保持它们之间的比例。

结合技术

这些自定义可以通过重写多个方法并根据需要添加回调来组合成一个单一的训练器类。

常见问题

我如何将自定义训练器传递给YOLO？

将您的自定义训练器类（而非实例）传递给 trainer 参数中指定。 model.train():

from ultralytics import YOLO

model = YOLO("yolo26n.pt")
model.train(data="coco8.yaml", trainer=MyCustomTrainer)

字段 YOLO 类在内部处理训练器实例化。请参阅 高级自定义 页面，了解有关训练器架构的更多详细信息。

我可以覆盖哪些 BaseTrainer 方法？

可供定制的主要方法：

方法目的
validate()运行验证并返回指标
build_optimizer()构建优化器
save_model()保存训练检查点
get_model()返回模型实例
get_validator()返回验证器实例
get_dataloader()构建数据加载器
preprocess_batch()预处理输入批次
label_loss_items()格式化损失项用于日志记录

有关完整的API参考，请参阅 BaseTrainer 文档.

我能否使用回调函数而不是子类化训练器？

是的，对于更简单的自定义， 回调函数 通常足够。可用的回调事件包括 on_train_start, on_train_epoch_start, on_train_epoch_end, on_fit_epoch_end和 on_model_save. 这些允许您在不进行子类化的情况下介入训练循环。上面的骨干网络冻结示例演示了这种方法。

我如何在不子类化模型的情况下自定义损失函数？

如果您的更改更简单（例如调整损失增益），您可以直接修改超参数：

model.train(data="coco8.yaml", box=10.0, cls=1.5, dfl=2.0)

对于损失函数的结构性更改（例如添加类别权重），您需要像类别权重部分所示那样对损失函数和模型进行子类化。



📅 1 个月前创建 ✏️ 1 个月前更新
raimbekovmonuralpszr

评论