Настройка тренера
Конвейер 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 ,5 на класс |
ap | AP IoU ,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 ,5 |
metrics/mAP50-95(B) | mAP IoU ,5:0,95 |
Замораживание и размораживание магистрали
Перенос обучения рабочие процессы часто выигрывают от замораживания предварительно обученной основы на первые N эпох, что позволяет детектирующей головке адаптироваться перед точная настройка всю сеть. 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 обратный вызов срабатывает в начале каждой эпохи и размораживает все параметры по окончании периода замораживания.
Выбор того, что заморозить
freeze=10замораживает первые 10 слоев (обычно это основная часть в 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) по-прежнему применяется в дополнение к базовым скоростям обучения для каждой группы. Скорости обучения основной части и головы будут следовать одному и тому же графику снижения, сохраняя соотношение между ними на протяжении всего обучения.
Комбинирование техник
Эти настройки можно объединить в один класс тренера, переопределив несколько методов и добавив обратные вызовы по мере необходимости.
Часто задаваемые вопросы
Как передать пользовательский тренер в 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. Они позволяют подключиться к циклу обучения без создания подклассов. Пример замораживания основного алгоритма, приведенный выше, демонстрирует этот подход.
Как настроить функцию потерь без создания подкласса модели?
Если ваше изменение более простое (например, корректировка потерь и прибыли), вы можете изменить гиперпараметры напрямую:
model.train(data="coco8.yaml", box=10.0, cls=1.5, dfl=2.0)
Для структурных изменений потери (таких как добавление весов классов) необходимо создать подкласс потери и модели, как показано в разделе «Веса классов».
