Personalizando o Treinador

Q: Which BaseTrainer methods can I override?

Principais métodos disponíveis para personalização: Para a referência completa da API, consulte a documentação do BaseTrainer.

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

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_end e on_model_save. Estes permitem-lhe intervir no ciclo de treino sem necessidade de subclassificação. O exemplo de congelamento do backbone acima demonstra esta abordagem.

O pipeline de treinamento Ultralytics é construído em torno de BaseTrainer e treinadores específicos para tarefas como DetectionTrainer. Essas classes gerenciam o loop de treinamento, validação, checkpointing e registro de forma nativa. Quando você precisa de mais controle — como rastrear métricas personalizadas, ajustar o peso da perda ou implementar agendamentos de taxa de aprendizado — você pode criar uma subclasse do treinador e sobrescrever métodos específicos.

Este guia apresenta cinco personalizações comuns:

Registro de métricas personalizadas (pontuação F1) ao final de cada época
Adicionando pesos de classe para lidar com o desequilíbrio de classes
Salvando o melhor modelo com base em uma métrica diferente
Congelamento do backbone para as primeiras N épocas, e depois descongelamento
Especificando taxas de aprendizado por camada

Pré-requisitos

Antes de ler este guia, certifique-se de estar familiarizado com os fundamentos de treinamento de modelos YOLO e Customização Avançada página, que abrange o(a) BaseTrainer arquitetura.

Como Funcionam os Treinadores Personalizados

O YOLO a classe do modelo aceita um trainer parâmetro no train() método. Isto permite-lhe passar a sua própria classe de treinador 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)

Seu treinador personalizado herda todas as funcionalidades de DetectionTrainer, então você só precisa sobrescrever os métodos específicos que deseja personalizar.

Logging de Métricas Personalizadas

O validação o passo calcula precisão, recall, e mAP. Se precisar de métricas adicionais como por classe score F1, sobrescrever 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.

Métricas Disponíveis

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
`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 revocação médias
`map50`, `map`	Métricas de AP médio

Adicionando Pesos de Classe

Se seu conjunto de dados tiver classes desbalanceadas (por exemplo, um defeito raro na inspeção de fabricação), você pode aumentar o peso das classes sub-representadas na função de perda. Isso faz com que o modelo penalize mais severamente as classificações incorretas em classes raras.

Para personalizar a função de perda, crie subclasses das classes de perda, do modelo e do treinador:

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)

Calculando Pesos a partir do Conjunto de Dados

Você pode calcular os pesos das classes automaticamente a partir da distribuição de rótulos do seu 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]

Salvando o Melhor Modelo por Métrica Personalizada

O treinador salva best.pt com base no fitness, que por padrão é 0.9 × mAP@0.5:0.95 + 0.1 × mAP@0.5. Para usar uma métrica diferente (como mAP@0.5 ou recall), substitua validate() e retorne sua métrica escolhida como o valor de aptidão. O integrado save_model() irá então usá-lo 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 Disponíveis

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

Congelamento e Descongelamento do Backbone

Aprendizado por transferência fluxos de trabalho frequentemente se beneficiam de congelar o backbone pré-treinado pelas primeiras N épocas, permitindo que o cabeçalho de detecção se adapte antes ajuste fino toda a rede. A Ultralytics oferece um freeze parâmetro para congelar camadas no início do treinamento, e você pode usar um callback para descongelá-los 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 freeze=10 parâmetro congela as primeiras 10 camadas (o backbone) no início do treinamento. O on_train_epoch_start o callback é acionado no início de cada época e descongela todos os parâmetros assim que o período de congelamento é concluído.

Escolhendo o Que Congelar

freeze=10 congela as 10 primeiras camadas (tipicamente o backbone em arquiteturas YOLO)
freeze=[0, 1, 2, 3] congela camadas específicas por índice
Maior FREEZE_EPOCHS os valores dão à cabeça mais tempo para se adaptar antes que o backbone mude

Taxas de Aprendizagem por Camada

Diferentes partes da rede podem se beneficiar de diferentes taxas de aprendizado. Uma estratégia comum é usar uma taxa de aprendizado mais baixa para o backbone pré-treinado para preservar os recursos aprendidos, enquanto permite que a cabeça de detecção se adapte mais rapidamente com uma taxa mais 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)

Agendador de Taxa de Aprendizado

O agendador de taxa de aprendizagem integrado (cosine ou linear) ainda se aplica sobre as taxas de aprendizado base por grupo. As taxas de aprendizado do backbone e do head seguirão o mesmo cronograma de decaimento, mantendo a proporção entre elas durante todo o treinamento.

Combinação de Técnicas

Essas personalizações podem ser combinadas em uma única classe de treinador, substituindo múltiplos métodos e adicionando callbacks conforme necessário.

FAQ

Como passo um treinador personalizado para o YOLO?

Passe sua classe de treinador personalizada (não uma instância) para o trainer em model.train():

from ultralytics import YOLO

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

O YOLO a classe lida com a instanciação do treinador internamente. Consulte o Customização Avançada página para mais detalhes sobre a arquitetura do treinador.

Quais métodos do BaseTrainer posso sobrescrever?

Principais métodos disponíveis para personalização:

Método	Propósito
`validate()`	Executar validação e retornar métricas
`build_optimizer()`	Construa o otimizador
`save_model()`	Salvar checkpoints de treinamento
`get_model()`	Retornar a instância do modelo
`get_validator()`	Retornar a instância do validador
`get_dataloader()`	Construir o dataloader
`preprocess_batch()`	Pré-processar lote de entrada
`label_loss_items()`	Formatar itens de perda para registo

Para a referência completa da API, consulte BaseTrainer documentação.

Posso usar callbacks em vez de subclassificar o treinador?

Sim, para personalizações mais simples, 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_end, e on_model_save. Isso permite que você se conecte ao loop de treinamento sem a necessidade de subclasse. O exemplo de congelamento do backbone acima demonstra essa abordagem.

Como personalizo a função de perda sem subclassificar o modelo?

Se sua alteração for mais simples (como ajustar os 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 função de perda (como adicionar pesos de classe), é necessário criar uma subclasse da função de perda e do modelo, conforme mostrado na seção de pesos de classe.

📅 Criado há 1 mês ✏️ Atualizado há 1 mês