Link to this sectionPersonalizar o Trainer#
O pipeline de treinamento da Ultralytics é construído em torno do BaseTrainer e de treinadores específicos de tarefa como o DetectionTrainer. Essas classes lidam com o loop de treinamento, validação, criação de checkpoints e registro de logs nativamente. Quando precisares de mais controle — como rastrear métricas personalizadas, ajustar a ponderação de perdas ou implementar agendamentos de taxa de aprendizado — podes criar uma subclasse do trainer e substituir métodos específicos.
Este guia percorre sete personalizações comuns:
- Registrar métricas personalizadas (pontuação F1) ao final de cada época
- Adicionar pesos de classe para lidar com o desequilíbrio de classes
- Salvar o melhor modelo com base em uma métrica diferente
- Congelar a backbone pelas primeiras N épocas, depois descongelar
- Especificar taxas de aprendizado por camada
- Sincronizar BatchNorm entre GPUs para treinamento multi-GPU
- Configurar o clipping de gradiente para ajuste de estabilidade
Antes de ler este guia, certifica-te de que estás familiarizado com os conceitos básicos de treinamento de modelos YOLO e a página de Personalização Avançada, que aborda a arquitetura do BaseTrainer.
Link to this sectionComo funcionam os Trainers Personalizados#
A classe de modelo YOLO aceita um parâmetro trainer no método train(). Isso permite-te passar a tua própria classe de trainer que estende o comportamento padrão:
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)O teu trainer personalizado herda todas as funcionalidades do DetectionTrainer, por isso só precisas de substituir os métodos específicos que desejas personalizar.
Link to this sectionRegistrando Métricas Personalizadas#
O passo de validação calcula precisão, recall e mAP. Se precisares de métricas adicionais, como pontuação F1 por classe, substitui o método 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)Isso registra a pontuação F1 média em todas as classes e um detalhamento por classe após cada execução de validação.
O validador fornece acesso a muitas 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 precisão, recall, F1, TP, FP e FN |
p | Precisão por classe |
r | Recall por classe |
ap50 | AP em IoU 0.5 por classe |
ap | AP em IoU 0.5:0.95 por classe |
mp, mr | Precisão e recall médios |
map50, map | Métricas mAP médias |
Link to this sectionAdicionando Pesos de Classe#
Se o teu conjunto de dados tiver classes desequilibradas (por exemplo, um defeito raro na inspeção de fabricação), podes aumentar o peso das classes sub-representadas na função de perda. Isso faz com que o modelo penalize mais fortemente classificações incorretas em classes raras.
Para personalizar a perda, cria subclasses para as classes de perda, o modelo e o 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)Podes calcular os pesos de classe automaticamente a partir da distribuição de rótulos do teu conjunto de dados. Uma abordagem comum é a ponderação por frequência 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]Classes personalizadas como WeightedDetectionModel são armazenadas no checkpoint por referência. Quando definidas em um script de treinamento, elas pertencem ao módulo __main__, portanto, carregar best.pt a partir de um script diferente gera o erro AttributeError: Can't get attribute 'WeightedDetectionModel' on <module '__main__'>.
Define classes personalizadas em um módulo dedicado para que permaneçam importáveis, e garante que esse módulo esteja no teu PYTHONPATH no momento do carregamento.
# weighted_model.py
from ultralytics.nn.tasks import DetectionModel
class WeightedDetectionModel(DetectionModel):
"""Detection model that uses class-weighted loss."""
...# inference script
from weighted_model import WeightedDetectionModel # noqa: F401 - must be importable at checkpoint load time
from ultralytics import YOLO
model = YOLO("runs/detect/train/weights/best.pt")
metrics = model.val()Link to this sectionSalvando o Melhor Modelo por Métrica Personalizada#
O trainer salva o best.pt com base na adequação (fitness), que para detecção assume como padrão o mAP@0.5:0.95 (pesos [0.0, 0.0, 0.0, 1.0] para [P, R, mAP@0.5, mAP@0.5:0.95]). Para usar uma métrica diferente (como mAP@0.5 ou recall), substitui o validate() e retorna a métrica escolhida como o valor de adequação. O save_model() embutido usá-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."""
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 comuns disponíveis em self.metrics após a validação incluem:
| Chave | Descrição |
|---|---|
metrics/precision(B) | Precisão |
metrics/recall(B) | Recall |
metrics/mAP50(B) | mAP em IoU 0.5 |
metrics/mAP50-95(B) | mAP em IoU 0.5:0.95 |
Link to this sectionCongelando e Descongelando a Backbone#
Fluxos de trabalho de aprendizado por transferência frequentemente se beneficiam ao congelar a backbone pré-treinada pelas primeiras N épocas, permitindo que a cabeça de detecção se adapte antes do ajuste fino de toda a rede. A Ultralytics fornece um parâmetro freeze para congelar camadas no início do treinamento, e podes usar um callback para descongelá-las 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 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)O parâmetro freeze=10 congela as primeiras 10 camadas (a backbone) no início do treinamento. O callback on_train_epoch_start é disparado no início de cada época e descongela todos os parâmetros uma vez que o período de congelamento esteja completo.
freeze=10congela as primeiras 10 camadas (tipicamente a backbone nas arquiteturas YOLO)freeze=[0, 1, 2, 3]congela camadas específicas por índice- Valores de
FREEZE_EPOCHSmais altos dão à cabeça mais tempo para se adaptar antes que a backbone mude
Link to this sectionTaxas de Aprendizado por Camada#
Diferentes partes da rede podem beneficiar de diferentes taxas de aprendizado. Uma estratégia comum é usar uma taxa de aprendizado menor para a backbone pré-treinada, para preservar características aprendidas, enquanto permite que a cabeça de detecção se adapte mais rapidamente com uma taxa maior:
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)Link to this sectionVariante RT-DETR#
Para o RT-DETR o padrão é o mesmo com dois refinamentos. O comprimento da backbone é lido de model.yaml["backbone"], portanto o mesmo trainer funciona em variantes do RT-DETR (RT-DETR-L, RT-DETR-X, backbones ResNet-50/101) sem codificar valores de camadas. Os parâmetros também são divididos em grupos de pesos, BatchNorm e vieses dentro de cada seção para que o weight decay seja excluído dos parâmetros de BatchNorm e vieses, correspondendo à política do trainer padrão. Isso é especialmente útil para ajuste fino do RT-DETR, onde a cabeça do decodificador é tipicamente inicializada aleatoriamente, enquanto a backbone carrega características pré-treinadas que se beneficiam de uma taxa de aprendizado menor:
import torch
from torch import nn
from ultralytics import RTDETR
from ultralytics.models.rtdetr.train import RTDETRTrainer
from ultralytics.utils import LOGGER, colorstr
from ultralytics.utils.torch_utils import unwrap_model
class RTDETRBackboneLRTrainer(RTDETRTrainer):
"""RT-DETR trainer with a lower learning rate for backbone parameters."""
backbone_lr_ratio = 0.1 # backbone learning rate as a fraction of head learning rate
def build_optimizer(self, model, name="auto", lr=0.001, momentum=0.9, decay=1e-5, iterations=1e5):
"""Build an AdamW optimizer with six param groups: head and backbone x {weight, bn, bias}."""
# Resolve optimizer name; "auto" maps to AdamW with RT-DETR-style defaults
canonical = {"Adam", "Adamax", "AdamW", "NAdam", "RAdam", "auto"}
name = {x.lower(): x for x in canonical}.get(name.lower(), name)
if name == "auto":
name, lr, momentum = "AdamW", 1e-4, 0.9
self.args.warmup_bias_lr = 0.0 # RT-DETR warms biases from 0, unlike YOLO's 0.1
if name not in {"Adam", "Adamax", "AdamW", "NAdam", "RAdam"}:
raise NotImplementedError(f"This trainer only supports AdamW-family optimizers; got {name}")
# Identify backbone parameters from model.yaml and route each param into a (section, kind) group
unwrapped = unwrap_model(model)
backbone_len = len(unwrapped.yaml["backbone"])
norm_types = tuple(v for k, v in nn.__dict__.items() if "Norm" in k)
groups = {f"{s}_{k}": [] for s in ("head", "backbone") for k in ("weight", "bn", "bias")}
for module_name, module in unwrapped.named_modules():
for param_name, param in module.named_parameters(recurse=False):
if not param.requires_grad:
continue
fullname = f"{module_name}.{param_name}" if module_name else param_name
parts = fullname.split(".")
section = (
"backbone"
if len(parts) > 1 and parts[0] == "model" and parts[1].isdigit() and int(parts[1]) < backbone_len
else "head"
)
if "bias" in param_name:
kind = "bias"
elif isinstance(module, norm_types) or "logit_scale" in fullname:
kind = "bn"
else:
kind = "weight"
groups[f"{section}_{kind}"].append(param)
# Build the optimizer with per-group lr and weight decay; backbone groups use lr * backbone_lr_ratio
backbone_lr = lr * self.backbone_lr_ratio
param_groups = [
{"params": groups["head_weight"], "lr": lr, "weight_decay": decay, "param_group": "weight"},
{"params": groups["head_bn"], "lr": lr, "weight_decay": 0.0, "param_group": "bn"},
{"params": groups["head_bias"], "lr": lr, "weight_decay": 0.0, "param_group": "bias"},
{"params": groups["backbone_weight"], "lr": backbone_lr, "weight_decay": decay, "param_group": "weight"},
{"params": groups["backbone_bn"], "lr": backbone_lr, "weight_decay": 0.0, "param_group": "bn"},
{"params": groups["backbone_bias"], "lr": backbone_lr, "weight_decay": 0.0, "param_group": "bias"},
]
param_groups = [pg for pg in param_groups if pg["params"]] # drop empty groups
optimizer = getattr(torch.optim, name)(param_groups, betas=(momentum, 0.999))
LOGGER.info(
f"{colorstr('optimizer:')} {name}(lr={lr}, backbone_lr={backbone_lr}) with parameter groups\n"
f" Head: {len(groups['head_bn'])} bn, {len(groups['head_weight'])} weight(decay={decay}), "
f"{len(groups['head_bias'])} bias (lr={lr})\n"
f" Backbone: {len(groups['backbone_bn'])} bn, {len(groups['backbone_weight'])} weight(decay={decay}), "
f"{len(groups['backbone_bias'])} bias (lr={backbone_lr})"
)
return optimizer
model = RTDETR("rtdetr-l.pt")
model.train(data="coco8.yaml", epochs=20, trainer=RTDETRBackboneLRTrainer)Um ponto de partida comum é backbone_lr_ratio = 0.1, correspondendo à configuração original do RT-DETR com sua backbone HGNetV2. A literatura sugere escalar a razão inversamente com o tamanho da backbone e a escala dos dados de pré-treinamento: backbones grandes pré-treinadas em conjuntos de dados muito extensos (por exemplo, ViT-L/H treinadas com DINO, CLIP ou MAE em centenas de milhões de imagens) tipicamente usam razões menores como 0.01 ou abaixo para preservar características bem aprendidas, enquanto backbones menores com pré-treinamento mais leve toleram razões maiores como 0.5 ou mais.
O agendador de taxa de aprendizado embutido (cosine ou linear) ainda se aplica sobre as taxas de aprendizado base por grupo. Tanto as taxas de aprendizado da backbone quanto da cabeça seguirão o mesmo cronograma de decaimento, mantendo a proporção entre elas durante todo o treinamento.
Essas personalizações podem ser combinadas em uma única classe de trainer ao substituir múltiplos métodos e adicionar callbacks conforme necessário.
Link to this sectionSincronização de BatchNorm para Treinamento Multi-GPU#
Ao treinar em múltiplas GPUs com DistributedDataParallel, as camadas BatchNorm2d padrão calculam estatísticas independentemente em cada GPU. Para ajuste fino de RT-DETR e outras receitas que usam tamanhos de batch pequenos por GPU, as estatísticas de batch por GPU podem ser ruidosas. O SyncBatchNorm do PyTorch sincroniza a média e a variância em todos os ranks para uma estatística de batch global única, o que frequentemente melhora a convergência ao custo de um pequeno overhead de comunicação inter-GPU.
A conversão tem de acontecer depois que o modelo estiver na GPU, mas antes que o DDP o envolva. O gancho mais limpo para isso é set_model_attributes(), que o BaseTrainer chama exatamente nessa janela:
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 guarda world_size > 1 garante que o trainer seja seguro para uso também em execuções com GPU única; em uma GPU única a conversão é ignorada e o treinamento prossegue com BatchNorm2d normal. O mesmo padrão funciona para o YOLO ao mudar a classe pai para DetectionTrainer.
| Cenário | Recomendação |
|---|---|
| Treinamento multi-GPU, batch pequeno por GPU (≤ 16) | Ativar |
| Treinamento multi-GPU, batch grande por GPU (≥ 32) | Opcional; benefício menor |
| Treinamento em GPU única | Não aplicável (ignorado) |
Link to this sectionConfiguração de Clipping de Gradiente#
O trainer padrão corta gradientes para max_norm=10.0 em optimizer_step(), um valor flexível ajustado para modelos YOLO onde gradientes raramente o excedem. Detectores da família DETR (RT-DETR, DEIM, DINO) tipicamente usam valores muito mais rigorosos como 0.1 para estabilizar as camadas de cross-attention do decodificador, onde as magnitudes de gradiente podem sofrer picos. Para substituir o valor de corte, cria uma subclasse do trainer e substitui optimizer_step():
import torch
from ultralytics import RTDETR
from ultralytics.models.rtdetr.train import RTDETRTrainer
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:
torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=self.clip_grad_norm)
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 o YOLO ao mudar a classe pai para DetectionTrainer (from ultralytics.models.yolo.detect import DetectionTrainer) e carregar um checkpoint YOLO com YOLO("yolo26n.pt"). O corpo do optimizer_step permanece inalterado.
| Família de arquitetura | max_norm típico |
|---|---|
| Família RT-DETR / DEIM / DETR | 0.1 |
| YOLO (padrão da Ultralytics) | 10.0 |
| Desativar clipping | 0 |
Link to this sectionFAQ#
Link to this sectionComo passo um trainer personalizado para o YOLO?#
Passa a tua classe de trainer personalizada (não uma instância) para o parâmetro trainer no model.train():
from ultralytics import YOLO
model = YOLO("yolo26n.pt")
model.train(data="coco8.yaml", trainer=MyCustomTrainer)A classe YOLO lida com a instanciação do trainer internamente. Consulta a página de Personalização Avançada para mais detalhes sobre a arquitetura do trainer.
Link to this sectionQuais métodos do BaseTrainer posso substituir?#
Principais métodos disponíveis para personalização:
| Método | Objetivo |
|---|---|
validate() | Executa a validação e retorna as métricas |
build_optimizer() | Constrói o otimizador |
save_model() | Salva os checkpoints do treinamento |
get_model() | Retorna a instância do modelo |
get_validator() | Retorna a instância do validador |
get_dataloader() | Constrói o dataloader |
preprocess_batch() | Pré-processa o batch de entrada |
label_loss_items() | Formata itens de perda para registro (logging) |
Para a referência completa da API, veja a documentação de BaseTrainer.
Link to this sectionPosso usar callbacks em vez de criar subclasses do treinador?#
Sim, para personalizações mais simples, callbacks costumam ser suficientes. Eventos de callback disponíveis incluem on_train_start, on_train_epoch_start, on_train_epoch_end, on_fit_epoch_end e on_model_save. Eles permitem que você se conecte ao loop de treinamento sem precisar criar subclasses. O exemplo de congelamento de backbone acima demonstra essa abordagem.
Link to this sectionComo personalizo a função de perda sem criar uma subclasse do modelo?#
Se a sua alteração for mais simples (como ajustar ganhos de perda), você pode modificar os hiperparâmetros diretamente:
model.train(data="coco8.yaml", box=10.0, cls=1.5, dfl=2.0)Para alterações estruturais na perda (como adicionar pesos de classe), você precisa criar subclasses da perda e do modelo, conforme mostrado na seção de pesos de classe.