[ YOLOv8 ] Custom Dataset 학습(Google Colab)

https://jinwoo-jung.tistory.com/68

[ YOLOv8 ] Custom Dataset 구축(Roboflow)

---

 

구축한 Custom Dataset을 기반으로 YOLOv8 Model을 학습시켜 보자. Local 환경에서 학습시켜도 되지만, Google Colab을 활용하여 학습시켰다.

 

 

먼저 Colab 파일을 하나 생성한 뒤 GPU를 활용하도록 런타임 유형을 변경한다. !nvidia-smi를 통해 현재 그래픽 카드 버전을 확인한다(필수는 아님)

 

 

그다음 YOLOv8을 위해 ultralytics를 설치한다. 이때 꼭 버전을 맞춰 줘야 학습이 원활하게 진행되기 때문에 주의 바란다. 이후 오른쪽과 같이 Ultralytics가 잘 설치되었는지 확인한다. 아래 버전이 뜨면은 잘 설치 된 것이다.

 

 

학습과 YOLOv8에 필요한 패키지를 import 한 뒤, pre-trained model을 불러와 정상적으로 동작하는 환경 설정이 다 되었는지 확인한다. 오른쪽과 같이 pre-trained model의 detection 결과가 잘 뜬다면 잘 설정 한 것이다.

 

이후 이전에 생성한 Custom Dataset을 설치하는 Code를 전부 복사하여 실행시킨다. 해당 Code는 이전 포스팅에서 설명하였으니 참고 바란다. 

 

 

학습시키기 이전 해당 Dataset의 yaml파일을 수정 해 줘야 한다. 해당 파일 내에는 구축한 Dataset의 정보가 담겨 있는데, 이때 test/train/valid images의 경로를 해당 환경에 맞게 변경 해 줘야 한다.

 

 

이후 각 환경에 맞는 경로 및 Parameter를 토대로 모델을 학습시키면 된다. 필자는 기본 모델을 yolov8s.pt로 설정하여 해당 모델을 학습시켰다.

 

 

학습이 완료 되었다면, 생성된 모델을 기반으로 test/images를 이용하여 검출 성능을 확인할 수 있다. 

 

 

Detection 결과는 다음과 같이 확인할 수 있으며, 학습시킨 모델이 Person 객체를 잘 검출한 것을 확인할 수 있다.