Personnalisation du formateur

Q: Can I use callbacks instead of subclassing the trainer?

Oui, pour des personnalisations plus simples, les rappels sont souvent suffisants. Les événements de rappel disponibles incluent on_train_start, on_train_epoch_start, on_train_epoch_end, on_fit_epoch_end et on_model_save. Ceux-ci vous permettent de vous connecter à la boucle d'entraînement sans sous-classer. L'exemple de gel de la colonne vertébrale ci-dessus illustre cette approche.

Le pipeline Ultralytics s'articule autour de BaseTrainer et des formateurs spécialisés dans des tâches spécifiques, tels que DetectionTrainerCes classes gèrent la boucle d'entraînement, la validation, la création de points de contrôle et la connexion dès leur installation. Si vous avez besoin d'un contrôle accru (suivi de métriques personnalisées, ajustement de la pondération des pertes ou mise en œuvre de calendriers de taux d'apprentissage), vous pouvez créer une sous-classe du programme d'entraînement et remplacer des méthodes spécifiques.

Ce guide présente cinq personnalisations courantes :

Enregistrement des métriques personnalisées (score F1) à la fin de chaque époque
Ajout de pondérations de classe pour gérer le déséquilibre des classes
Enregistrer le meilleur modèle en fonction d'un indicateur différent
Gel de la colonne vertébrale pendant les N premières époques, puis dégel
Spécification des taux d'apprentissage par couche

Prérequis

Avant de lire ce guide, assurez-vous de bien connaître les bases de formation YOLO et le Personnalisation avancée page, qui couvre le BaseTrainer architecture.

Comment fonctionnent les entraîneurs personnalisés

L'argument YOLO La classe modèle accepte un trainer paramètre dans le train() méthode. Cela vous permet de passer votre 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."""

    pass  # Add your customizations here


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

Votre formateur personnalisé hérite de toutes les fonctionnalités de DetectionTrainer, vous n'avez donc besoin de remplacer que les méthodes spécifiques que vous souhaitez personnaliser.

Enregistrement des métriques personnalisées

L'argument validation étape de calcul précision, rappel, et mAP. Si vous avez besoin de mesures supplémentaires telles que par classe Score F1, remplacer 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)

Cela enregistre le score F1 moyen pour toutes les classes et une ventilation par classe après chaque cycle de validation.

Métriques disponibles

Le validateur donne accès à de nombreuses métriques via self.validator.metrics.box:

Attribut	Description
`f1`	Score F1 par classe
`p`	Précision par classe
`r`	Rappel par classe
`ap50`	AP IoU ,5 par classe
`ap`	AP IoU ,5 : 0,95 par classe
`mp`, `mr`	Précision et rappel moyens
`map50`, `map`	AP moyennes

Ajouter des pondérations de classe

Si votre ensemble de données comporte des classes déséquilibrées (par exemple, un défaut rare dans le contrôle de fabrication), vous pouvez pondérer les classes sous-représentées dans la fonction de perte. Cela permet au modèle de pénaliser plus lourdement les erreurs de classification sur les classes rares.

Pour personnaliser la perte, créez des sous-classes pour les classes de perte, le modèle et le formateur :

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)

Calcul des pondérations à partir d'un ensemble de données

Vous pouvez calculer automatiquement les pondérations des classes à partir de la distribution des étiquettes de votre ensemble de données. Une approche courante consiste à utiliser 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]

Enregistrer le meilleur modèle à l'aide d'une métrique personnalisée

Le formateur enregistre best.pt basé sur la condition physique, dont la valeur par défaut est 0.9 × mAP@0.5:0.95 + 0.1 × mAP@0.5. Pour utiliser une autre mesure (comme mAP@0.5 ou rappel), passer outre validate() et renvoie la métrique choisie comme valeur d'aptitude. La fonction intégrée save_model() 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)

Métriques disponibles

Mesures courantes disponibles dans self.metrics après validation, comprennent :

