YOLOv6 custom Train

YOLOv5는 Ultralytics에서 제작되었지만 YOLOv6는 Meituan Vision AI Department 에서 제작했다. 

특이하게 YOLOv6는 YOLOv7보다 늦게 나왔다. 순서대로 정리하면 5 → 7 →6 순서이다. 

여기서는 성능향상 또는 모델의 구조 변경을 다루기 보다는 실제로 custom dataset을 학습 시키는것을 

중점적으로 볼예정이다. 

YOLOv6도 우분투 20.04 LTS, 파이참 가상환경에서 실행하도록 하겠다.이과정은 YOLOv5와 별반 다르지 않다. 

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

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

---

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

URL : https://github.com/meituan/YOLOv6을 입력해 git clone을 진행한다.

디렉토리는 사용자가 편한곳으로 설정한다. (YOLOv5에서 설명한 내용이므로 생략하겠다.)

---

YOLOv6는 python=3.8 환경을 요구한다. 따라서 그에 따른 패키지들을 설치해주면 된다.

Pytorch는 아래의 명령어로 입력을 해준다.

$ conda install pytorch==1.9.1 torchvision==0.10.1 torchaudio==0.9.1 cudatoolkit=10.2 -c pytorch

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

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

이미 설치가 되어 있거나 따로 설치하고 싶은 패키지는 앞에 주석처리 한뒤 명령어를 입력한다. 

$ pip install -r requirement.txt

requirement.txt

---

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

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

dataset 파일 구조는 아래의 사진과 같이 준비를 해야한다. 

이때 images의 .jpg파일과 labels .txt파일 이 일치해야한다. 

custom yolov6 파일구성

---

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

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

custom.yaml

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

is_coco : coco dataset일 경우 True로 아닐경우 False 

nc : class의 갯수 

names : class의 이름 

---

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

YOLOv6에도 모델의 종류가 많다. YOLOv6는 n, s, m, l 4가지 버전이 존재한다. n는 가장 성능이 낮지만 FPS가 가장 높고, l은 성능이 가장 높지만 FPS가 가장 낮은 특징을 가진다.

(-N6, S6 등 이미지의 크기로 분류하는 것 같다.) 본인의 컴퓨터 사양에 맞는 모델을 다운받도록 하자 

yolov6 pretrained weights

이제 학습을 시작하면 된다. 특이하게도 wandb 지원이 아닌 tensorboard를 지원해준다. tensorboard도

실험 과정을 추적하고 기록해주는 tool이다. 

학습을 시작하기 전에 중요한 학습 관련  인자들을 살펴보도록하겠다. 

Parameter	value	Description
--data-path	. /data/coco.yaml	모델을 구성하는 dataset의 경로(yaml 파일의 위치)
--conf-file	. /configs/yolov6s.py	모델 구성 파일
--img-size	640	이미지 크기
--batch-size	32	batch-size
--epochs	400	학습 횟수
--workers	8	dataloader의 갯수
--device	0	사용할수 있는 GPU 갯수(0 / 0,1,2,3 / CPU)
--eval-interval	20	지정한 간격에서 모델 평가

 

---

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

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

$ python tools/train.py --batch 16 --conf configs/yolov6s.py --epochs 400 --img-size 640 --data data/custom1.yaml --device 0

---

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

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

---

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

$ python tools/infer.py --weights runs/train/exp/weights/best.pt --source runs/inference/exp/images/image1.jpg

yolov6 detection result

---

Reference

https://github.com/meituan/YOLOv6

GitHub - meituan/YOLOv6: YOLOv6: a single-stage object detection framework dedicated to industrial applications.

https://docs.ultralytics.com/models/yolov6/?h=yolov6

YOLOv6

https://blog.roboflow.com/how-to-train-yolov6-on-a-custom-dataset/

How to Train YOLOv6 on a Custom Dataset

https://learnopencv.com/yolov6-custom-dataset-training/

YOLOv6 Custom Dataset Training – Underwater Trash Detection

---

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

감사하겠습니다.