Trainer 사용자 지정#
Ultralytics 학습 파이프라인은 BaseTrainer 및 DetectionTrainer과 같은 작업별 trainer를 중심으로 구축됩니다. 이러한 클래스는 기본적으로 학습 루프, 검증, 체크포인트 저장 및 로깅을 처리합니다. 사용자 지정 metric 추적, loss 가중치 조정 또는 learning rate schedule 구현 등 더 세밀한 제어가 필요한 경우 trainer를 subclassing하고 특정 메서드를 재정의할 수 있습니다.
이 가이드에서는 다음과 같은 일반적인 사용자 지정 7가지를 설명합니다:
- 각 epoch 종료 시 사용자 지정 metric(F1 score) 로깅
- 클래스 불균형을 처리하기 위한 class weight 추가
- 다른 metric을 기준으로 최고 성능 모델 저장
- 처음 N개 epoch 동안 backbone 고정 후 고정 해제
- 레이어별 learning rate 지정
- 다중 GPU 학습을 위한 GPU 간 BatchNorm 동기화
- 안정성 조정을 위한 gradient clipping 구성
이 가이드를 읽기 전에 YOLO 모델 학습의 기본 사항과 BaseTrainer 아키텍처를 다루는 고급 사용자 지정 페이지를 숙지하시기 바랍니다.
Custom Trainer의 작동 방식#
YOLO 모델 클래스는 train() 메서드에서 trainer parameter를 허용합니다. 이를 통해 기본 동작을 확장한 자체 trainer 클래스를 전달할 수 있습니다:
from ultralytics import YOLO
from ultralytics.models.yolo.detect import DetectionTrainer
class CustomTrainer(DetectionTrainer):
"""A custom trainer that extends DetectionTrainer with additional functionality."""
# Add your customizations here
model = YOLO("yolo26n.pt")
model.train(data="coco8.yaml", epochs=10, trainer=CustomTrainer)사용자 지정 trainer는 DetectionTrainer의 모든 기능을 상속하므로, 사용자 지정하려는 특정 메서드만 재정의하면 됩니다.
사용자 지정 Metric 로깅#
validation 단계에서는 precision, recall 및 mAP를 계산합니다. 클래스별 F1 score와 같은 추가 metric이 필요한 경우 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)이렇게 하면 각 validation 실행 후 validation에 포함된 모든 클래스의 평균 F1 score와 클래스별 세부 결과가 로깅됩니다.
validator는 self.validator.metrics.box을 통해 다양한 metric에 액세스할 수 있도록 제공합니다:
| Attribute | 설명 |
|---|---|
f1 | 클래스별 F1 score |
image_metrics | precision, recall, F1, TP, FP 및 FN을 포함하는 이미지별 metric dictionary |
p | 클래스별 precision |
r | 클래스별 recall |
ap50 | 클래스별 IoU 0.5에서의 AP |
ap | 클래스별 IoU 0.5:0.95에서의 AP |
mp, mr | 평균 precision 및 recall |
map50, map | 평균 AP metric |
Class Weight 추가#
분류 loss에 정규화된 역빈도 가중치를 적용하려면 0.0과 1.0 사이에 cls_pw을 설정합니다. 직접 지정한 비율이 필요한 경우에만 기존 가중치 계산을 재정의합니다:
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 for the production loss owner."""
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으로 정규화하여 모델에 저장하고, 기존 detection loss가 이를 적용하도록 합니다. 위의 인덱스를 사용하려면 최소 두 개의 클래스가 있는 dataset이 필요합니다.
사용자 지정 Metric을 기준으로 최고 성능 모델 저장#
trainer는 fitness를 기준으로 best.pt을 저장하며, detection의 경우 기본값은 mAP@0.5:0.95입니다(가중치 [0.0, 0.0, 0.0, 1.0]는 [P, R, mAP@0.5, mAP@0.5:0.95]에 적용됩니다). 다른 metric(예: mAP@0.5 또는 recall)을 사용하려면 validate()를 재정의하고 선택한 metric을 fitness 값으로 반환합니다. 그러면 내장된 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()은 기본 metric을 사용하여 best_fitness을 업데이트하므로, 이를 호출하기 전에 이전 값을 저장해야 합니다.
validation 후 self.metrics에서 사용할 수 있는 일반적인 metric은 다음과 같습니다:
| 키 | 설명 |
|---|---|
metrics/precision(B) | Precision |
metrics/recall(B) | Recall |
metrics/mAP50(B) | IoU 0.5에서의 mAP |
metrics/mAP50-95(B) | IoU 0.5:0.95에서의 mAP |
Backbone 고정 및 고정 해제#
Transfer learning workflow에서는 처음 N개 epoch 동안 pretrained backbone을 고정하는 것이 효과적인 경우가 많습니다. 이렇게 하면 전체 네트워크를 fine-tuning하기 전에 detection head가 적응할 수 있습니다. Ultralytics는 학습 시작 시 레이어를 고정하는 freeze parameter를 제공하며, callback을 사용하여 N개 epoch 후 고정을 해제할 수 있습니다:
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 parameter는 학습 시작 시 처음 10개 레이어(인덱스 09)를 고정하며, 이는 YOLO26 backbone의 대부분을 포함합니다. backbone은 레이어 010에 걸쳐 있으므로 freeze=10은 마지막 C2PSA block(레이어 10)을 학습 가능한 상태로 둡니다. 전체 backbone을 고정하려면 freeze=11를 사용합니다. on_train_epoch_start callback은 각 epoch 시작 시 실행되며 고정 기간이 끝나면 지정된 레이어의 고정을 해제합니다. 이때 영구적으로 고정된 DFL 및 distillation-teacher parameter는 그대로 유지됩니다.
freeze=10은 처음 10개 레이어(인덱스 0~9)를 고정합니다(YOLO26 backbone의 대부분이며, 레이어 10의 마지막 C2PSA block까지 포함하려면freeze=11사용).freeze=[0, 1, 2, 3]은 인덱스로 특정 레이어를 고정합니다.FREEZE_EPOCHS값이 클수록 backbone이 변경되기 전에 head가 적응할 시간이 늘어납니다.
레이어별 Learning Rate#
네트워크의 서로 다른 부분에는 서로 다른 learning rate가 효과적일 수 있습니다. 일반적인 전략은 pretrained backbone에 더 낮은 learning rate를 사용하여 학습된 feature를 보존하고, detection head에는 더 높은 rate를 사용하여 더 빠르게 적응하도록 하는 것입니다:
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을 parent로 사용하는 동일한 override를 적용하고 RTDETR("rtdetr-l.pt")로 checkpoint를 로드합니다.
다중 GPU 학습을 위한 동기화된 BatchNorm#
DistributedDataParallel을 사용하여 여러 GPU에서 학습할 때 기본 BatchNorm2d 레이어는 각 GPU에서 독립적으로 통계를 계산합니다. RT-DETR fine-tuning 및 GPU별 batch size가 작은 다른 recipe에서는 GPU별 batch 통계가 불안정할 수 있습니다. PyTorch의 SyncBatchNorm은 모든 rank에서 평균과 분산을 동기화하여 단일 전역 batch 통계를 생성합니다. 이 방식은 GPU 간 통신 overhead가 약간 발생하는 대신 convergence를 개선하는 경우가 많습니다.
변환은 모델을 GPU에 올린 후, DDP가 모델을 wrapping하기 전에 수행해야 합니다. 이를 위한 가장 깔끔한 hook은 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 guard는 single-GPU 실행에서도 trainer를 안전하게 사용할 수 있도록 합니다. 단일 GPU에서는 변환을 건너뛰고 일반 BatchNorm2d로 학습을 진행합니다. parent class를 DetectionTrainer로 변경하면 YOLO에도 동일한 패턴을 적용할 수 있습니다.
| 시나리오 | 권장 사항 |
|---|---|
| 다중 GPU 학습, GPU별 작은 batch(≤ 16) | 활성화 |
| 다중 GPU 학습, GPU별 큰 batch(≥ 32) | 선택 사항; 효과 미미 |
| 단일 GPU 학습 | 해당 없음(건너뜀) |
구성 가능한 Gradient Clipping#
기본 trainer는 optimizer_step()에서 gradient를 max_norm=10.0으로 clipping합니다. 이는 gradient가 이 값을 초과하는 경우가 드문 YOLO 모델에 맞춰 설정된 여유 있는 값입니다. DETR 계열 detector(RT-DETR, DEIM, DINO)는 일반적으로 0.1와 같이 훨씬 더 엄격한 값을 사용하여 gradient 크기가 급증할 수 있는 decoder의 cross-attention 레이어를 안정화합니다. clip 값을 재정의하려면 trainer를 subclassing하고 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 # max gradient norm; set to 0 to disable clipping
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)parent class를 DetectionTrainer(from ultralytics.models.yolo.detect import DetectionTrainer)로 변경하고 YOLO("yolo26n.pt")로 YOLO checkpoint를 로드하면 동일한 trainer를 YOLO에 사용할 수 있습니다. optimizer_step 본문은 변경되지 않습니다.
| 아키텍처 계열 | 일반적인 max_norm |
|---|---|
| RT-DETR / DEIM / DETR 계열 | 0.1 |
| YOLO(Ultralytics 기본값) | 10.0 |
| clipping 비활성화 | 0 |
FAQ#
사용자 지정 trainer 클래스(인스턴스가 아님)를
model.train()의trainerparameter에 전달합니다: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클래스가 trainer 인스턴스화를 내부적으로 처리합니다. trainer 아키텍처에 대한 자세한 내용은 고급 사용자 지정 페이지를 참조하시기 바랍니다.사용자 지정에 사용할 수 있는 주요 메서드:
메서드 목적 validate()validation을 실행하고 metric 반환 build_optimizer()optimizer 구성 save_model()학습 checkpoint 저장 get_model()model 인스턴스 반환 get_validator()validator 인스턴스 반환 get_dataloader()dataloader 생성 preprocess_batch()입력 batch 전처리 label_loss_items()로깅을 위한 loss 항목 형식 지정 전체 API reference는
BaseTrainerdocumentation을 참조하시기 바랍니다.예. 더 간단한 사용자 지정에는 callback이 충분한 경우가 많습니다. 사용 가능한 callback event에는
on_train_start,on_train_epoch_start,on_train_epoch_end,on_fit_epoch_end및on_model_save가 있습니다. 이를 사용하면 subclassing 없이 학습 루프에 hook을 연결할 수 있습니다. 위의 backbone 고정 예제에서 이 접근 방식을 보여줍니다.변경 사항이 더 간단한 경우(예: loss gain 조정) hyperparameter를 직접 수정할 수 있습니다:
from ultralytics import YOLO model = YOLO("yolo26n.pt") model.train(data="coco8.yaml", box=10.0, cls=1.5, dfl=2.0)YOLO26에서는 detection head가
reg_max: 1를 사용하므로dfl이 로깅된l1_loss을 조정합니다.reg_max > 1이 있는 모델에서는dfl_loss를 조정합니다.