Clé	Description
`metrics/precision(B)`	Précision
`metrics/recall(B)`	Rappel
`metrics/mAP50(B)`	mAP IoU ,5
`metrics/mAP50-95(B)`	mAP IoU ,5 : 0,95

Gel et dégel de la colonne vertébrale

Apprentissage par transfert Les flux de travail bénéficient souvent du gel de la structure pré-entraînée pendant les N premiers cycles, ce qui permet à la tête de détection de s'adapter avant ajustement l'ensemble du réseau. Ultralytics un freeze paramètre pour geler les couches au début de l'entraînement, et vous pouvez utiliser un callback 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)

L'argument freeze=10 Le paramètre gèle les 10 premières couches (la colonne vertébrale) au début de l'entraînement. Le on_train_epoch_start Le rappel se déclenche au début de chaque époque et débloque tous les paramètres une fois la période de gel terminée.

Choisir ce qu'il faut congeler

freeze=10 gèle les 10 premières couches (généralement la colonne vertébrale dans YOLO )
freeze=[0, 1, 2, 3] gèle des couches spécifiques par index
Plus élevé FREEZE_EPOCHS les valeurs donnent à la tête plus de temps pour s'adapter avant que la colonne vertébrale ne change

Taux d'apprentissage par couche

Différentes parties du réseau peuvent bénéficier de différents taux d'apprentissage. Une stratégie courante consiste à utiliser un taux d'apprentissage plus faible pour la structure pré-entraînée 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 = []

        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)

Planificateur du taux d'apprentissage

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 de la colonne vertébrale et de la tête suivront le même calendrier de décroissance, en maintenant le rapport entre eux tout au long de la formation.

Combinaison de techniques

Ces personnalisations peuvent être combinées dans une seule classe de formateur en remplaçant plusieurs méthodes et en ajoutant des rappels si nécessaire.

FAQ

Comment transmettre un entraîneur personnalisé à YOLO?

Transmettez votre classe de formateur personnalisée (et non une instance) à la trainer dans model.train():

from ultralytics import YOLO

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

L'argument YOLO La classe gère l'instanciation du formateur en interne. Voir le Personnalisation avancée page pour plus de détails sur l'architecture du formateur.

Quelles méthodes BaseTrainer puis-je remplacer ?

Principales méthodes disponibles pour la personnalisation :

Méthode	Objectif
`validate()`	Exécuter la validation et renvoyer les métriques
`build_optimizer()`	Construire l'optimiseur
`save_model()`	Enregistrer les points de contrôle de formation
`get_model()`	Renvoyer l'instance du modèle
`get_validator()`	Renvoyer l'instance du validateur
`get_dataloader()`	Construire le chargeur de données
`preprocess_batch()`	Pré-traitement du lot d'entrée
`label_loss_items()`	Éléments de perte de format pour la journalisation

Pour la référence API complète, consultez le BaseTrainer documentation.

Puis-je utiliser des rappels au lieu de sous-classer le formateur ?

Oui, pour des personnalisations plus simples, rappels sont souvent suffisants. Les événements de rappel disponibles incluent on_train_start, on_train_epoch_start, on_train_epoch_end, on_fit_epoch_end, et on_model_save. Ceux-ci vous permettent de vous connecter à la boucle de formation sans créer de sous-classe. L'exemple de gel de la colonne vertébrale ci-dessus illustre cette approche.

Comment personnaliser la fonction de perte sans créer de sous-classe du modèle ?

Si votre modification est plus simple (comme l'ajustement des gains de perte), vous pouvez modifier directement les hyperparamètres:

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

Pour apporter des modifications structurelles à la perte (telles que l'ajout de pondérations de classe), vous devez créer une sous-classe de la perte et du modèle, comme indiqué dans la section sur les pondérations de classe.

📅 Créé il y a 6 jours ✏️ Mis à jour il y a 0 jours