yolov5 를 customize해서 나만의 object detection 모델을 만들어보기

친구가 정리해준 내용을 바탕으로 제가 모델 training을 담당했습니다

 

개발환경 : google colab

                런타임 유형 gpu

                내 드라이브를 mount 시킨 상태

 

사전작업

+ 내 드라이브에서 colab 파일 생성 

 

 

 

+ 내 드라이브 mount 하기 

 

 

이 버튼을 클릭 하거나

 

코드를 실행하고 뜨는 브라우저에 들어가면  있는 코드를 복붙 해주면 된다.

 

마운트를 완료하면 내 드라이브에 있는 파일들을 사용하거나  볼 수 있다

 

---

colab에서 pwd가 /content/drive/My\ Drive 일 때를 기준으로 작성한 것입니다.

 

1. 먼저 yolov5 를 clone 

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

위 코드를 실행시키면 내 드라이브에 yolov5라는 폴더가 생긴다.

%cd yolov5 
!pip install -r requirements.txt # install dependencies 

yolov5 폴더로 이동하여 requirement.txt에 포함되어 있는 것들을 모두 다운 받아준다.

 

2. customized 모델에 사용할 데이터셋 만들기

 

(이 경우 저는 사진을 모아 직접 라벨링을 했습니다 -> labelimg 라이브러리 사용) 

developeryoung.tistory.com/16 에서 더 자세히 볼 수 있음

 

3.  data.yaml 파일 생성 

 

data.yaml : 모델 학습을 위한 작은 튜토리얼 같은 것 (내 드라이브/dataset 아래에 만들어준다)

data.yaml 내용

파일의 내용은 그림과 같이 적어주면 된다. (우리는 car라는 클래스만 사용했기 때문에 1개)

 

이때 디렉토리 이름은 자기가 만들어놓은 것과 같이 해야하는게 중요 !!

 

nc는 자신이 학습시키고자 하는 클래스의 수(number)고

names에는 그 클래스의  이름을 배열로 적어주면 됨

여러개일 경우 ['class1','class2'] 처럼..

 

 

4.  labels/images 폴더 정리해주기

 

2번에서 만들어놨던 내 데이터셋들을 모두 정리해줘야한다.

 

나는 내 드라이브에서 images 폴더와 labels 폴더를 만들었다. 

images/train 에는 훈련시키고자 하는 image들을 넣고

images/val 에는 validation에 사용되는 image들, 

labels/train  훈련시키는 image의 바운딩 박스 정보가 저장된 txt파일들

labels/val 에는  validation에 사용되는  image의 바운딩 박스 정보 txt파일들 을 모두 업로드 해준다.

 

 

 

5. 모델 선택하기

 

yolov5/models 에 여러 파일들이 있다. 그 중 하나 선택하여 해당 파일 내용중 nc 를 자신이 학습시키고자 하는 클래스 개수로 바꾼다.

(나는 yolov5s.yaml를 사용했음, 뭔 차인지는 자세히는 모름)

 

 

 

yolov5s.yaml 내용

6. training 시키기 !!

 

training 시키기 전에 항상 이 코드를 실행해준다 (왜인지 이걸 Install 안하면 항상 train 코드가 에러 남)

 

%cd /content/drive/My\ Drive/yolov5

코드를 통해 yolov5 디렉토리로 이동

 

!python train.py --img 640 --batch 30 --epochs 100 --data /content/drive/My\ Drive/dataset/data.yaml --cfg ./models/yolov5s.yaml --weights yolov5s.pt --name [내가 정한 폴더명]

 

를 실행하면 훈련이 시작된다.

 

 

argument에 대한 설명은 github.com/ultralytics/yolov5에서 볼 수 있다.

 

훈련이 시작되면

 

이런게 출력되고

 

학습이 완료되면 어느 폴더에 .pt 파일이 생성되었는 지 알려준다.