[YOLO] YOLOv5 모델 학습 / 예제 코드

 

[YOLO] Roboflow로 라벨링하기 (Multiclass Classification)

위의 글에서 이어, Roboflow로 라벨링 작업을 한 데이터셋을 바탕으로 실제 YOLOv5 모델을 학습시켜보겠습니다. 

 

 YOLOv5 깃허브 클론

git clone https://github.com/ultralytics/yolov5
cd yolov5
pip install -r requirements.txt

 

 

학습 예제 코드 

아래의 코드로 생각보다 간단히 학습을 돌릴 수 있습니다.

import subprocess

subprocess.run([
    "python", "train.py",
    "--img", "416",                       # 이미지 크기
    "--batch", "16",                      # 배치 크기
    "--epochs", "50",                     # 학습 횟수
    "--data", "dataset/data.yaml",        # 데이터 설정 파일 경로
    "--weights", "yolov5s.pt",            
    "--name", "custom-model",   
    "--device", "0",                      # GPU 사용 (GPU 0번 사용)
    "--exist-ok"                         
])

 

• yolov5s.pt: 경량화된 YOLOv5 모델
• 결과는 /runs/train/custom-model 폴더에 저장됨

 

 

학습 결과 확인

• results.png: Precision/Recall 그래프 등 시각화
• weights/best.pt: 가장 성능 좋은 모델 (추론할 때 사용)

 

학습이 끝나면 위와 같은 결과 파일이 생성됩니다.

 

YOLOv5는 학습 중 매 epoch마다 validation 데이터를 기준으로 모델의 성능을 평가합니다.

이때 mAP(mean Average Precision) 등의 지표를 기준으로 가장 성능이 좋았던 시점의 모델을 runs/train/모델이름/weights/best.pt에 저장합니다. best.pt 모델을 이용해 실제 이미지에 대해 객체 탐지를 수행할 수 있습니다.