Skip to main content

​
YOLO 학습

이 문서에서는 Ultralytics를 사용하여 커스텀 데이터셋으로 YOLO 탐지 모델을 학습하는 전 과정을 실습합니다.
1
환경 준비
2
pip install ultralytics
3
데이터셋 준비
4
YOLO 형식의 데이터셋 구조를 준비합니다.
5
dataset/
├── images/
│   ├── train/
│   │   ├── img001.jpg
│   │   └── img002.jpg
│   └── val/
│       ├── img101.jpg
│       └── img102.jpg
├── labels/
│   ├── train/
│   │   ├── img001.txt
│   │   └── img002.txt
│   └── val/
│       ├── img101.txt
│       └── img102.txt
└── dataset.yaml
6
dataset.yaml 작성
7
# dataset.yaml
path: /path/to/dataset   # 데이터셋 루트 경로
train: images/train       # 학습 이미지 경로
val: images/val           # 검증 이미지 경로

nc: 3                     # 클래스 수
names: ['person', 'helmet', 'no_helmet']  # 클래스 이름
8
모델 학습
9
from ultralytics import YOLO

# 사전학습 모델 로드
model = YOLO('yolo11m.pt')

# 학습 실행
results = model.train(
    data='dataset.yaml',
    epochs=100,
    imgsz=640,
    batch=16,
    lr0=0.01,            # 초기 학습률
    lrf=0.01,            # 최종 학습률 비율
    optimizer='AdamW',
    patience=20,          # Early Stopping patience
    device=0,             # GPU 번호 (CPU는 'cpu')
    project='runs/detect',
    name='safety_helmet',
)
10
하이퍼파라미터 상세
11
파라미터기본값설명epochs100학습 에폭 수imgsz640입력 이미지 크기batch16배치 크기 (-1로 자동 설정)lr00.01초기 학습률optimizerSGD옵티마이저 (SGD, Adam, AdamW)patience50Early Stopping 인내 에폭mosaic1.0Mosaic 증강 확률mixup0.0MixUp 증강 확률close_mosaic10마지막 N 에폭에서 Mosaic 비활성화
12
학습 결과 확인
13
# 학습 결과 디렉터리 확인
# runs/detect/safety_helmet/
# ├── weights/
# │   ├── best.pt        # 최적 모델
# │   └── last.pt        # 마지막 에폭 모델
# ├── results.csv        # 에폭별 지표
# ├── results.png        # 학습 곡선
# ├── confusion_matrix.png
# ├── F1_curve.png
# ├── PR_curve.png
# └── val_batch0_pred.jpg  # 검증 예측 시각화
14
학습된 모델로 검증
15
# 최적 모델로 검증
model = YOLO('runs/detect/safety_helmet/weights/best.pt')

# 검증 세트 평가
metrics = model.val(data='dataset.yaml')

print(f"mAP@50: {metrics.box.map50:.4f}")
print(f"mAP@50:95: {metrics.box.map:.4f}")
print(f"Precision: {metrics.box.mp:.4f}")
print(f"Recall: {metrics.box.mr:.4f}")

​
CLI 사용법

Python API 대신 명령줄에서도 학습할 수 있습니다. 
# 학습
yolo detect train data=dataset.yaml model=yolo11m.pt epochs=100 imgsz=640

# 검증
yolo detect val data=dataset.yaml model=runs/detect/train/weights/best.pt

# 추론
yolo detect predict model=runs/detect/train/weights/best.pt source=test_images/

​
학습 모범 사례

항목추천이유
모델 크기YOLO11m속도/성능 균형
이미지 크기640YOLO 표준, 필요 시 1280
배치 크기GPU 메모리에 맞게batch=-1로 자동 탐색
에폭100~300데이터 규모에 비례
증강Mosaic + 기본 증강기본 설정 충분
옵티마이저AdamWSGD 대비 빠른 수렴
학습 데이터의 이미지와 레이블 파일명이 정확히 일치해야 합니다. img001.jpg에 대응하는 레이블 파일은 img001.txt입니다. 파일명이 다르면 해당 이미지는 배경(Negative Sample)으로 처리됩니다.

​
트러블슈팅

  1. dataset.yaml의 경로가 정확한지 확인하세요. 2) 레이블 파일의 좌표가 0~1 사이의 정규화 값인지 확인하세요. 3) 레이블 파일의 클래스 인덱스가 0부터 시작하는지 확인하세요.
batch 값을 절반으로 줄이거나, batch=-1로 자동 설정하세요. 이미지 크기(imgsz)를 640에서 480으로 줄이는 것도 효과적입니다.
model='runs/detect/train/weights/last.pt'를 사용하고 resume=True를 추가하면 중단된 지점부터 학습을 재개합니다.