Настройка Тренера
Конвейер обучения Ultralytics построен вокруг BaseTrainer и специализированные тренеры, такие как DetectionTrainer. Эти классы обеспечивают цикл обучения, валидацию, контрольные точки и логирование «из коробки». Если вам нужен больший контроль — отслеживание пользовательских метрик, корректировка весов потерь или реализация графиков скорости обучения — вы можете создать подкласс тренера и переопределить определенные методы.
В этом руководстве рассматриваются пять распространенных настроек:
- Логирование пользовательских метрик (F1-мера) в конце каждой эпохи
- Добавление весов классов для обработки дисбаланса классов
- Сохранение лучшей модели на основе другой метрики
- Заморозка базовой сети на первые N эпох с последующей разморозкой
- Задание скоростей обучения для каждого слоя.
Предварительные требования
Прежде чем читать это руководство, убедитесь, что вы знакомы с основами обучение моделей YOLO и Расширенная настройка странице, которая охватывает BaseTrainer архитектура.
Как работают пользовательские тренеры
Параметр YOLO класс модели принимает trainer в train() метод. Это позволяет передать собственный класс тренера, который расширяет поведение по умолчанию:
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)
Ваш пользовательский тренажер наследует всю функциональность от DetectionTrainer, поэтому вам нужно переопределить только те конкретные методы, которые вы хотите настроить.
Регистрация пользовательских метрик
Параметр валидация шаг вычисляет точность, полнотой (recall)и mAP. Если вам нужны дополнительные метрики, такие как поклассовые F1 score, переопределите 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)
Это регистрирует средний показатель F1 по всем классам и разбивку по классам после каждого запуска валидации.
Доступные метрики
Валидатор предоставляет доступ ко многим метрикам через self.validator.metrics.box:
| Атрибут | Описание |
|---|---|
f1 | F1-мера для каждого класса |
p | Точность по классам |
r | Полнота для каждого класса |
ap50 | AP при IoU 0.5 по классам |
ap | AP при IoU 0.5:0.95 по классам |
mp, mr | Средняя точность и полнота |
map50, map | Средние метрики AP |
Добавление весов классов
Если ваш набор данных содержит несбалансированные классы (например, редкий дефект при производственном контроле), вы можете увеличить вес недопредставленных классов в функции потерь. Это заставляет модель сильнее наказывать за ошибочные классификации редких классов.
Чтобы настроить функцию потерь, создайте подклассы классов потерь, модели и тренера:
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)
Вычисление весов из набора данных
Вы можете автоматически вычислять веса классов на основе распределения меток в вашем наборе данных. Распространенный подход — взвешивание по обратной частоте:
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]
Сохранение лучшей модели по пользовательской метрике
Модуль обучения сохраняет best.pt на основе показателя пригодности, который по умолчанию равен 0.9 × mAP@0.5:0.95 + 0.1 × mAP@0.5. Чтобы использовать другую метрику (например, mAP@0.5 или полноты), переопределить validate() и возвращает выбранную метрику в качестве значения пригодности. Встроенный save_model() будет затем использовать его автоматически:
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)
Доступные метрики
Доступные метрики в self.metrics после валидации включают:
| Ключ | Описание |
|---|---|
metrics/precision(B) | Точность |
metrics/recall(B) | Полнота (Recall) |
metrics/mAP50(B) | mAP при IoU 0.5 |
metrics/mAP50-95(B) | mAP при IoU 0.5:0.95 |
Замораживание и размораживание Backbone
Трансферное обучение рабочие процессы часто выигрывают от замораживания предварительно обученного бэкбона на первые N эпох, позволяя detect-голове адаптироваться до дообучение всю сеть. Ultralytics предоставляет freeze параметр для замораживания слоев в начале обучения, и вы можете использовать callback чтобы разморозить их после N эпох:
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)
Параметр freeze=10 параметр замораживает первые 10 слоев (основу) в начале обучения. Этот on_train_epoch_start Callback срабатывает в начале каждой эпохи и размораживает все параметры после завершения периода заморозки.
Выбор того, что заморозить
freeze=10замораживает первые 10 слоев (обычно это backbone в архитектурах YOLO)freeze=[0, 1, 2, 3]замораживает определённые слои по индексу- Выше
FREEZE_EPOCHSзначения дают головной части больше времени на адаптацию до изменения базовой сети
Послойные скорости обучения
Различные части сети могут выигрывать от разных скоростей обучения. Распространенная стратегия заключается в использовании более низкой скорости обучения для предварительно обученного базового блока для сохранения изученных признаков, в то время как головка обнаружения может адаптироваться быстрее с более высокой скоростью:
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)
Планировщик скорости обучения
Встроенный планировщик скорости обучения (cosine или linear) по-прежнему применяется поверх базовых скоростей обучения для каждой группы. Скорости обучения как для backbone, так и для head будут следовать одному и тому же графику затухания, поддерживая соотношение между ними на протяжении всего обучения.
Комбинирование методов
Эти настройки могут быть объединены в единый класс тренера путем переопределения нескольких методов и добавления колбэков по мере необходимости.
Часто задаваемые вопросы
Как мне передать пользовательский тренер в YOLO?
Передайте ваш пользовательский класс тренера (не экземпляр) в trainer параметре в model.train():
from ultralytics import YOLO
model = YOLO("yolo26n.pt")
model.train(data="coco8.yaml", trainer=MyCustomTrainer)
Параметр YOLO класс обрабатывает создание экземпляра тренера внутри. См. Расширенная настройка странице для получения более подробной информации об архитектуре тренера.
Какие методы BaseTrainer я могу переопределить?
Ключевые методы, доступные для настройки:
| Метод | Цель |
|---|---|
validate() | Выполнить валидацию и вернуть метрики |
build_optimizer() | Сконструировать оптимизатор |
save_model() | Сохранять контрольные точки обучения |
get_model() | Вернуть экземпляр модели |
get_validator() | Вернуть экземпляр валидатора |
get_dataloader() | Построение загрузчика данных |
preprocess_batch() | Предварительная обработка входного пакета |
label_loss_items() | Форматирование элементов потерь для логирования |
Полную справочную информацию по API см. в BaseTrainer документацию.
Можно ли использовать колбэки вместо наследования класса тренера?
Да, для более простых настроек, Колбэки часто достаточны. Доступные события обратного вызова включают on_train_start, on_train_epoch_start, on_train_epoch_end, on_fit_epoch_endи on_model_save. Они позволяют вам подключаться к циклу обучения без создания подклассов. Пример замораживания backbone, приведенный выше, демонстрирует этот подход.
Как мне настроить функцию потерь без создания подкласса модели?
Если ваше изменение проще (например, корректировка коэффициентов потерь), вы можете изменить гиперпараметры напрямую:
model.train(data="coco8.yaml", box=10.0, cls=1.5, dfl=2.0)
Для структурных изменений в функции потерь (например, добавление весов классов) необходимо создать подкласс функции потерь и модели, как показано в разделе весов классов.
