Personalizar o Trainer#
O pipeline de treino da Ultralytics é construído em torno de BaseTrainer e de trainers específicos para cada tarefa, como DetectionTrainer. Estas classes tratam, de imediato, do ciclo de treino, da validação, dos checkpoints e do registo. Quando precisares de mais controlo — acompanhar métricas personalizadas, ajustar a ponderação da loss ou implementar agendamentos da taxa de aprendizagem — podes criar uma subclasse do trainer e substituir métodos específicos.
Este guia aborda sete personalizações comuns:
- Registar métricas personalizadas (pontuação F1) no final de cada época
- Adicionar pesos às classes para lidar com o desequilíbrio entre classes
- Guardar o melhor modelo com base numa métrica diferente
- Congelar o backbone durante as primeiras N épocas e depois descongelá-lo
- Especificar taxas de aprendizagem por camada
- Sincronizar o BatchNorm entre GPUs para treino com várias GPUs
- Configurar o gradient clipping para ajustar a estabilidade
Antes de leres este guia, certifica-te de que conheces os conceitos básicos de treino de modelos YOLO e a página Personalização avançada, que aborda a arquitetura BaseTrainer.
Como funcionam os Custom Trainers#
A classe de modelo YOLO aceita um parâmetro trainer no método train(). Isto permite passar a tua própria classe de trainer, que estende o comportamento predefinido:
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)O teu trainer personalizado herda todas as funcionalidades de DetectionTrainer, por isso só precisas de substituir os métodos específicos que queres personalizar.
Registar métricas personalizadas#
A etapa de validação calcula precision, recall e mAP. Se precisares de métricas adicionais, como a pontuação F1 por classe, substitui 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
mean_f1 = float(np.mean(f1_per_class)) if len(f1_per_class) 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)]
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)Isto regista a pontuação F1 média de todas as classes representadas na validação e uma discriminação por classe após cada execução da validação.
O validador fornece acesso a várias métricas através de self.validator.metrics.box:
| Atributo | Descrição |
|---|---|
f1 | Pontuação F1 por classe |
image_metrics | Dicionário de métricas por imagem com precision, recall, F1, TP, FP e FN |
p | Precision por classe |
r | Recall por classe |
ap50 | AP com IoU 0.5 por classe |
ap | AP com IoU 0.5:0.95 por classe |
mp, mr | Precision e recall médios |
map50, map | Métricas de AP médias |
Adicionar pesos às classes#
Define cls_pw entre 0.0 e 1.0 para aplicar pesos normalizados de frequência inversa à loss de classificação. Substitui o cálculo de pesos existente apenas quando precisares de proporções definidas manualmente:
import numpy as np
from ultralytics import YOLO
from ultralytics.models.yolo.detect import DetectionTrainer
class WeightedTrainer(DetectionTrainer):
"""Detection trainer with hand-picked class-weight ratios."""
def compute_class_weights(self, class_counts):
"""Return custom per-class weights for the production loss owner."""
weights = np.ones_like(class_counts)
weights[0] = 2.0
weights[1] = 3.0
return weights
model = YOLO("yolo26n.pt")
model.train(data="custom.yaml", epochs=10, cls_pw=1.0, trainer=WeightedTrainer)set_class_weights() normaliza estes valores para uma média de 1.0 e armazena-os no modelo, onde a loss de deteção existente os aplica. Os índices acima requerem um dataset com pelo menos duas classes.
Guardar o melhor modelo por métrica personalizada#
O trainer guarda best.pt com base no fitness, que, para deteção, por predefinição corresponde a mAP@0.5:0.95 (os pesos [0.0, 0.0, 0.0, 1.0] para [P, R, mAP@0.5, mAP@0.5:0.95]). Para utilizares uma métrica diferente (como mAP@0.5 ou recall), substitui validate() e devolve a métrica escolhida como valor de fitness. O save_model() integrado irá então utilizá-la 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."""
previous_best = self.best_fitness
metrics, fitness = super().validate()
if metrics is None:
return metrics, fitness
fitness = metrics["metrics/mAP50(B)"]
self.best_fitness = fitness if previous_best is None else max(previous_best, fitness)
return metrics, fitness
model = YOLO("yolo26n.pt")
model.train(data="coco8.yaml", epochs=20, trainer=CustomSaveTrainer)BaseTrainer.validate() atualiza best_fitness usando a métrica predefinida, por isso guarda o valor anterior antes de o chamares.
As métricas comuns disponíveis em self.metrics após a validação incluem:
| Chave | Descrição |
|---|---|
metrics/precision(B) | Precision |
metrics/recall(B) | Recall |
metrics/mAP50(B) | mAP com IoU 0.5 |
metrics/mAP50-95(B) | mAP com IoU 0.5:0.95 |
Congelar e descongelar o backbone#
Os fluxos de trabalho de transfer learning beneficiam frequentemente de congelar o backbone pré-treinado durante as primeiras N épocas, permitindo que a cabeça de deteção se adapte antes de fazer o fine-tuning de toda a rede. A Ultralytics fornece um parâmetro freeze para congelar camadas no início do treino, e podes usar um callback para as descongelar após 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 the user-requested layers after FREEZE_EPOCHS."""
if trainer.epoch == FREEZE_EPOCHS:
user_freeze = [x for x in trainer.freeze_layer_names if x not in {".dfl", "teacher_model."}]
LOGGER.info(f"Epoch {trainer.epoch}: Unfreezing requested layers for fine-tuning")
for name, param in trainer.model.named_parameters():
if (
not param.requires_grad
and ".dfl" not in name
and "teacher_model." not in name
and any(x in name for x in user_freeze)
):
param.requires_grad = True
LOGGER.info(f" Unfroze: {name}")
trainer.freeze_layer_names = [x for x in trainer.freeze_layer_names if x not in user_freeze]
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)O parâmetro freeze=10 congela as primeiras 10 camadas (índices 0-9) no início do treino, abrangendo a maior parte do backbone do YOLO26. O backbone estende-se pelas camadas 0-10, por isso freeze=10 deixa o bloco C2PSA final (camada 10) treinável; usa freeze=11 para congelar todo o backbone. O callback on_train_epoch_start é executado no início de cada época e descongela as camadas solicitadas quando o período de congelamento termina, preservando os parâmetros permanentemente congelados do DFL e do professor de destilação.
freeze=10congela as primeiras 10 camadas, índices 0-9 (a maior parte do backbone do YOLO26; usafreeze=11para incluir o bloco C2PSA final na camada 10)freeze=[0, 1, 2, 3]congela camadas específicas por índice- Valores mais altos de
FREEZE_EPOCHSdão à cabeça mais tempo para se adaptar antes de o backbone mudar
Taxas de aprendizagem por camada#
Diferentes partes da rede podem beneficiar de taxas de aprendizagem diferentes. Uma estratégia comum é usar uma taxa de aprendizagem mais baixa para o backbone pré-treinado, preservando as características aprendidas, e permitir que a cabeça de deteção se adapte mais rapidamente com uma taxa mais alta:
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."""
backbone_lr_ratio = 0.1
def build_optimizer(self, model, name="auto", lr=0.001, momentum=0.9, decay=1e-5, iterations=1e5):
"""Reuse the trainer optimizer and lower its backbone parameter-group rates."""
optimizer = super().build_optimizer(model, name, lr, momentum, decay, iterations)
unwrapped = unwrap_model(model)
backbone_len = len(unwrapped.yaml["backbone"])
backbone = {
id(p)
for name, p in unwrapped.named_parameters()
if any(name.startswith(f"model.{i}.") for i in range(backbone_len))
}
groups = []
for group in optimizer.param_groups:
head_params = [p for p in group["params"] if id(p) not in backbone]
backbone_params = [p for p in group["params"] if id(p) in backbone]
if head_params:
groups.append({**group, "params": head_params})
if backbone_params:
groups.append({**group, "params": backbone_params, "lr": group["lr"] * self.backbone_lr_ratio})
optimizer.param_groups = groups
LOGGER.info(f"PerLayerLR: {len(backbone)} backbone params at {self.backbone_lr_ratio}x the head rate")
return optimizer
model = YOLO("yolo26n.pt")
model.train(data="coco8.yaml", epochs=20, trainer=PerLayerLRTrainer)Variante RT-DETR#
Para RT-DETR, usa a mesma substituição, com RTDETRTrainer como classe-pai, e carrega o checkpoint com RTDETR("rtdetr-l.pt").
BatchNorm sincronizado para treino com várias GPUs#
Ao treinar em várias GPUs com DistributedDataParallel, as camadas BatchNorm2d predefinidas calculam estatísticas de forma independente em cada GPU. No fine-tuning de RT-DETR e noutras receitas que usam tamanhos de batch pequenos por GPU, as estatísticas por GPU podem ser ruidosas. O SyncBatchNorm do PyTorch sincroniza a média e a variância entre todos os ranks para obter uma única estatística de batch global, o que frequentemente melhora a convergência ao custo de uma pequena sobrecarga de comunicação entre GPUs.
A conversão tem de ocorrer depois de o modelo estar na GPU, mas antes de o DDP o envolver. O hook mais simples para isto é set_model_attributes(), que BaseTrainer chama exatamente nesse intervalo:
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)A verificação world_size > 1 garante que o trainer também pode ser usado com segurança em execuções numa única GPU; numa única GPU, a conversão é ignorada e o treino prossegue com o BatchNorm2d normal. O mesmo padrão funciona para YOLO, alterando a classe-pai para DetectionTrainer.
| Cenário | Recomendação |
|---|---|
| Treino com várias GPUs, batch pequeno por GPU (≤ 16) | Ativar |
| Treino com várias GPUs, batch grande por GPU (≥ 32) | Opcional; benefício reduzido |
| Treino numa única GPU | Não aplicável (ignorado) |
Gradient clipping configurável#
O trainer predefinido limita os gradientes a max_norm=10.0 em optimizer_step(), um valor permissivo ajustado para modelos YOLO, nos quais os gradientes raramente o ultrapassam. Os detetores da família DETR (RT-DETR, DEIM, DINO) usam normalmente valores muito mais baixos, como 0.1, para estabilizar as camadas de cross-attention do decoder, onde a magnitude dos gradientes pode aumentar abruptamente. Para substituir o valor de clipping, cria uma subclasse do trainer e substitui optimizer_step():
import torch
from ultralytics import RTDETR
from ultralytics.models.rtdetr.train import RTDETRTrainer
from ultralytics.utils.torch_utils import TORCH_2_0
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:
kwargs = {"foreach": False} if self.device.type == "npu" and TORCH_2_0 else {}
torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=self.clip_grad_norm, **kwargs)
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)O mesmo trainer funciona para YOLO, alterando a classe-pai para DetectionTrainer (from ultralytics.models.yolo.detect import DetectionTrainer) e carregando um checkpoint YOLO com YOLO("yolo26n.pt"). O corpo de optimizer_step não é alterado.
| Família de arquiteturas | max_norm típico |
|---|---|
| Família RT-DETR / DEIM / DETR | 0.1 |
| YOLO (predefinição da Ultralytics) | 10.0 |
| Desativar o clipping | 0 |
Perguntas frequentes#
Passa a tua classe de trainer personalizada (não uma instância) ao parâmetro
traineremmodel.train():from ultralytics import YOLO from ultralytics.models.yolo.detect import DetectionTrainer class MyCustomTrainer(DetectionTrainer): """A custom trainer that extends DetectionTrainer.""" model = YOLO("yolo26n.pt") model.train(data="coco8.yaml", trainer=MyCustomTrainer)A classe
YOLOtrata internamente da instanciação do trainer. Consulta a página Personalização avançada para obteres mais detalhes sobre a arquitetura do trainer.Principais métodos disponíveis para personalização:
Método Finalidade validate()Executar a validação e devolver métricas build_optimizer()Construir o otimizador save_model()Guardar checkpoints de treino get_model()Devolver a instância do modelo get_validator()Devolver a instância do validador get_dataloader()Construir o dataloader preprocess_batch()Pré-processar o batch de entrada label_loss_items()Formatar itens da loss para registo Para obteres a referência completa da API, consulta a documentação de
BaseTrainer.Sim, para personalizações mais simples, os callbacks são frequentemente suficientes. Os eventos de callback disponíveis incluem
on_train_start,on_train_epoch_start,on_train_epoch_end,on_fit_epoch_endeon_model_save. Estes permitem associar lógica ao ciclo de treino sem criar uma subclasse. O exemplo de congelamento do backbone acima demonstra esta abordagem.Se a tua alteração for mais simples (como ajustar os ganhos da loss), podes modificar diretamente os hiperparâmetros:
from ultralytics import YOLO model = YOLO("yolo26n.pt") model.train(data="coco8.yaml", box=10.0, cls=1.5, dfl=2.0)No YOLO26,
dflescala ol1_lossregistado porque a cabeça de deteção usareg_max: 1; nos modelos comreg_max > 1, escaladfl_loss.