[YOLO] YOLOv5-custom 학습(신호등 인지)

 

이번 포스트는 기존 YOLOv3-Darknet 사용법 이후로 YOLOv5(Pytorch)로 버전을 바꾸기 위해 테스트를 진행했습니다.

*현재 작성자 기준 세팅

1. Nvidia driver - 470

2. Cuda - 11.2

3. CuDnn - 8.1.1

4.  Pytorch - 1.8.0 + cu111

 

!!주의!!) 위 세팅 기준으로 작성하였기에 사용자 분들은 현재 GPU와 기본세팅에 맞춰서 준비하셔야합니다.

사용법 참고) https://docs.ultralytics.com/yolov5/tutorials/train_custom_data/#before-you-start

Train Custom Data

위 개발자의 Yolov5-custom 학습 Tutorial을 참고하였습니다.

 

 

git clone https://github.com/ultralytics/yolov5.git #yolov5 Github
pip3 install -r requirements.txt #install

이번 포스트에서도 Roboflow를 통해서 이미지 + 이미지 Bounding Box 데이터를 활용했습니다.

Roboflow

여기서 다운받은 train 폴더를 git clone해서 받아온 yolov5 바로 하위 폴더에 위치하고 같이 압축 되어있는 data.yaml은 ./yolov5/data/ 폴더 안에 놓습니다.

*여기서 이미지 파일의 위치는 편하게 두셔도 됩니다.

data.yaml

중요한것은 data.yaml파일에서 train, valid, test 사용시 위치경로만 잘 맞춰주시면 됩니다.

- nc : class 수

- names : class 명 list

 

파일 경로만 잘 맞는다면 train.py 실행시 문제가 없습니다.

python3 train.py --img 416 --batch 16 --epochs 50 --data ./data/data.yaml --cfg ./models/yolov5s.yaml --weight yolov5s-cls.pt --device 0 --name "작성자 마음대로"

- img: 이미지 사이즈

- batch: 하나의 소그룹에 속하는 데이터 수

- epochs: train 수

- data: 데이터 파일 위치 (data.yaml 위치)

- cfg: Yolov5의 모델 정보.yaml

- weight: yolov5 모델 파일

- device: single-GPU, Multi-GPU 관련 설정(Single일 경우 --device 0, Multi일 경우 --device 0,1)

- name: 학습 후 모델 폴더명

 

세부설명은 향후 바로 업데이트 하겠습니다.