YOLO 모델의 전이 학습(Transfer Learning), Custom Dataset을 이용한 YOLO 학습 방법

이번 글에서는 YOLO(You Only Look Once) 모델의 전이 학습(Transfer Learning) 개념과 Custom Dataset을 사용하여 YOLO 모델을 학습하는 방법에 대해 설명하겠습니다. 컴퓨터 비전 분야에서는 주로 이미지 객체 인식과 분류 작업에서 YOLO 모델이 널리 사용됩니다. YOLO 모델은 빠른 실시간 검출 성능으로 유명하지만, 특정 데이터셋에 맞게 Fine-tuning하는 과정이 필요합니다.

전이 학습(Transfer Learning) 개요

전이 학습이란 이미 학습된 모델을 기반으로 새로운 데이터셋에 대해 재학습하는 방법입니다. 이는 특히 데이터가 적은 경우 매우 유용합니다. YOLO와 같은 대규모 모델은 일반적으로 대량의 데이터로 미리 학습되어 있으며, 이러한 모델을 기반으로 특정 작업에 맞게 수정하여 성능을 극대화할 수 있습니다.

전이 학습의 장점

  • 학습 시간 단축: 대량의 데이터로 학습한 모델을 사용하므로 빠르게 학습할 수 있습니다.
  • 적은 데이터량으로도 성능 향상: 작은 데이터셋으로도 좋은 성능을 낼 수 있습니다.
  • 더 좋은 일반화 성능: 일반적으로 사전 학습된 모델은 이미 많은 패턴을 학습했기 때문에 일반화 능력이 뛰어납니다.

YOLO 모델 개요

YOLO는 한 번의 신경망 평가로 전체 이미지에서 객체를 동시에 탐지하는 방식입니다. YOLO의 구조는 다소 복잡하지만, 기본적으로 이미지를 격자(grid)로 나누고 각 격자에서 특정 객체가 있을 확률을 예측합니다. YOLOv5 모델이 현재 가장 많이 사용되고 있으며, 파이썬에서 YOLOv5로 전이 학습을 수행할 수 있습니다.

YOLO의 기본 개념

  • 격자(cell): 입력 이미지를 SxS의 그리드로 나눕니다.
  • 바운딩 박스 및 점수: 각 격자에서 최대 B개의 바운딩 박스를 예측하고, 각 박스에 대해 신뢰도 점수를 부여합니다.
  • 클래스 확률: 각 격자에 대해 K개의 클래스 확률을 예측합니다.

Custom Dataset 준비하기

YOLO 모델을 재학습하기 위해서는 Custom Dataset이 필요합니다. 이는 이미지와 각 이미지에 대한 주석(annotation) 파일로 구성됩니다. 주석 파일은 일반적으로 YOLO 형식으로 준비됩니다.

YOLO 데이터셋 포맷

YOLO에서 사용되는 데이터셋 포맷은 다음과 같습니다:

  • 이미지 파일: 학습에 사용할 이미지 파일들.
  • 라벨 파일: 각 이미지에 대한 바운딩 박스의 정보를 포함하는 텍스트 파일. 각 줄은 형식이 아래와 같습니다: 
    클래스_id x_center y_center width height

    여기서 (x_center, y_center)는 해당 바운딩 박스의 중심 좌표, width와 height는 바운딩 박스의 너비와 높이를 나타냅니다.

예제 데이터셋

예를 들어, 자동차와 사람을 탐지하는 Custom Dataset을 준비할 수 있습니다. 여기서 샘플 데이터셋을 확인할 수 있습니다.

YOLOv5 설치 및 환경 설정

YOLOv5를 설치하기 위해서는 PyTorch와 기타 종속성 라이브러리가 필요합니다. 아래의 단계를 통해 설치할 수 있습니다:

1. 환경 설정

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

2. 데이터셋 구성

Custom Dataset 디렉토리 구조는 다음과 같이 설정합니다:

yolov5/
│
├── datasets/
│   ├── my_custom_dataset/
│   │   ├── images/
│   │   ├── labels/
│   │   └── data.yaml

3. data.yaml 파일 예시

data.yaml 파일은 YOLO 모델에 데이터셋 정보를 제공합니다. 파일은 다음과 같은 형식을 가져야 합니다:

train: datasets/my_custom_dataset/images/train
val: datasets/my_custom_dataset/images/val

nc: 2
names: ['car', 'person']

YOLOv5 모델 훈련하기

위의 준비가 끝났다면, YOLOv5 모델을 학습할 수 있습니다. 아래 명령어를 사용하여 학습을 시작합니다.

훈련 명령어

python train.py --img 640 --batch 16 --epochs 50 --data datasets/my_custom_dataset/data.yaml --weights yolov5s.pt

여기서 각 인자의 의미는 다음과 같습니다:

  • --img: 입력 이미지의 크기. 일반적으로 640 또는 416을 사용합니다.
  • --batch: 배치 사이즈입니다.
  • --epochs: 전체 훈련 반복 수입니다.
  • --data: 데이터셋에 대한 yaml 파일 경로입니다.
  • --weights: 초기 가중치 파일입니다. 최전의 pretrained weights를 사용할 수 있습니다.

모델 성능 평가

훈련이 완료되면, 모델의 성능을 평가할 수 있습니다.

python val.py --weights runs/train/exp/weights/best.pt --data datasets/my_custom_dataset/data.yaml --img 640

평가 결과 확인

성능 평가는 mAP(Mean Average Precision) 및 IOU(Intersection Over Union) 등의 지표를 통해 확인할 수 있습니다. 에폭 반복 후 트레이닝 로그에 이러한 지표가 기록됩니다.

결론

YOLO 모델의 전이 학습을 통해 Custom Dataset을 이용하여 원하는 객체를 탐지할 수 있는 강력한 모델을 구축할 수 있습니다. YOLOv5의 기능을 활용하여 많은 경우에 훌륭한 성능을 얻을 수 있음을 알았습니다. 본 글에서는 기본적인 과정만을 다루었지만, 하이퍼파라미터 조정이나 데이터 증강 등의 추가적인 기술을 통해 성능을 더욱 높일 수 있습니다.

참고 자료