Personalizzazione del Trainer

Q: Which BaseTrainer methods can I override?

Metodi chiave disponibili per la personalizzazione: Per il riferimento API completo, consultare la documentazione di BaseTrainer.

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

Sì, per personalizzazioni più semplici, i callback sono spesso sufficienti. Gli eventi di callback disponibili includono on_train_start, on_train_epoch_start, on_train_epoch_end, on_fit_epoch_end e on_model_save. Questi consentono di agganciarsi al ciclo di addestramento senza dover sottoclassare. L'esempio di congelamento del backbone sopra riportato dimostra questo approccio.

La pipeline di addestramento di Ultralytics è costruita attorno a BaseTrainer e trainer specifici per task come DetectionTrainer. Queste classi gestiscono il ciclo di addestramento, la validazione, il checkpointing e la registrazione in modo predefinito. Quando hai bisogno di maggiore controllo — monitorare metriche personalizzate, regolare la ponderazione della loss o implementare schemi di learning rate — puoi sottoclassificare il trainer e sovrascrivere metodi specifici.

Questa guida illustra cinque personalizzazioni comuni:

Registrazione di metriche personalizzate (punteggio F1) alla fine di ogni epoca
Aggiungere pesi di classe per gestire lo squilibrio di classe
Salvataggio del modello migliore basato su una metrica diversa
Congelare il backbone per le prime N epoche, quindi scongelarlo
Specificare i tassi di apprendimento per strato

Prerequisiti

Prima di leggere questa guida, assicurati di avere familiarità con le basi di addestramento di modelli YOLO che l' Personalizzazione avanzata pagina, che copre il BaseTrainer architettura.

Come funzionano i Custom Trainer

Il YOLO la classe del modello accetta un trainer parametro in the train() metodo. Questo ti permette di passare la tua classe di trainer che estende il comportamento predefinito:

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)

Il tuo trainer personalizzato eredita tutte le funzionalità da DetectionTrainer, quindi devi solo override i metodi specifici che desideri personalizzare.

Registrazione di metriche personalizzate

Il convalida il passaggio calcola precisione, richiamo, e mAP. Se hai bisogno di metriche aggiuntive come quelle per classe Punteggio F1, override 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)

Questo registra il punteggio F1 medio su tutte le classi e una ripartizione per classe dopo ogni esecuzione di validazione.

Metriche Disponibili

Il validatore fornisce accesso a molte metriche tramite self.validator.metrics.box:

Attributo	Descrizione
`f1`	Punteggio F1 per classe
`p`	Precisione per classe
`r`	Recall per classe
`ap50`	AP con IoU 0.5 per classe
`ap`	AP con IoU 0.5:0.95 per classe
`mp`, `mr`	Precisione e recall medie
`map50`, `map`	Metriche AP medie

Aggiunta di pesi di classe

Se il tuo dataset presenta classi sbilanciate (ad esempio, un difetto raro nell'ispezione di produzione), puoi aumentare il peso delle classi sottorappresentate nella funzione di perdita. Ciò fa sì che il modello penalizzi più severamente le errate classificazioni delle classi rare.

Per personalizzare la funzione di perdita, sottoclassa le classi di perdita, il modello e il 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)

Calcolo dei pesi dal dataset

È possibile calcolare automaticamente i pesi delle classi dalla distribuzione delle etichette del proprio dataset. Un approccio comune è la ponderazione per frequenza inversa:

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]

Salvataggio del modello migliore tramite metrica personalizzata

Il trainer salva best.pt basato sulla fitness, che di default è 0.9 × mAP@0.5:0.95 + 0.1 × mAP@0.5. Per utilizzare una metrica diversa (come mAP@0.5 o richiamo), sovrascrivi validate() e restituire la metrica scelta come valore di fitness. Il modulo integrato save_model() lo utilizzerà quindi automaticamente:

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)

Metriche Disponibili

Metriche comuni disponibili in self.metrics dopo la convalida includono:

