Personalización del Entrenador

Q: Which BaseTrainer methods can I override?

Métodos clave disponibles para personalización: Para la referencia completa de la API, consulte la documentación de BaseTrainer.

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

Sí, para personalizaciones más sencillas, los callbacks suelen ser suficientes. Los eventos de callback disponibles incluyen on_train_start, on_train_epoch_start, on_train_epoch_end, on_fit_epoch_end y on_model_save. Estos le permiten conectarse al bucle de entrenamiento sin necesidad de subclases. El ejemplo de congelación del backbone anterior demuestra este enfoque.

El pipeline de entrenamiento de Ultralytics se basa en BaseTrainer y entrenadores específicos para cada tarea, como DetectionTrainer. Estas clases gestionan el bucle de entrenamiento, la validación, los puntos de control y el registro de forma predeterminada. Cuando necesite más control —como el seguimiento de métricas personalizadas, el ajuste de la ponderación de la pérdida o la implementación de programas de tasa de aprendizaje— puede crear una subclase del entrenador y anular métodos específicos.

Esta guía detalla cinco personalizaciones comunes:

Registro de métricas personalizadas (puntuación F1) al final de cada época
Adición de pesos de clase para manejar el desequilibrio de clases
Guardar el mejor modelo basado en una métrica diferente
Congelar el backbone durante las primeras N épocas, luego descongelar
Especificación de tasas de aprendizaje por capa

Prerrequisitos

Antes de leer esta guía, asegúrate de estar familiarizado con los conceptos básicos de entrenamiento de modelos YOLO como el Personalización avanzada página, que cubre el BaseTrainer arquitectura.

Cómo Funcionan los Entrenadores Personalizados

El YOLO la clase del modelo acepta un trainer parámetro en la train() método. Esto le permite pasar su propia clase de entrenador que extiende el comportamiento predeterminado:

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)

Su entrenador personalizado hereda toda la funcionalidad de DetectionTrainer, por lo que solo necesita sobrescribir los métodos específicos que desea personalizar.

Registro de métricas personalizadas

El validación el paso calcula precisión, exhaustividad, y mAP. Si necesita métricas adicionales como por clase Puntuación F1, sobrescribir 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)

Esto registra la puntuación F1 media en todas las clases y un desglose por clase después de cada ejecución de validación.

Métricas Disponibles

El validador proporciona acceso a numerosas métricas a través de self.validator.metrics.box:

Atributo	Descripción
`f1`	Puntuación F1 por clase
`p`	Precisión por clase
`r`	Recall por clase
`ap50`	AP en IoU 0.5 por clase
`ap`	AP en IoU 0.5:0.95 por clase
`mp`, `mr`	Precisión y recall promedio
`map50`, `map`	Métricas de AP promedio

Añadir Pesos de Clase

Si su conjunto de datos tiene clases desequilibradas (por ejemplo, un defecto raro en la inspección de fabricación), puede ponderar más las clases subrepresentadas en la función de pérdida. Esto hace que el modelo penalice más severamente las clasificaciones erróneas en clases raras.

Para personalizar la función de pérdida, subclase las clases de pérdida, el modelo y el entrenador:

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)

Cálculo de pesos a partir del conjunto de datos

Puede calcular los pesos de clase automáticamente a partir de la distribución de etiquetas de su conjunto de datos. Un enfoque común es la ponderación por frecuencia 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]

Guardando el mejor modelo por métrica personalizada

El entrenador guarda best.pt basado en la aptitud, que por defecto es 0.9 × mAP@0.5:0.95 + 0.1 × mAP@0.5. Para usar una métrica diferente (como mAP@0.5 o recall), anular validate() y devolver la métrica elegida como valor de aptitud. El incorporado save_model() lo utilizará automáticamente:

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étricas Disponibles

Métricas comunes disponibles en self.metrics después de la validación incluyen:

