Personalização do treinador

Q: Which BaseTrainer methods can I override?

Principais métodos disponíveis para personalização: Para obter 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 costumam ser 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. Eles permitem que você se conecte ao loop de treinamento sem subclassificação. O exemplo de congelamento de backbone acima demonstra essa abordagem.

O pipeline Ultralytics é construído em torno de BaseTrainer e formadores especializados em tarefas específicas, como DetectionTrainer. Essas classes lidam com o ciclo de treino, validação, pontos de verificação e registo de saída prontos para uso. Quando precisar de mais controlo — rastrear métricas personalizadas, ajustar a ponderação de perdas ou implementar cronogramas de taxa de aprendizagem — pode criar uma subclasse do treinador e substituir métodos específicos.

Este guia apresenta cinco 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 elas
Guardar o melhor modelo com base numa métrica diferente
Congelar a espinha dorsal durante os primeiros N períodos e, em seguida, descongelar
Especificando taxas de aprendizagem por camada

Pré-requisitos

Antes de ler este guia, certifique-se de que está familiarizado com os conceitos básicos de treinamento de YOLO e Customização Avançada página, que abrange o BaseTrainer arquitetura.

Como funcionam os treinadores personalizados

O YOLO A classe modelo aceita um trainer parâmetro no train() método. Isso permite que você passe 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)

O seu treinador personalizado herda todas as funcionalidades de DetectionTrainer, então só precisa substituir os métodos específicos que deseja personalizar.

Logging de Métricas Personalizadas

O validação passo calcula precisão, recall, e mAP. Se precisar de métricas adicionais, como por classe score F1, substituir 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 regista a pontuação F1 média em todas as classes e uma análise 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 turma
`ap50`	AP IoU ,5 por classe
`ap`	AP IoU ,5:0,95 por aula
`mp`, `mr`	Precisão e recuperação médias
`map50`, `map`	AP médias

Adicionar pesos de classe

Se o seu conjunto de dados tiver classes desequilibradas (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 erradas em classes raras.

Para personalizar a perda, crie subclasses das classes de perda, modelo e 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)

Cálculo de 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]

Guardar o melhor modelo por métrica personalizada

O treinador salva best.pt com base na aptidão física, cujo valor 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), substituir validate() e retorne a métrica escolhida como o valor de adequação. O save_model() irá então utilizá-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 IoU ,5
`metrics/mAP50-95(B)`	mAP IoU ,5:0,95

Congelar e descongelar a espinha dorsal

Aprendizagem por transferência Os fluxos de trabalho muitas vezes se beneficiam do congelamento da estrutura pré-treinada durante os primeiros N épocas, permitindo que o cabeçote de detecção se adapte antes de ajuste fino toda a rede. Ultralytics um freeze parâmetro para congelar camadas no início do treino, e 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 O parâmetro congela as primeiras 10 camadas (a espinha dorsal) no início do treino. 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 primeiras 10 camadas (normalmente a espinha dorsal nas 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 a coluna vertebral mude

Taxas de aprendizagem por camada

Diferentes partes da rede podem beneficiar de diferentes taxas de aprendizagem. Uma estratégia comum é usar uma taxa de aprendizagem mais baixa para a estrutura pré-treinada, a fim de preservar as características aprendidas, permitindo que o cabeçote 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 Aprendizagem

O programador de taxa de aprendizagem integrado (cosine ou linear) continua a aplicar-se além das taxas de aprendizagem básicas por grupo. Tanto as taxas de aprendizagem da espinha dorsal como as da cabeça seguirão o mesmo calendário de decaimento, mantendo a proporção entre elas ao longo do treino.

Combinação de técnicas

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

FAQ

Como passo um treinador personalizado para YOLO?

Passe a 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 obter mais detalhes sobre a arquitetura do treinador.

Quais métodos do BaseTrainer posso substituir?

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()`	Guardar pontos de verificação do treino
`get_model()`	Retornar a instância do modelo
`get_validator()`	Retornar a instância do validador
`get_dataloader()`	Crie o carregador de dados
`preprocess_batch()`	Pré-processar lote de entrada
`label_loss_items()`	Formatar itens perdidos para registo

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

Posso usar callbacks em vez de subclassificar o treinador?

Sim, para personalizações mais simples, retornos de chamada são frequentemente suficientes. Os eventos de retorno de chamada 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 ciclo de treinamento sem criar subclasses. O exemplo de congelamento da espinha dorsal acima demonstra essa abordagem.

Como 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), 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), é necessário criar uma subclasse da perda e do modelo, conforme mostrado na secção de pesos de classe.

📅 Criado há 6 dias ✏️ Atualizado há 0 dias