[Detection] Yolov5 custom dataset 학습 및 추론방법

2022년 5월 24일 기준 작성된 글입니다. 수정사항이 있으면 댓글로 말씀 부탁드립니다.

 

실행환경

os: window10 

python: 3.8

pytorch: 1.10.0

cuda: 11.3

cudnn: 8.2.1 for cuda 11.x

---

1. 가상환경 생성

#가상환경 생성
conda create -n yolov5 python=3.8

#가상환경 실행
conda activate yolov5

2. pytorch 설치 

cuda 11.3에 맞는 pytorch 설치

# CUDA 11.3
conda install pytorch==1.10.0 torchvision==0.11.0 torchaudio==0.10.0 cudatoolkit=11.3 -c pytorch -c conda-forge

 

3. pytorch cuda 사용가능 여부 확인

conda activate yolov5
python
>>import torch 
>>torch.cuda.is_available()

True를 반환하면 Good!

 

4.  YOLOv5 git clone

https://github.com/ultralytics/yolov5

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

#clone
git clone https://github.com/ultralytics/yolov5

#yolov5 폴더로 이동 
cd yolov5

내 컴퓨터(GPU)에 맞는 dependency를 맞추기 위해 yolov5에서 권장하는 requirement.txt 파일을 설치하지 않는다!!

pip instasll requirement.txt 

하지말 것!!

 

5. 디렉토리 구조

yolov5는 custom dataset을 트레이닝하기 위해 다음과 같이 디렉토리 구조를 만들어줘야 한다.

나는 labelimg 라는 라벨링 툴을 이용해 yolo 형식으로 저장한 뒤 디렉토리 구조에 맞게 이미지(.jpg)와 라벨링파일(.txt)을 넣어주었다.

 

- yolov5 디렉토리 내에 coco128이라는 폴더 하나 생성

- coco128 폴더내에 images, labels 폴더 2개 생성

- images, labels 각 폴더내에 train2017 폴더 생성

- images/train2017 폴더에 원본이미지 file 넣기

- labels/train2017 폴더에 라벨링 txt file 넣기

 

최종 디렉토리 구조는 다음과 같다.

• yolov5
  • coco128
    • images
      • train2017
    • labels
      • train2017

 

 

6. data 폴더 내 coco128.yaml파일 수정

위와 같이 나머지 경로는 다 주석처리하고 

train: 위 사진 참조

val: 위 사진 참조

test: 주석처리 할것

nc: 내가 학습시키고자 하는 class 개수

names: 내가 학습시키고자 하는 class 이름

 

 

7. yolov5 폴더 내 train.py 파일 들어가서 num_worker default=0 으로 수정

train.py

위 line의 default=8로 되어있을텐데 0으로 수정한다. 

(위 과정을 진행하지 않으면 페이징 오류 발생!)

 

8. 내가 사용할 pre-trained 모델 확인

yolov5에는 n,s,m,l,x 5가지 버전이 있다. n부터 가장 가벼운 모델 x가 가장 무거운 모델이다. 

n는 모델이 가볍지만 mAP가 가장 낮고 x는 모델이 무겁지만 mAP가 높다.

나는 이번에 yolov5s 모델을 사용 할 것이다. 

mAP 비교

 

9. Training

python train.py --img 640 --batch 16 --epochs 100 --data coco128.yaml --weights yolov5s.pt --device 0

학습이 완료되면 yolov5 폴더 내에 runs 라는 폴더가 생성된다. 

yolov5/runs/train/ 폴더에 들어가보면 exp2, exp3 등 트레이닝 횟수대로 폴더가 생성되어 있는데 

최신 폴더에 들어가보면 weight 폴더에 best.pt 파일을 weight file로 쓰면 된다!! 

 

 

10.Inference

python detect.py --source=C:/Users/MSDL-DESK-02/yolov5/data/images/egg.jpg --weights=C:/Users/MSDL-DESK-02/yolov5/runs/train/yolov5_egg/weights/best.pt --img 416 --conf 0.5 --save-txt

--img : 네트워크에 입력되는 이미지 크기

--batch : 배치 크기
--epochs : epoch 횟수
--data : 학습에 참조될 yaml file 
--weights : pre-trained model에 전이학습할 경우 (yolov5s.pt 등을 넣어주면 자동으로 다운로드 됨)
--name : 학습된 모델의 저장이름

 

끝 :)