Trainer 사용자 지정#
Ultralytics 학습 파이프라인은 BaseTrainer과 DetectionTrainer과 같은 작업별 trainer를 중심으로 구축되어 있습니다. 이러한 클래스는 학습 루프, 검증, 체크포인트 저장 및 로깅을 기본적으로 처리합니다. 사용자 지정 지표 추적, loss 가중치 조정 또는 learning rate schedule 구현과 같이 더 세밀한 제어가 필요한 경우 trainer를 서브클래싱하고 특정 메서드를 재정의할 수 있습니다.
이 가이드에서는 다음과 같은 일반적인 사용자 지정 방법 7가지를 설명합니다:
- 각 epoch 종료 시 사용자 지정 지표(F1 점수) 로깅
- 클래스 불균형을 처리하기 위한 클래스 가중치 추가
- 다른 지표를 기준으로 최적 모델 저장
- 처음 N개 에포크 동안 백본 고정 후 고정 해제
- 레이어별 학습률 지정
- 다중 GPU 학습을 위한 GPU 간 BatchNorm 동기화
- 안정성 조정을 위한 그래디언트 클리핑 구성
이 가이드를 읽기 전에 YOLO 모델 학습의 기본 사항과 BaseTrainer 아키텍처를 다루는 고급 사용자 지정 페이지를 숙지하시기 바랍니다.
사용자 지정 트레이너의 작동 방식#
YOLO 모델 클래스는 train() 메서드에서 trainer 매개변수를 받습니다. 이를 통해 기본 동작을 확장하는 자체 트레이너 클래스를 전달할 수 있습니다:
from ultralytics import YOLO
from ultralytics.models.yolo.detect import DetectionTrainer
class CustomTrainer(DetectionTrainer):
"""A custom trainer that extends DetectionTrainer with additional functionality."""
# 여기에 사용자 지정 내용을 추가합니다
model = YOLO("yolo26n.pt")
model.train(data="coco8.yaml", epochs=10, trainer=CustomTrainer)사용자 지정 트레이너는 DetectionTrainer의 모든 기능을 상속하므로, 사용자 지정하려는 특정 메서드만 재정의하면 됩니다.
사용자 지정 지표 로깅#
검증 단계에서는 정밀도, 재현율, mAP을 계산합니다. 클래스별 F1 점수와 같은 추가 지표가 필요하면 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
mean_f1 = float(np.mean(f1_per_class)) if len(f1_per_class) 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)]
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 점수 |
image_metrics | 정밀도, 재현율, F1, TP, FP, FN이 포함된 이미지별 지표 딕셔너리 |
p | 클래스별 정밀도 |
r | 클래스별 재현율 |
ap50 | 클래스별 IoU 0.5 기준 AP |
ap | 클래스별 IoU 0.5:0.95 기준 AP |
mp, mr | 평균 정밀도 및 재현율 |
map50, map | 평균 AP 지표 |
클래스 가중치 추가#
분류 손실에 정규화된 역빈도 가중치를 적용하려면 cls_pw을 0.0과 1.0 사이의 값으로 설정합니다. 직접 지정한 비율이 필요한 경우에만 기존 가중치 계산을 재정의합니다:
import numpy as np
from ultralytics import YOLO
from ultralytics.models.yolo.detect import DetectionTrainer
class WeightedTrainer(DetectionTrainer):
"""Detection trainer with hand-picked class-weight ratios."""
def compute_class_weights(self, class_counts):
"""Return custom per-class weights that the detection loss applies."""
weights = np.ones_like(class_counts)
weights[0] = 2.0
weights[1] = 3.0
return weights
model = YOLO("yolo26n.pt")
model.train(data="custom.yaml", epochs=10, cls_pw=1.0, trainer=WeightedTrainer)set_class_weights()은 이 값을 평균 1.0으로 정규화하고 모델에 저장하며, 기존 탐지 손실에서 해당 값을 적용합니다. 위 인덱스를 사용하려면 데이터 세트에 클래스가 두 개 이상 있어야 합니다.
사용자 지정 지표를 기준으로 최적 모델 저장#
트레이너는 적합도를 기준으로 best.pt을 저장합니다. 탐지 작업에서 적합도는 기본적으로 mAP@0.5:0.95이며, [P, R, mAP@0.5, mAP@0.5:0.95]에 [0.0, 0.0, 0.0, 1.0] 가중치를 적용합니다. 다른 지표(예: 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."""
previous_best = self.best_fitness
metrics, fitness = super().validate()
if metrics is None:
return metrics, fitness
fitness = metrics["metrics/mAP50(B)"]
self.best_fitness = fitness if previous_best is None else max(previous_best, fitness)
return metrics, fitness
model = YOLO("yolo26n.pt")
model.train(data="coco8.yaml", epochs=20, trainer=CustomSaveTrainer)BaseTrainer.validate()은 기본 지표를 사용해 best_fitness을 업데이트하므로, 메서드를 호출하기 전에 이전 값을 저장합니다.
검증 후 self.metrics에서 사용할 수 있는 일반적인 지표는 다음과 같습니다:
| 키 | 설명 |
|---|---|
metrics/precision(B) | 정밀도 |
metrics/recall(B) | 재현율 |
metrics/mAP50(B) | IoU 0.5 기준 mAP |
metrics/mAP50-95(B) | IoU 0.5:0.95 기준 mAP |
백본 고정 및 고정 해제#
전이 학습 워크플로에서는 처음 N개 에포크 동안 사전 학습된 백본을 고정하여 탐지 헤드가 적응하도록 한 다음, 전체 네트워크를 파인튜닝하는 방식이 효과적인 경우가 많습니다. Ultralytics는 학습 시작 시 레이어를 고정하는 freeze 매개변수를 제공하며, 콜백을 사용해 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 the user-requested layers after FREEZE_EPOCHS."""
if trainer.epoch == FREEZE_EPOCHS:
user_freeze = [x for x in trainer.freeze_layer_names if x not in {".dfl", "teacher_model."}]
LOGGER.info(f"Epoch {trainer.epoch}: Unfreezing requested layers for fine-tuning")
for name, param in trainer.model.named_parameters():
if (
not param.requires_grad
and ".dfl" not in name
and "teacher_model." not in name
and any(x in name for x in user_freeze)
):
param.requires_grad = True
LOGGER.info(f" Unfroze: {name}")
trainer.freeze_layer_names = [x for x in trainer.freeze_layer_names if x not in user_freeze]
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개 레이어(인덱스 0~9)를 고정하며, 이 범위에는 YOLO26 백본의 대부분이 포함됩니다. 백본은 레이어 0~10으로 구성되므로 freeze=10은 마지막 C2PSA 블록(레이어 10)을 학습 가능한 상태로 둡니다. 백본 전체를 고정하려면 freeze=11를 사용합니다. on_train_epoch_start 콜백은 각 에포크 시작 시 실행되며, 고정 기간이 끝나면 요청된 레이어의 고정을 해제하는 동시에 영구 고정된 DFL 및 증류 교사 매개변수는 그대로 유지합니다.
freeze=10은 처음 10개 레이어(인덱스 0~9, YOLO26 백본의 대부분)를 고정합니다. 마지막 C2PSA 블록(레이어 10)까지 포함하려면freeze=11을 사용합니다.freeze=[0, 1, 2, 3]은 인덱스를 기준으로 특정 레이어를 고정합니다.FREEZE_EPOCHS값이 클수록 백본이 변경되기 전에 헤드가 적응할 시간이 늘어납니다.
레이어별 학습률#
네트워크의 각 부분에 서로 다른 학습률을 적용하면 효과적일 수 있습니다. 일반적인 전략은 사전 학습된 백본에 더 낮은 학습률을 적용해 학습된 특징을 보존하고, 탐지 헤드에는 더 높은 학습률을 적용해 빠르게 적응시키는 것입니다:
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."""
backbone_lr_ratio = 0.1
def build_optimizer(self, model, name="auto", lr=0.001, momentum=0.9, decay=1e-5, iterations=1e5):
"""Reuse the trainer optimizer and lower its backbone parameter-group rates."""
optimizer = super().build_optimizer(model, name, lr, momentum, decay, iterations)
unwrapped = unwrap_model(model)
backbone_len = len(unwrapped.yaml["backbone"])
backbone = {
id(p)
for name, p in unwrapped.named_parameters()
if any(name.startswith(f"model.{i}.") for i in range(backbone_len))
}
groups = []
for group in optimizer.param_groups:
head_params = [p for p in group["params"] if id(p) not in backbone]
backbone_params = [p for p in group["params"] if id(p) in backbone]
if head_params:
groups.append({**group, "params": head_params})
if backbone_params:
groups.append({**group, "params": backbone_params, "lr": group["lr"] * self.backbone_lr_ratio})
optimizer.param_groups = groups
LOGGER.info(f"PerLayerLR: {len(backbone)} backbone params at {self.backbone_lr_ratio}x the head rate")
return optimizer
model = YOLO("yolo26n.pt")
model.train(data="coco8.yaml", epochs=20, trainer=PerLayerLRTrainer)RT-DETR 변형#
RT-DETR에서는 RTDETRTrainer을 부모 클래스로 지정하고 RTDETR("rtdetr-l.pt")로 체크포인트를 불러오는 동일한 재정의 방법을 사용합니다.
다중 GPU 학습을 위한 동기화 BatchNorm#
DistributedDataParallel을 사용하여 여러 GPU에서 학습할 때 기본 BatchNorm2d 레이어는 각 GPU에서 독립적으로 통계를 계산합니다. RT-DETR 파인튜닝 및 GPU당 배치 크기가 작은 다른 레시피에서는 GPU별 배치 통계가 불안정해질 수 있습니다. PyTorch의 SyncBatchNorm은 모든 랭크의 평균과 분산을 동기화하여 단일 전역 배치 통계를 계산합니다. 이 방식은 GPU 간 통신 오버헤드가 조금 발생하는 대신 수렴 성능을 개선하는 경우가 많습니다.
변환은 모델을 GPU에 올린 후, DDP로 래핑하기 전에 수행해야 합니다. 이를 위한 가장 적절한 훅은 set_model_attributes()이며, BaseTrainer은 정확히 해당 시점에 이 훅을 호출합니다:
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)world_size > 1 가드를 사용하면 단일 GPU 실행에서도 트레이너를 안전하게 사용할 수 있습니다. 단일 GPU에서는 변환을 건너뛰고 일반 BatchNorm2d로 학습을 진행합니다. 부모 클래스를 DetectionTrainer로 변경하면 YOLO에도 동일한 패턴을 적용할 수 있습니다.
| 시나리오 | 권장 사항 |
|---|---|
| 다중 GPU 학습, GPU당 배치 크기가 작음(≤ 16) | 활성화 |
| 다중 GPU 학습, GPU당 배치 크기가 큼(≥ 32) | 선택 사항, 효과는 미미함 |
| 단일 GPU 학습 | 해당 없음(건너뜀) |
구성 가능한 그래디언트 클리핑#
기본 트레이너는 optimizer_step()에서 그래디언트를 max_norm=10.0으로 클리핑합니다. 이는 그래디언트가 해당 값을 초과하는 경우가 드문 YOLO 모델에 맞춰 설정된 여유 있는 값입니다. DETR 계열 탐지기(RT-DETR, DEIM, DINO)는 일반적으로 0.1와 같이 훨씬 낮은 값을 사용해 그래디언트 크기가 급증할 수 있는 디코더의 크로스 어텐션 레이어를 안정화합니다. 클리핑 값을 재정의하려면 트레이너를 서브클래싱하고 optimizer_step()을 재정의합니다:
import torch
from ultralytics import RTDETR
from ultralytics.models.rtdetr.train import RTDETRTrainer
from ultralytics.utils.torch_utils import TORCH_2_0
class CustomClipTrainer(RTDETRTrainer):
"""RT-DETR trainer with configurable gradient clipping."""
clip_grad_norm = 0.1 # 최대 그래디언트 노름입니다. 클리핑을 비활성화하려면 0으로 설정합니다.
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:
kwargs = {"foreach": False} if self.device.type == "npu" and TORCH_2_0 else {}
torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=self.clip_grad_norm, **kwargs)
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)부모 클래스를 DetectionTrainer(from ultralytics.models.yolo.detect import DetectionTrainer)으로 변경하고 YOLO("yolo26n.pt")로 YOLO 체크포인트를 불러오면 동일한 트레이너를 YOLO에도 사용할 수 있습니다. optimizer_step 본문은 변경되지 않습니다.
| 아키텍처 계열 | 일반적인 max_norm |
|---|---|
| RT-DETR / DEIM / DETR 계열 | 0.1 |
| YOLO(Ultralytics 기본값) | 10.0 |
| 클리핑 비활성화 | 0 |
자주 묻는 질문#
사용자 지정 트레이너 클래스(인스턴스가 아님)를
model.train()의trainer매개변수에 전달합니다:from ultralytics import YOLO from ultralytics.models.yolo.detect import DetectionTrainer class MyCustomTrainer(DetectionTrainer): """A custom trainer that extends DetectionTrainer.""" model = YOLO("yolo26n.pt") model.train(data="coco8.yaml", trainer=MyCustomTrainer)YOLO클래스는 트레이너 인스턴스화를 내부적으로 처리합니다. 트레이너 아키텍처에 대한 자세한 내용은 고급 사용자 지정 페이지를 참조하세요.사용자 지정에 사용할 수 있는 주요 메서드:
메서드 용도 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가 있습니다. 이 이벤트를 사용하면 서브클래싱하지 않고도 학습 루프에 훅을 연결할 수 있습니다. 위의 백본 고정 예제에서 이 방법을 확인할 수 있습니다.손실 게인 조정과 같이 변경 사항이 간단하다면 하이퍼파라미터를 직접 수정할 수 있습니다:
from ultralytics import YOLO model = YOLO("yolo26n.pt") model.train(data="coco8.yaml", box=10.0, cls=1.5, dfl=2.0)YOLO26에서는 탐지 헤드가
reg_max: 1를 사용하므로dfl이 기록된l1_loss을 스케일링합니다.reg_max > 1을 사용하는 모델에서는dfl_loss를 스케일링합니다.