自定义 Trainer#
Ultralytics 训练流水线围绕 BaseTrainer 和诸如 DetectionTrainer 等特定任务的训练器构建。这些类开箱即用，可处理训练循环、验证、检查点保存和日志记录。当你需要更多控制时（例如跟踪自定义指标、调整损失权重或实现学习率调度），你可以继承训练器并重写特定方法。
本指南介绍了七种常见的自定义操作：
- 在每个 epoch 结束时记录自定义指标（F1 分数）
- 添加类权重以处理类别不平衡
- 根据不同的指标保存最佳模型
- 在前 N 个 epoch 中冻结骨干网络，然后解冻
- 指定每层的学习率
- 在多 GPU 训练中跨 GPU 同步 BatchNorm
- 配置梯度裁剪以进行稳定性调优
前提条件
在阅读本指南之前，请确保你熟悉训练 YOLO 模型的基础知识以及高级自定义页面，该页面涵盖了 BaseTrainer 架构。
自定义训练器的工作原理#
YOLO 模型类在 train() 方法中接受一个 trainer 参数。这允许你传入扩展默认行为的自定义 trainer 类：
from ultralytics import YOLO
from ultralytics.models.yolo.detect import DetectionTrainer
class CustomTrainer(DetectionTrainer):
"""A custom trainer that extends DetectionTrainer with additional functionality."""
# Add your customizations here
model = YOLO("yolo26n.pt")
model.train(data="coco8.yaml", epochs=10, trainer=CustomTrainer)你的自定义训练器继承了来自 DetectionTrainer 的所有功能，因此你只需重写想要自定义的特定方法。
记录自定义指标#
验证步骤会计算精度、召回率和 mAP。如果你需要其他指标（如每个类别的 F1 score），请重写 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 分数 |
image_metrics | 包含精确率、召回率、F1、TP、FP 和 FN 的单图像指标字典 |
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]加载带有自定义类的模型
诸如 WeightedDetectionModel 之类的自定义类通过引用存储在检查点中。当在训练脚本中定义时，它们属于 __main__ 模块，因此从不同脚本加载 best.pt 会引发 AttributeError: Can't get attribute 'WeightedDetectionModel' on <module '__main__'>。
在专用模块中定义自定义类以便它们保持可导入状态，并确保该模块在加载时位于你的 PYTHONPATH 上。
# weighted_model.py
from ultralytics.nn.tasks import DetectionModel
class WeightedDetectionModel(DetectionModel):
"""Detection model that uses class-weighted loss."""# inference script
from weighted_model import WeightedDetectionModel # noqa: F401 - must be importable at checkpoint load time
from ultralytics import YOLO
model = YOLO("runs/detect/train/weights/best.pt")
metrics = model.val()通过自定义指标保存最佳模型#
训练器根据适应度（fitness）保存 best.pt，对于检测任务，适应度默认为 mAP@0.5:0.95（对 [P, R, mAP@0.5, mAP@0.5:0.95] 赋予权重 [0.0, 0.0, 0.0, 1.0]）。要使用不同的指标（如 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 中冻结预训练的骨干网络，从而允许检测头在微调整个网络之前进行适应。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 层（索引 0-9），这涵盖了大部分 YOLO26 骨干网络。骨干网络跨越第 0-10 层，因此 freeze=10 使得最终的 C2PSA 块（第 10 层）保持可训练状态；使用 freeze=11 可以冻结整个骨干网络。on_train_epoch_start 回调在每个 epoch 开始时触发，并在冻结期结束后解冻所有参数。
选择要冻结的内容
freeze=10冻结前 10 层，即索引 0-9（大部分 YOLO26 骨干网络；使用freeze=11以包含第 10 层的最终 C2PSA 块）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 = []
unwrapped = unwrap_model(model)
backbone_len = len(unwrapped.yaml["backbone"]) # YOLO26 backbone spans layers 0-10 (C2PSA at layer 10)
for k, v in unwrapped.named_parameters():
if not v.requires_grad:
continue
is_backbone = any(k.startswith(f"model.{i}.") for i in range(backbone_len))
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)RT-DETR 变体#
对于 RT-DETR，模式相同但有两个改进。骨干网络长度从 model.yaml["backbone"] 读取，因此同一个训练器可以跨 RT-DETR 变体（RT-DETR-L、RT-DETR-X、ResNet-50/101 骨干网络）工作，而无需硬编码层数。参数还在每个部分内被拆分为权重、BatchNorm 和偏置组，以便从 BatchNorm 参数和偏置中排除权重衰减，这与默认训练器的策略相匹配。这对于 RT-DETR 微调特别有用，其中解码器头通常是随机初始化的，而骨干网络携带受益于较低学习率的预训练特征：
import torch
from torch import nn
from ultralytics import RTDETR
from ultralytics.models.rtdetr.train import RTDETRTrainer
from ultralytics.utils import LOGGER, colorstr
from ultralytics.utils.torch_utils import unwrap_model
class RTDETRBackboneLRTrainer(RTDETRTrainer):
"""RT-DETR trainer with a lower learning rate for backbone parameters."""
backbone_lr_ratio = 0.1 # backbone learning rate as a fraction of head learning rate
def build_optimizer(self, model, name="auto", lr=0.001, momentum=0.9, decay=1e-5, iterations=1e5):
"""Build an AdamW optimizer with six param groups: head and backbone x {weight, bn, bias}."""
# Resolve optimizer name; "auto" maps to AdamW with RT-DETR-style defaults
canonical = {"Adam", "Adamax", "AdamW", "NAdam", "RAdam", "auto"}
name = {x.lower(): x for x in canonical}.get(name.lower(), name)
if name == "auto":
name, lr, momentum = "AdamW", 1e-4, 0.9
self.args.warmup_bias_lr = 0.0 # RT-DETR warms biases from 0, unlike YOLO's 0.1
if name not in {"Adam", "Adamax", "AdamW", "NAdam", "RAdam"}:
raise NotImplementedError(f"This trainer only supports AdamW-family optimizers; got {name}")
# Identify backbone parameters from model.yaml and route each param into a (section, kind) group
unwrapped = unwrap_model(model)
backbone_len = len(unwrapped.yaml["backbone"])
norm_types = tuple(v for k, v in nn.__dict__.items() if "Norm" in k)
groups = {f"{s}_{k}": [] for s in ("head", "backbone") for k in ("weight", "bn", "bias")}
for module_name, module in unwrapped.named_modules():
for param_name, param in module.named_parameters(recurse=False):
if not param.requires_grad:
continue
fullname = f"{module_name}.{param_name}" if module_name else param_name
parts = fullname.split(".")
section = (
"backbone"
if len(parts) > 1 and parts[0] == "model" and parts[1].isdigit() and int(parts[1]) < backbone_len
else "head"
)
if "bias" in param_name:
kind = "bias"
elif isinstance(module, norm_types) or "logit_scale" in fullname:
kind = "bn"
else:
kind = "weight"
groups[f"{section}_{kind}"].append(param)
# Build the optimizer with per-group lr and weight decay; backbone groups use lr * backbone_lr_ratio
backbone_lr = lr * self.backbone_lr_ratio
param_groups = [
{"params": groups["head_weight"], "lr": lr, "weight_decay": decay, "param_group": "weight"},
{"params": groups["head_bn"], "lr": lr, "weight_decay": 0.0, "param_group": "bn"},
{"params": groups["head_bias"], "lr": lr, "weight_decay": 0.0, "param_group": "bias"},
{"params": groups["backbone_weight"], "lr": backbone_lr, "weight_decay": decay, "param_group": "weight"},
{"params": groups["backbone_bn"], "lr": backbone_lr, "weight_decay": 0.0, "param_group": "bn"},
{"params": groups["backbone_bias"], "lr": backbone_lr, "weight_decay": 0.0, "param_group": "bias"},
]
param_groups = [pg for pg in param_groups if pg["params"]] # drop empty groups
optimizer = getattr(torch.optim, name)(param_groups, betas=(momentum, 0.999))
LOGGER.info(
f"{colorstr('optimizer:')} {name}(lr={lr}, backbone_lr={backbone_lr}) with parameter groups\n"
f" Head: {len(groups['head_bn'])} bn, {len(groups['head_weight'])} weight(decay={decay}), "
f"{len(groups['head_bias'])} bias (lr={lr})\n"
f" Backbone: {len(groups['backbone_bn'])} bn, {len(groups['backbone_weight'])} weight(decay={decay}), "
f"{len(groups['backbone_bias'])} bias (lr={backbone_lr})"
)
return optimizer
model = RTDETR("rtdetr-l.pt")
model.train(data="coco8.yaml", epochs=20, trainer=RTDETRBackboneLRTrainer)选择 `backbone_lr_ratio`
一个常见的起点是 backbone_lr_ratio = 0.1，它与带有 HGNetV2 骨干网络的原始 RT-DETR 设置相匹配。文献建议将该比率与骨干网络大小和预训练数据规模成反比缩放：在数亿张图像上预训练的大型骨干网络（例如使用 DINO、CLIP 或 MAE 训练的 ViT-L/H）通常使用诸如 0.01 或更小的较小比率来保留学得良好的特征，而具有较轻预训练的较小骨干网络则能容忍诸如 0.5 或更高的较大比率。
学习率调度器
内置的学习率调度器（cosine 或 linear）仍然适用于每组基础学习率之上。骨干网络和检测头的学习率都将遵循相同的衰减调度，在整个训练过程中保持它们之间的比率。
组合技术
可以通过覆盖多个方法并根据需要添加回调，将这些自定义组合到单个训练器类中。
多 GPU 训练的同步 BatchNorm#
在使用 DistributedDataParallel 在多个 GPU 上进行训练时，默认的 BatchNorm2d 层会在每个 GPU 上独立计算统计信息。对于 RT-DETR 微调和其他使用较小每 GPU 批次大小的配方来说，每 GPU 的批次统计信息可能会有噪声。PyTorch 的 SyncBatchNorm 跨所有进程同步均值和方差以获得单个全局批次统计信息，这通常以少量的进程间 GPU 通信开销为代价来改善收敛性。
转换必须在模型位于 GPU 上之后但在 DDP 包装它之前进行。最干净的钩子是 set_model_attributes()，BaseTrainer 正是在这个窗口中调用它：
from torch import nn
from ultralytics import RTDETR
from ultralytics.models.rtdetr.train import RTDETRTrainer
class SyncBNTrainer(RTDETRTrainer):
"""RT-DETR trainer that converts BatchNorm to SyncBatchNorm for multi-GPU training."""
def set_model_attributes(self):
"""Run the parent setup, then convert BN to SyncBatchNorm when training on multiple GPUs."""
super().set_model_attributes()
if self.world_size > 1:
self.model = nn.SyncBatchNorm.convert_sync_batchnorm(self.model)
model = RTDETR("rtdetr-l.pt")
model.train(data="coco8.yaml", epochs=20, device=[0, 1], trainer=SyncBNTrainer)world_size > 1 卫语句确保训练器在单 GPU 运行中也是安全的；在单个 GPU 上，转换会被跳过，并且训练以常规的 BatchNorm2d 继续进行。通过将父类切换为 DetectionTrainer，同样的模式也适用于 YOLO。
何时使用 SyncBatchNorm
| 场景 | 建议 |
|---|---|
| 多 GPU 训练，小单 GPU 批次 (≤ 16) | 启用 |
| 多 GPU 训练，大单 GPU 批次 (≥ 32) | 可选；收益较小 |
| 单 GPU 训练 | 不适用（已跳过） |
可配置的梯度裁剪#
默认的 trainer 在 optimizer_step() 中将梯度裁剪到 max_norm=10.0，这是一个为 YOLO 模型调优的宽松值，因为这些模型的梯度很少超过该值。DETR 系列检测器（RT-DETR、DEIM、DINO）通常使用更严格的值，例如 0.1，以稳定解码器的交叉注意力层，这里的梯度幅度可能会激增。要覆盖裁剪值，请继承 trainer 并覆盖 optimizer_step()：
import torch
from ultralytics import RTDETR
from ultralytics.models.rtdetr.train import RTDETRTrainer
class CustomClipTrainer(RTDETRTrainer):
"""RT-DETR trainer with configurable gradient clipping."""
clip_grad_norm = 0.1 # max gradient norm; set to 0 to disable clipping
def optimizer_step(self):
"""Run an optimizer step with a configurable gradient-norm clip."""
self.scaler.unscale_(self.optimizer)
if self.clip_grad_norm > 0:
torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=self.clip_grad_norm)
self.scaler.step(self.optimizer)
self.scaler.update()
self.optimizer.zero_grad()
if self.ema:
self.ema.update(self.model)
model = RTDETR("rtdetr-l.pt")
model.train(data="coco8.yaml", epochs=20, trainer=CustomClipTrainer)通过将父类切换为 DetectionTrainer（from ultralytics.models.yolo.detect import DetectionTrainer）并使用 YOLO("yolo26n.pt") 加载 YOLO 检查点，同样的训练器也适用于 YOLO。optimizer_step 主体保持不变。
典型的 `clip_grad_norm` 值
| 架构系列 | 典型的 max_norm |
|---|---|
| RT-DETR / DEIM / DETR 系列 | 0.1 |
| YOLO (Ultralytics 默认) | 10.0 |
| 禁用裁剪 | 0 |
常见问题解答#
将你的自定义 trainer 类（而不是实例）传递给
model.train()中的trainer参数：from ultralytics import YOLO model = YOLO("yolo26n.pt") model.train(data="coco8.yaml", trainer=MyCustomTrainer)YOLO类在内部处理训练器实例化。有关训练器架构的更多详细信息，请参阅高级自定义页面。可用于自定义的关键方法：
方法 用途 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。这些允许你在不进行子类化的情况下挂钩到训练循环中。上面的骨干网络冻结示例演示了这种方法。