Personalización del entrenador
El proceso Ultralytics de Ultralytics se basa en BaseTrainer y formadores especializados en tareas específicas, 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 (por ejemplo, para realizar un seguimiento de métricas personalizadas, ajustar la ponderación de pérdidas o implementar programas de velocidad de aprendizaje), puede crear una subclase del entrenador y anular métodos específicos.
Esta guía explica cinco personalizaciones comunes:
- Registro de métricas personalizadas (puntuación F1) al final de cada época.
- Añadir pesos a las clases para gestionar el desequilibrio entre ellas.
- Guardar el mejor modelo basado en una métrica diferente
- Congelar la columna vertebral durante las primeras N épocas y luego descongelarla.
- Especificar las 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 YOLO como el Personalización avanzada página, que cubre el BaseTrainer arquitectura.
Cómo funcionan los entrenadores personalizados
El YOLO La clase modelo acepta un trainer parámetro en la train() método. Esto le permite pasar su propia clase de entrenador que amplía 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)
Tu entrenador personalizado hereda todas las funciones de DetectionTrainer, por lo que solo es necesario sobrescribir los métodos específicos que desee personalizar.
Registro de métricas personalizadas
El validación paso calcula precisión, exhaustividad, y mAP. Si necesita métricas adicionales, como por clase Puntuación F1, anular 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 muchas 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 | Recuerdo por clase |
ap50 | AP IoU ,5 por clase |
ap | AP IoU ,5:0,95 por clase |
mp, mr | Precisión y recuperación medias |
map50, map | AP medias |
Añadir ponderaciones de clase
Si su conjunto de datos tiene clases desequilibradas (por ejemplo, un defecto poco frecuente en la inspección de fabricación), puede aumentar la ponderación de las clases infrarrepresentadas en la función de pérdida. Esto hace que el modelo penalice más severamente las clasificaciones erróneas en las clases poco frecuentes.
Para personalizar la pérdida, subclasifique 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 ponderaciones a partir del conjunto de datos
Puede calcular automáticamente las ponderaciones de clase 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]
Guardar el mejor modelo por métrica personalizada
El entrenador salva best.pt basado en el estado físico, cuyo valor predeterminado es 0.9 × mAP@0.5:0.95 + 0.1 × mAP@0.5. Para utilizar una métrica diferente (como mAP@0.5 o recall), anular validate() y devuelve la métrica elegida como valor de aptitud. El integrado 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, se incluyen:
| Clave | Descripción |
|---|---|
metrics/precision(B) | Precisión |
metrics/recall(B) | Recall |
metrics/mAP50(B) | mAP IoU ,5 |
metrics/mAP50-95(B) | mAP IoU ,5:0,95 |
Congelación y descongelación de la columna vertebral
Aprendizaje por transferencia Los flujos de trabajo suelen beneficiarse de la congelación de la estructura preentrenada durante las primeras N épocas, lo que permite que el cabezal de detección se adapte antes de ajuste fino toda la red. Ultralytics un freeze parámetro para congelar capas al inicio del entrenamiento, y puede utilizar un callback 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 El parámetro congela las primeras 10 capas (la columna vertebral) al inicio del entrenamiento. El on_train_epoch_start La devolución de llamada se activa al comienzo de cada época y descongela todos los parámetros una vez finalizado el período de congelación.
Elegir qué congelar
freeze=10congela las primeras 10 capas (normalmente la columna vertebral en YOLO )freeze=[0, 1, 2, 3]congela capas específicas por índice- Mayor
FREEZE_EPOCHSLos valores dan más tiempo a la cabeza para adaptarse antes de que cambie la columna vertebral.
Tasas de aprendizaje por capa
Las diferentes partes de la red pueden beneficiarse de diferentes tasas de aprendizaje. Una estrategia habitual consiste en utilizar una tasa de aprendizaje más baja para la estructura preentrenada con el fin de conservar las características aprendidas, al tiempo que 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)
Programador de la tasa de aprendizaje
El programador de velocidad de aprendizaje integrado (cosine o linear) sigue aplicándose además de las tasas de aprendizaje básicas por grupo. Tanto la tasa de aprendizaje de la columna vertebral como la de la cabeza seguirán el mismo programa de decaimiento, manteniendo la proporción entre ellas durante todo el entrenamiento.
Combinación de técnicas
Estas personalizaciones se pueden combinar en una única clase de entrenador sobrescribiendo varios métodos y añadiendo devoluciones de llamada 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 gestiona la instanciación del entrenador internamente. Consulte el Personalización avanzada página para obtener más detalles sobre la arquitectura del entrenador.
¿Qué métodos de BaseTrainer puedo sobrescribir?
Métodos clave disponibles para la personalización:
| Método | Propósito |
|---|---|
validate() | Ejecutar la validación y devolver las métricas. |
build_optimizer() | Construir el optimizador |
save_model() | Guardar puntos de control de entrenamiento |
get_model() | Devuelve la instancia del modelo. |
get_validator() | Devuelve la instancia del validador. |
get_dataloader() | Crear el cargador de datos |
preprocess_batch() | Preprocesar lote de entrada |
label_loss_items() | Formatear elementos perdidos para su registro |
Para obtener la referencia completa de la API, consulte el BaseTrainer documentación.
¿Puedo usar devoluciones de llamada en lugar de crear una subclase del entrenador?
Sí, para personalizaciones más sencillas, devoluciones de llamada suelen ser suficientes. Los eventos de devolución de llamada disponibles incluyen on_train_start, on_train_epoch_start, on_train_epoch_end, on_fit_epoch_end, y on_model_save. Estos te permiten conectarte al bucle de entrenamiento sin crear subclases. El ejemplo anterior de congelación de backbone ilustra este enfoque.
¿Cómo puedo personalizar la función de pérdida sin crear una subclase del modelo?
Si el cambio es más sencillo (como ajustar las ganancias por pérdidas), puede modificar los hiperparámetros directamente:
model.train(data="coco8.yaml", box=10.0, cls=1.5, dfl=2.0)
Para realizar cambios estructurales en la pérdida (como añadir ponderaciones de clase), es necesario crear una subclase de la pérdida y el modelo, tal y como se muestra en la sección sobre ponderaciones de clase.
