[YOLOv5] YOLOv5 커스텀 데이터 학습

YOLOv5란?

YOLOv5(You Only Look Once)는 CNN을 기반으로 object detection을 수행하기 위하여 생성된 모델이다. 타 모델에 비해서 기존에 존재한 카테고리 이외의 새로운 카테고리를 쉽게 학습하고 검출할 수 있는 일반화를 잘하는 경향이 있으며, 모델이 가볍기에 실시간 처리에 적합한다.
YOLOv5 공식 문서 - https://github.com/ultralytics/yolov5.git

YOLOv5 커스텀 데이터 학습

학습을 위한 준비 사항(환경 설정)

필자의 환경

• OS : Windows 11 Home
• Python : Python 3.8

1. YOLOv5 깃허브 클론

cd YOLOv5
git clone https://github.com/ultralytics/yolov5.git

YOLOv5라는 폴더를 생성한 뒤 이동하고, 폴더 안에 클론한다.

2. 필수 라이브러리 설치

cd YOLOv5
pip install -r requirements.txt

YOLOv5/requirements.txt에 들어가보면 필요한 라이브러리가 적혀져 있다. 여기에서 필요한 부분을 설치하도록 한다.

데이터셋 준비

데이터 셋을 준비한다. 이 때 이미지를 전처리를 한 뒤에 사용하도록 하고, 밑에 있는 폴더처럼 데이터를 구성한다.

dataset_root
	- train
    	- images
        	data1.png
            data1.png
            ...
        - label
        	data1.txt
            data2.txt
            ...
    - val
    	- images
        - label
    - test
    	- images
        - label

데이터를 준비한 뒤에 데이터를 YOLOv5가 참고할 수 있도록 yaml 파일을 설정한다.

# custom_dataset.yaml

train: ./path_to/train/images
val: ./path_to/val/images
test: ./path_to/test/images

nc: 3 # number of classes

names: [ 'class1', 'class2', 'class3'] # class names

train, val, test에 각 데이터의 경로를 적어주고, nc에 검출할 클래스의 갯수를, names에 검출할 클래스의 이름을 리스트로 전달한다.

학습

데이터셋이 준비 되었다면 학습을 진행하면 된다. 학습을 할 때 명령줄 인수로 필요한 것들을 설정한다. 필자는 CUDA 세팅이 되어 있지 않기에 CUDA를 설정하지 않았고, 필요하다면 --device 를 이용하여 설정하도록 한다.

python train.py --img 800 --batch 16 --epochs 10 --data custom_dataset.yaml --weights yolov5s.pt

train.py를 이용하여 데이터를 학습시킨다. 이 때 batch와 epoch을 설정해주고, data는 이전에 설정하였던 custom_dataset.yaml을 이용하도록 한다. weights는 어떤 모델을 이용하여 학습을 할 것인지 설정한다. 학습이 끝나게 되면 모델이 YOLOv5/runs/train/exp/weights/best.pt 경로로 생성이 된다.

평가

학습된 모델을 사용하여 모델을 평가할 수 있다.

python detect.py --source ./path_to/image.png --weights ./runs/train/exp/weights.best.pt

명령줄 인수로 --source는 테스트에 이용할 이미지 경로,
--weights는 참조할 모델이므로 학습이 완료된 모델을 이용하였다.

데이터 추가 학습

데이터를 추가적으로 학습시키려면 기존에 학습하였던 모델을 학습할 때 이용하면 된다. 즉, --weights에서 yolov5s.pt에서 이전에 나온 best.pt로 설정해주면 된다.

python train.py --img 800 --batch 16 --epochs 10 --data custom_dataset.yaml --weights path_to/best.pt