YOLOv5 모델 아키텍쳐 분석 | YOLOv5 Training 및 Inference

 

Object Detection을 하면서 가장 쉽게, 많이 사용하는 모델인 Yolo v5 모델 아키텍쳐를 간단하게 분석을 해 보았다.

Yolo v5 Github : https://github.com/ultralytics/yolov5

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

---

YOLOv5 모델 아키텍쳐 분석

1) Backbone

Yolo v5의 Backbone은 이미지에서 feature map을 추출하는 부분으로 CSPNet기반의 CSP-Darknet을 사용한다.

CSPNet은 CNN의 학습 능력을 향상 시킬 수 있는 Backbone으로 DenseNet에 CSP 기법을 적용시켜 정확도를 유지 시키면서 연산량을 줄여 경량화를 유지하였다. 때문에 CSPNet은 ImageNet과 MS COCO 같은 컴퓨터 비전의 대표적인 데이터에서 성능을 향상 시켰다.

Yolo v5는 Yolo v4와 마찬가지로 이러한 장점을 가진 CSPNet을 이용하여 각 계층의 연산량을 균등하게 분배해서 연산 bottleneck을 없애고 CNN layer의 연산 활용을 업그레이드 시켰다.

Yolo v5가 다른 Yolo 모델들과 다른 이유는 backbone을 다른 모델들과 달리 depth multiple과 width multiple를 기준으로 하여 크기별로 yolo v5 s, yolo v5 m yolo v5 l, yolo v5 x로 나눈다는 것이다. yolo v5-s의 depth&width_multiple이 가장 작고(depth_multiple : 0.33, width_multiple : 0.50), x의 depth&width_multiple은 (1.33, 1.25)로 가장 크다. 따라서 아래 그림과 같이 yolo v5 s가 가장 빠르지만 정확도가 비교적 떨어지며 yolo v5 x는 가장 느리지만 정확도는 향상된다.

Yolo v5 모델 별 성능

 

 

2) Head

Yolo v5의 Head는 물체의 위치를 찾는 부분으로 추출된 feature map을 바탕으로 위치를 찾는다.

Anchor Box를 설정한 후 최종적으로 Bounding Box를 생성하는데, 8 pixel의 정보를 가진 작은 물체, 16 pixel의 정보를 가진 중간 물체, 32 pixel의 정보를 가진 큰 물체를 인식을 할 수 있는 3가지 scale의 bounding box를 생성하게 된다. 각 스케일에서 3개의 앵커 박스를 사용하므로 총 9개의 앵커박스를 사용한다.

 

 

YOLOv5 Training 및 Inference

1) Yolo v5 학습 환경 세팅

(1) 아나콘다 환경 생성 및 필요한 라이브러리 설치

 

Yolo v5 Github에서 제공하는 requirenemts.txt를 참고하여 필요한 라이브러리를 설치한다.

 

(2) pre-trained model 다운로드

https://github.com/ultralytics/yolov5/releases에서 학습에 사용할 pre-trained model을 다운 받을 수 있다.

 

(3) yaml 파일 제작

학습 데이터의 경로, 클래스 개수, 클래스 종류가 적힌 yaml 파일을 제작한다.

 

train, val, nc, names를 학습 데이터셋에 맞춰서 수정한다. 또 학습에 사용할 모델의 yaml 파일에서도 nc를 학습하고자 하는 클래스 수에 맞춰 수정을 해 준다.

 

(4) Yolo v5 데이터셋 만들기

아래 그림과 같이 데이터셋 폴더를 생성한다. data 폴더에 images, labels 폴더와 함께 윗 단계에서 생성한 yaml파일도 함께 넣어준다.

 

 

(5) Yolo v5 모델 학습

Yolo v5 Github에서 제공하는 train.py를 사용하여 모델 학습을 진행할 수 있다.

 

--img : img 사이즈를 지정할 수 있다.

--batch : batch 사이즈를 지정할 수 있다.

--epochs : 학습을 진행할 에폭 수를 지정할 수 있다.

--data : 위에서 생성한 yaml 파일의 경로를 지정해준다.

--cfg : 모델 학습에 사용할 모델의 yaml 파일의 경로를 지정해준다.

--weight : 모델 학습에 사용할 모델의 가중치 파일의 경로를 지정해준다.

--name : 학습 결과가 저장될 폴더의 이름을 지정해준다.

 

(6) Yolo v5 모델 테스트

Yolo v5 Github에서 제공하는 detect.py를 사용하여 학습 결과를 테스트 할 수 있다.

--source : 테스트 할 이미지의 경로를 지정해준다.

--weights : 학습이 완료된 가중치 파일의 경로를 지정해준다.

--conf : 바운딩 박스를 그릴 class score를 지정해준다.