Personnalisation du Trainer#
Le pipeline d'entraînement d'Ultralytics est articulé autour de BaseTrainer et de formateurs spécifiques aux tâches comme DetectionTrainer. Ces classes gèrent la boucle d'entraînement, la validation, la création de points de contrôle et la journalisation de manière autonome. Lorsque tu as besoin de plus de contrôle — suivre des métriques personnalisées, ajuster la pondération des pertes ou implémenter des planificateurs de taux d'apprentissage — tu peux sous-classer le formateur et surcharger des méthodes spécifiques.
Ce guide explore sept personnalisations courantes :
- Journalisation de métriques personnalisées (score F1) à la fin de chaque époque
- Ajout de poids de classes pour gérer le déséquilibre des classes
- Enregistrement du meilleur modèle basé sur une métrique différente
- Gel du tronc pour les N premières époques, puis dégel
- Spécification de taux d'apprentissage par couche
- Synchronisation de BatchNorm entre les GPU pour l'entraînement multi-GPU
- Configuration du découpage des gradients pour le réglage de la stabilité
Avant de lire ce guide, assure-toi de bien connaître les bases de l'entraînement de modèles YOLO et la page de Personnalisation Avancée, qui couvre l'architecture de BaseTrainer.
Comment fonctionnent les trainers personnalisés#
La classe de modèle YOLO accepte un paramètre trainer dans la méthode train(). Cela te permet de passer ta propre classe de formateur qui étend le comportement par défaut :
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)Ton formateur personnalisé hérite de toutes les fonctionnalités de DetectionTrainer, tu n'as donc besoin de surcharger que les méthodes spécifiques que tu souhaites personnaliser.
Journalisation de métriques personnalisées#
L'étape de validation calcule la précision, le rappel et le mAP. Si tu as besoin de métriques supplémentaires telles que le score F1 par classe, surcharge 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)Ceci enregistre le score F1 moyen sur toutes les classes et une répartition par classe après chaque exécution de validation.
Le validateur donne accès à de nombreuses métriques via self.validator.metrics.box :
| Attribut | Description |
|---|---|
f1 | Score F1 par classe |
image_metrics | Dictionnaire des métriques par image avec précision, rappel, F1, TP, FP et FN |
p | Précision par classe |
r | Rappel par classe |
ap50 | AP à IoU 0.5 par classe |
ap | AP à IoU 0.5:0.95 par classe |
mp, mr | Précision et rappel moyens |
map50, map | Métriques de mAP moyen |
Ajout de poids de classe#
Si ton jeu de données présente des classes déséquilibrées (par exemple, un défaut rare dans le contrôle de fabrication), tu peux surpondérer les classes sous-représentées dans la fonction de perte. Cela pousse le modèle à pénaliser plus sévèrement les erreurs de classification sur les classes rares.
Pour personnaliser la perte, crée une sous-classe des classes de perte, du modèle et du trainer :
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)Tu peux calculer automatiquement les poids des classes à partir de la distribution des labels de ton jeu de données. Une approche courante est la pondération par fréquence inverse :
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]Les classes personnalisées telles que WeightedDetectionModel sont stockées dans le point de contrôle par référence. Lorsqu'elles sont définies dans un script d'entraînement, elles appartiennent au module __main__, donc le chargement de best.pt à partir d'un script différent lève une erreur AttributeError: Can't get attribute 'WeightedDetectionModel' on <module '__main__'>.
Définis les classes personnalisées dans un module dédié afin qu'elles restent importables, et assure-toi que ce module se trouve sur ton PYTHONPATH au moment du chargement.
# 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()Sauvegarde du meilleur modèle par métrique personnalisée#
Le formateur enregistre best.pt en fonction de l'aptitude, qui par défaut pour la détection est mAP@0.5:0.95 (pondère [0.0, 0.0, 0.0, 1.0] pour [P, R, mAP@0.5, mAP@0.5:0.95]). Pour utiliser une métrique différente (comme mAP@0.5 ou le rappel), surcharge validate() et renvoie la métrique de ton choix comme valeur d'aptitude. Le save_model() intégré l'utilisera alors automatiquement :
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)Les métriques courantes disponibles dans self.metrics après la validation incluent :
| Clé | Description |
|---|---|
metrics/precision(B) | Précision |
metrics/recall(B) | Rappel |
metrics/mAP50(B) | mAP à IoU 0.5 |
metrics/mAP50-95(B) | mAP à IoU 0.5:0.95 |
Gel et dégel de la backbone#
Les flux de travail d'apprentissage par transfert bénéficient souvent du gel du tronc pré-entraîné pour les N premières époques, permettant à la tête de détection de s'adapter avant le réglage fin de l'ensemble du réseau. Ultralytics fournit un paramètre freeze pour geler les couches au début de l'entraînement, et tu peux utiliser un rappel pour les dégeler après N époques :
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)Le paramètre freeze=10 gèle les 10 premières couches (indices 0 à 9) au début de l'entraînement, ce qui couvre la majeure partie du tronc de YOLO26. Le tronc s'étend des couches 0 à 10, donc freeze=10 laisse le bloc C2PSA final (couche 10) entraînable ; utilise freeze=11 pour geler l'intégralité du tronc. Le rappel on_train_epoch_start se déclenche au début de chaque époque et dégèle tous les paramètres une fois la période de gel terminée.
freeze=10gèle les 10 premières couches, indices 0-9 (la majeure partie du tronc de YOLO26 ; utilisefreeze=11pour inclure le bloc C2PSA final à la couche 10)freeze=[0, 1, 2, 3]gèle des couches spécifiques par index- Des valeurs de
FREEZE_EPOCHSplus élevées donnent à la tête plus de temps pour s'adapter avant que le tronc ne change
Taux d'apprentissage par couche#
Différentes parties du réseau peuvent bénéficier de taux d'apprentissage différents. Une stratégie courante consiste à utiliser un taux d'apprentissage plus faible pour le tronc pré-entraîné afin de préserver les caractéristiques apprises, tout en permettant à la tête de détection de s'adapter plus rapidement avec un taux plus élevé :
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)Variante RT-DETR#
Pour RT-DETR, le schéma est le même avec deux améliorations. La longueur du tronc est lue à partir de model.yaml["backbone"] de sorte que le même formateur fonctionne sur toutes les variantes de RT-DETR (troncs RT-DETR-L, RT-DETR-X, ResNet-50/101) sans coder en dur le nombre de couches. Les paramètres sont également répartis en groupes de poids, de BatchNorm et de biais dans chaque section afin que la décadence du poids soit exclue des paramètres de BatchNorm et des biais, ce qui correspond à la politique du formateur par défaut. Ceci est particulièrement utile pour le réglage fin de RT-DETR, où la tête du décodeur est généralement initialisée aléatoirement tandis que le tronc transporte des caractéristiques pré-entraînées qui bénéficient d'un taux d'apprentissage plus faible :
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)Un point de départ courant est backbone_lr_ratio = 0.1, correspondant à la configuration originale de RT-DETR avec son tronc HGNetV2. La littérature suggère de mettre à l'échelle le ratio inversement par rapport à la taille du tronc et à l'échelle des données de pré-entraînement : les grands troncs pré-entraînés sur des ensembles de données très vastes (par exemple ViT-L/H entraîné avec DINO, CLIP ou MAE sur des centaines de millions d'images) utilisent généralement des ratios plus petits tels que 0.01 ou moins pour préserver les caractéristiques bien apprises, tandis que les troncs plus petits avec un pré-entraînement plus léger tolèrent des ratios plus grands tels que 0.5 ou plus.
Le planificateur de taux d'apprentissage intégré (cosine ou linear) s'applique toujours en plus des taux d'apprentissage de base par groupe. Les taux d'apprentissage du tronc et de la tête suivront tous deux le même calendrier de décroissance, maintenant le ratio entre eux tout au long de l'entraînement.
Ces personnalisations peuvent être combinées dans une seule classe de trainer en surchargeant plusieurs méthodes et en ajoutant des callbacks selon les besoins.
BatchNorm synchronisé pour l'entraînement multi-GPU#
Lors de l'entraînement sur plusieurs GPU avec DistributedDataParallel, les couches de BatchNorm2d par défaut calculent les statistiques indépendamment sur chaque GPU. Pour le réglage fin de RT-DETR et d'autres recettes qui utilisent de petites tailles de lots par GPU, les statistiques de lots par GPU peuvent être bruitées. Le SyncBatchNorm de PyTorch synchronise la moyenne et la variance sur tous les rangs pour une statistique de lot globale unique, ce qui améliore souvent la convergence au prix d'un léger surcoût de communication inter-GPU.
La conversion doit avoir lieu après que le modèle est sur le GPU mais avant que DDP ne l'enveloppe. Le point d'ancrage le plus propre pour cela est set_model_attributes(), que BaseTrainer appelle exactement dans cette fenêtre :
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)La protection world_size > 1 garantit que le formateur peut également être utilisé en toute sécurité lors d'exécutions à GPU unique ; sur un seul GPU, la conversion est ignorée et l'entraînement se poursuit avec le BatchNorm2d habituel. Le même schéma fonctionne pour YOLO en basculant la classe parente vers DetectionTrainer.
| Scénario | Recommandation |
|---|---|
| Entraînement multi-GPU, petit batch par GPU (≤ 16) | Activer |
| Entraînement multi-GPU, grand batch par GPU (≥ 32) | Optionnel ; bénéfice mineur |
| Entraînement sur un seul GPU | Non applicable (sauté) |
Clipping de gradient configurable#
Le formateur par défaut découpe les gradients à max_norm=10.0 dans optimizer_step(), une valeur souple ajustée pour les modèles YOLO où les gradients la dépassent rarement. Les détecteurs de la famille DETR (RT-DETR, DEIM, DINO) utilisent généralement des valeurs beaucoup plus strictes telles que 0.1 pour stabiliser les couches d'attention croisée du décodeur, où l'amplitude des gradients peut monter en flèche. Pour remplacer la valeur de découpage, sous-classe le formateur et surcharge 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)Le même formateur fonctionne pour YOLO en basculant la classe parente vers DetectionTrainer (from ultralytics.models.yolo.detect import DetectionTrainer) et en chargeant un point de contrôle YOLO avec YOLO("yolo26n.pt"). Le corps de optimizer_step reste inchangé.
| Famille d'architecture | max_norm typique |
|---|---|
| Famille RT-DETR / DEIM / DETR | 0.1 |
| YOLO (défaut Ultralytics) | 10.0 |
| Désactiver le clipping | 0 |
FAQ#
Méthodes clés disponibles pour la personnalisation :
Méthode Objectif validate()Exécute la validation et renvoie les métriques build_optimizer()Construit l'optimiseur save_model()Enregistre les points de contrôle d'entraînement get_model()Renvoie l'instance du modèle get_validator()Renvoie l'instance du validateur get_dataloader()Construit le dataloader preprocess_batch()Prétraite le lot d'entrée label_loss_items()Formate les éléments de perte pour la journalisation Pour la référence complète de l'API, consulte la documentation de
BaseTrainer.Oui, pour des personnalisations plus simples, les rappels suffisent souvent. Les événements de rappel disponibles incluent
on_train_start,on_train_epoch_start,on_train_epoch_end,on_fit_epoch_endeton_model_save. Ceux-ci te permettent de te brancher sur la boucle d'entraînement sans avoir à créer de sous-classe. L'exemple de gel du tronc ci-dessus illustre cette approche.Si ta modification est plus simple (comme l'ajustement des gains de perte), tu peux modifier directement les hyperparamètres :
model.train(data="coco8.yaml", box=10.0, cls=1.5, dfl=2.0)Pour des modifications structurelles de la perte (comme l'ajout de poids de classes), tu dois sous-classer la perte et le modèle comme indiqué dans la section sur les poids de classes.
Passe ta classe de formateur personnalisé (et non une instance) au paramètre
trainerdansmodel.train():La classe
YOLOgère l'instanciation du formateur en interne. Consulte la page de Personnalisation Avancée pour plus de détails sur l'architecture du formateur.