Chiave	Descrizione
`metrics/precision(B)`	Precisione
`metrics/recall(B)`	Recall
`metrics/mAP50(B)`	mAP con IoU 0.5
`metrics/mAP50-95(B)`	mAP con IoU 0.5:0.95

Congelamento e Scongelamento del Backbone

Transfer learning i workflow spesso beneficiano del congelamento del backbone pre-addestrato per le prime N epoche, permettendo alla detection head di adattarsi prima fine-tuning l'intera rete. Ultralytics fornisce un freeze parametro per congelare gli strati all'inizio dell'addestramento, e si può usare un callback per sbloccarli dopo N epoche:

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)

Il freeze=10 parametro congela i primi 10 strati (il backbone) all'inizio dell'addestramento. Il on_train_epoch_start il callback si attiva all'inizio di ogni epoca e sblocca tutti i parametri una volta completato il periodo di blocco.

Scelta di Cosa Congelare

freeze=10 congela i primi 10 strati (tipicamente il backbone nelle architetture YOLO)
freeze=[0, 1, 2, 3] congela specifici strati per indice
Maggiore FREEZE_EPOCHS i valori danno alla testa più tempo per adattarsi prima che il backbone cambi

Tassi di apprendimento per strato

Diverse parti della rete possono beneficiare di diversi tassi di apprendimento. Una strategia comune è utilizzare un tasso di apprendimento inferiore per il backbone pre-addestrato al fine di preservare le caratteristiche apprese, consentendo al contempo all'head di rilevamento di adattarsi più rapidamente con un tasso più elevato:

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)

Scheduler del Learning Rate

Lo scheduler del tasso di apprendimento integrato (cosine oppure linear) si applica comunque in aggiunta ai tassi di apprendimento base per gruppo. Sia i tassi di apprendimento del backbone che dell'head seguiranno lo stesso programma di decadimento, mantenendo il rapporto tra di essi durante l'intero addestramento.

Combinazione di Tecniche

Queste personalizzazioni possono essere combinate in un'unica classe di trainer sovrascrivendo più metodi e aggiungendo callback secondo necessità.

FAQ

Come si passa un custom trainer a YOLO?

Passa la tua classe di trainer personalizzata (non un'istanza) al trainer parametro in model.train():

from ultralytics import YOLO

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

Il YOLO la classe gestisce internamente l'istanziazione del trainer. Vedi il Personalizzazione avanzata pagina per maggiori dettagli sull'architettura del trainer.

Quali metodi di BaseTrainer posso sovrascrivere?

Metodi chiave disponibili per la personalizzazione:

Metodo	Scopo
`validate()`	Esegui la validazione e restituisci le metriche
`build_optimizer()`	Costruire l'ottimizzatore
`save_model()`	Salva i checkpoint di addestramento
`get_model()`	Restituisci l'istanza del modello
`get_validator()`	Restituisce l'istanza del validatore.
`get_dataloader()`	Crea il dataloader
`preprocess_batch()`	Pre-elaborare il batch di input
`label_loss_items()`	Formattazione degli elementi di perdita per il logging

Per il riferimento API completo, consulta il BaseTrainer documentazione.

Posso usare i callback invece di sottoclassare il trainer?

Sì, per personalizzazioni più semplici, callback sono spesso sufficienti. Gli eventi di callback disponibili includono on_train_start, on_train_epoch_start, on_train_epoch_end, on_fit_epoch_end, e on_model_save. Questi ti permettono di agganciarti al ciclo di addestramento senza sottoclassificazione. L'esempio di congelamento del backbone sopra dimostra questo approccio.

Come si personalizza la funzione di perdita senza sottoclassare il modello?

Se la tua modifica è più semplice (come la regolazione dei guadagni di perdita), puoi modificare direttamente gli iperparametri:

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

Per modifiche strutturali alla funzione di perdita (come l'aggiunta di pesi di classe), è necessario sottoclassare la funzione di perdita e il modello come mostrato nella sezione sui pesi di classe.

📅 Creato 1 mese fa ✏️ Aggiornato 1 mese fa