YOLOv5 custom Train

YOLOv5는 Ultralytics에서 제작되었다. 따라서 실행환경이 YOLOv3, YOLOv4와 달라 졌는데 

이번 YOLOv5는 우분투 20.04 LTS, 파이참 가상환경에서 실행하도록 하겠다. 

일단 아나나콘다 환경에서 아래의 명령어를 입력해 yolov5 가상환경을 생성하고 가상환경을 실행한다. 

$ conda create -n yolov5 pyton=3.8
$ conda activate yolov5

---

가상환경을 만든후 파이참에서 다음과 같이 'vsc에서 받기'를 선택하고

URL : https://github.com/ultralytics/yolov5 을 입력해 git clone을 진행한다. 디렉토리는 사용자가 편한곳으로 설정한다. 

yolov5 git clone

프로젝트가 생성이 되었으면 필요한 패키지 설치를 진행하도록 하겠다. 

YOLOv5의 공식문서에서  Python>=3.8.0 및 PyTorch>=1.8 환경에서 실행이 되므로 다음과 같이

파이토치 설치를 가상환경에 설치를 해준다.(버전은 1.8 이상만 충족이 되면 되자만 다른 패키지와도 호환이 되어야 합니다.) 

$ conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.3 -c pytorch

아래의 사진과 같이 requirements.txt 에 필요한 패키기 리스트가 나와있다. 

아래의 명령어를 입력해 한번에 설치가 가능하다. 필요한 패키지는 다 설치 해주도록 한다. 

$ pip install -r requirement.txt

requirement.txt

---

필요한 패키기 설치가 끝났으면 custom dataset을 정리해주어야 한다. 

custom dataset은 YOLO_mark(다른 annotation tool)을 사용해 학습할 이미지, txt파일(YOLO_format)을 준비한다. 

YOLOv3, YOLOv4와 다르게 dataset 파일 구조가 달라졌다. 아래의 사진과 같이 준비를 해야한다. 

imgaes에는 .jpg파일이 / labels파일에는 .txt파일이 들어가야한다. 이때 .jpg파일과 .txt파일 이 일치해야한다. 

custom yolov5 파일구성

---

다음은 custom.yaml파일을 만들어 준다. 모델의 구성파일이다. 

아래의 사진과 같이 구성을 해주면 된다. 

custom.yaml

path : train, validation, (test) 이미지와 annotation이 있는 폴더의 경로 

train : train 이미지의 경로 / val : validation 이미지의 경로 

nc : class의 갯수 

names : class의 이름 

---

pretrain 모델을 이용하면 좀 더 빠르고 정확한 학습을 진행할 수 있다. pretrained된 모델을 선택한다. 

YOLOv5에는 모델의 종류가 많다. Yolo v5는 s, m, l, x 4가지 버전이 존재한다. s는 가장 성능이 낮지만 FPS가 가장 높고, x는 성능이 가장 높지만 FPS가 가장 낮은 특징을 가진다. 본인의 컴퓨터 사양에 맞는 모델을 다운받도록 하자 

yolov5 pretrained model

이제 학습을 시작할 차례이다. 사전에 wandb와 같은 딥러닝 실험 과정을 손쉽게 Tracking하고, 시각화할 수 있는 Tool을 사용하면 더 쉽게 추적을 할수 있다. 관련 내용은 추후에 작성하도록 하겠다. 

우선 train.py에 들어가면 훈련관련 인자들을 확인 할수 있다. 중요한 인자들만 살펴 보도록 하겠다. 

train.py

Parameter	Value	Description
--weigths	ROOT(yolov5s.pt)	모델의 가중치 파일있는 경로
--cfg	' '	모델의 구조 파일이 있는 경로, (--weight 나 --cfg는 둘중 하나는 반드시 입력해야함)
--data	ROOT(data/coco128.yaml)	데이터에 대한 정보(train/val/test 경로, 클래스 개수, 클래스 이름)가 담긴 yaml파일의 경로 입력.
--hyp	ROOT(data/hyps/hyp.scratch-low.yaml)	하이퍼파라미터 경로
--epochs	100	학습 횟수
--batch-size	16	batchs-size, -1입력시 자동으로 설정
--imgz, --img, --img-size	640	train, validation image의 크기
--device	''	GPU의 갯수 1개이면 0 / 여러개이면 0,1,2,3 등 / CPU도 가능하다
--optimizer	SGD	optimizer 선택 SGD, Adam, AdamW 등

---

이제 학습을 시작할 차례이다. 아래의 명령어를 입력하면 훈련이 시작된다. 

yaml, weights 파일의 위치만 주의해서 적어주면 된다. 

$ python train.py --img 640 --batch 8 --epochs 200 --data custom.yaml  --weights yolov5s.pt --device 0 --project test1

---

wandb와 연동이 되어 있다면 학습 진행중 결과를 실시간으로 확인을 할수 있다. 

학습이 끝나면 다음과 같이 결과파일도 확인 할 수 있다. 

train result

confusion matrix

---

아래의 명령어를 입력해 detection을 진행 할수 있다. 

$ python detect.py --weights best.pt --source 0 (웹캠) / img.jpg (사진) / vid.mp4 (비디오)

webcam detection

---

Reference

https://wana.tistory.com/25

[Python] YoloV5 custom dataset 학습

https://ropiens.tistory.com/44

YOLO V5 환경 셋팅 및 모델 아키텍쳐 분석하기

https://github.com/ultralytics/yolov5

GitHub - ultralytics/yolov5: YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite

https://docs.ultralytics.com/yolov5/quickstart_tutorial/

Quickstart

https://goatlab.tistory.com/entry/YOLOv5

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

---

정확한 정보 전달보단 공부 겸 기록에 초점을 둔 글입니다. 틀린 내용이 있을 수 있습니다.
틀린 내용이나 다른 문제가 있으면 댓글 남겨주시거나 또는 이메일로 보내주시면

감사하겠습니다.