Clave	Descripción
`metrics/precision(B)`	Precisión
`metrics/recall(B)`	Recall
`metrics/mAP50(B)`	mAP con IoU 0.5
`metrics/mAP50-95(B)`	mAP con IoU 0.5:0.95

Congelación y Descongelación del Backbone

Aprendizaje por transferencia Los flujos de trabajo a menudo se benefician de congelar el backbone preentrenado durante las primeras N épocas, permitiendo que el cabezal de detección se adapte antes de. ajuste fino toda la red. Ultralytics proporciona una freeze parámetro para congelar capas al inicio del entrenamiento, y se puede utilizar un callback para descongelarlos después de N épocas:

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)

El freeze=10 parámetro congela las primeras 10 capas (el backbone) al inicio del entrenamiento. El on_train_epoch_start El callback se activa al inicio de cada época y descongela todos los parámetros una vez que el período de congelación ha finalizado.

Elegir qué congelar

freeze=10 congela las primeras 10 capas (típicamente el backbone en arquitecturas YOLO)
freeze=[0, 1, 2, 3] congela capas específicas por índice
Mayor FREEZE_EPOCHS los valores dan a la cabeza más tiempo para adaptarse antes de que el backbone cambie

Tasas de aprendizaje por capa

Diferentes partes de la red pueden beneficiarse de diferentes tasas de aprendizaje. Una estrategia común es utilizar una tasa de aprendizaje más baja para el *backbone* preentrenado para preservar las características aprendidas, mientras se permite que el cabezal de detección se adapte más rápidamente con una tasa más alta:

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)

Planificador de Tasa de Aprendizaje

El programador de tasa de aprendizaje incorporado (cosine o linear) sigue aplicándose además de las tasas de aprendizaje base por grupo. Tanto las tasas de aprendizaje del backbone como las del head seguirán el mismo programa de decaimiento, manteniendo la relación entre ellas a lo largo del entrenamiento.

Combinación de Técnicas

Estas personalizaciones se pueden combinar en una única clase de entrenador sobrescribiendo múltiples métodos y añadiendo callbacks según sea necesario.

Preguntas frecuentes

¿Cómo paso un entrenador personalizado a YOLO?

Pase su clase de entrenador personalizada (no una instancia) al trainer parámetro en model.train():

from ultralytics import YOLO

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

El YOLO la clase maneja la instanciación del entrenador internamente. Consulte el Personalización avanzada página para más detalles sobre la arquitectura del entrenador.

¿Qué métodos de BaseTrainer puedo sobrescribir?

Métodos clave disponibles para personalización:

Método	Propósito
`validate()`	Ejecutar validación y devolver métricas
`build_optimizer()`	Construir el optimizador
`save_model()`	Guardar puntos de control de entrenamiento
`get_model()`	Devolver la instancia del modelo
`get_validator()`	Devolver la instancia del validador
`get_dataloader()`	Construir el dataloader
`preprocess_batch()`	Preprocesar lote de entrada
`label_loss_items()`	Formatear elementos de pérdida para registro

Para la referencia completa de la API, consulte el BaseTrainer documentación.

¿Puedo usar callbacks en lugar de subclasear el entrenador?

Sí, para personalizaciones más sencillas, Callbacks suelen ser suficientes. Los eventos de callback disponibles incluyen on_train_start, on_train_epoch_start, on_train_epoch_end, on_fit_epoch_end, y on_model_save. Estos permiten integrarse en el bucle de entrenamiento sin necesidad de subclases. El ejemplo de congelación del backbone anterior demuestra este enfoque.

¿Cómo personalizo la función de pérdida sin subclasear el modelo?

Si su cambio es más simple (como ajustar las ganancias de pérdida), puede modificar los hiperparámetros directamente:

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

Para cambios estructurales en la función de pérdida (como la adición de pesos de clase), es necesario subclasificar la función de pérdida y el modelo, tal como se detalla en la sección de pesos de clase.

📅 Creado hace 1 mes ✏️ Actualizado hace 1 mes