[YOLO] 커스텀 데이터셋 학습시키기 - colab에서 Yolo5 모델

 

 

봉식이랑 다른 고양이들을 구분하려면 다른 고양이 사진들도 같이 학습시켜야하는데 완전... 잊어버리고 있었다.

 

그래서 일단 봉식이만을 인식하는 모델 먼저 만들어보기로 했다.

 

 

 

 

https://blog.roboflow.com/how-to-train-yolov5-on-a-custom-dataset/?ref=ultralytics 

How to Train YOLOv5 On a Custom Dataset

roboflow 블로그 글을 참고했다.

 

 

 

 

1. 드라이브 마운트

# 드라이브 마운트
from google.colab import drive
drive.mount('/content/drive')

 

봉식이 사진을 구글 드라이브에 넣어놓음 + 학습 완료된 모델을 구글 드라이브에 옮길 예정

이여서 구글 드라이브에 마운트를 해주었다.

 

 

 

2. 환경 셋팅

# 내 구글 드라이브로 이동
#%cd "/content/drive/MyDrive"

# Yolov5 github 레포지토리 clone
!git clone https://github.com/ultralytics/yolov5.git

# 필요한 모듈 설치
!pip install -U -r yolov5/requirements.txt

 

구글 드라이브에 설치해도 되는데 그냥... 용량이 아까워서 그러지는 않았다.

 

 

 

import torch

#파이토치 버전 확인, cuda device properties 확인
print('torch %s %s' % (torch.__version__, torch.cuda.get_device_properties(0) if torch.cuda.is_available() else 'CPU'))

 

torch를 import해주고 버전 확인, cuda 사용 가능 여부와 cuda device properties를 확인한다.

 

혹시 CPU라고 뜬다면 런타임 > 런타임 유형 변경에 들어가서 바꾸어주면 된다.

 

 

 

 

3. 데이터셋 다운로드

 

나는 드라이브에 올려놓아서 스킵했지만 roboflow에서 데이터셋을 export할 때 zip이 아닌 다운로드 코드를 받는 방식으로 진행했다면...

 

!mkdir datasets
%cd /content/datasets/

!curl -L "https://app.roboflow.com/ds/링크입력" > roboflow.zip; unzip roboflow.zip; rm roboflow.zip

 

이렇게 데이터를 받아오면 된다. 

 

혹은

 

!pip install roboflow

from roboflow import Roboflow
rf = Roboflow(api_key="APIkey입력")
project = rf.workspace().project("프로젝트명입력")
dataset = project.version(버전입력).download("yolov5")

이렇게 roboflow pip 패키지를 이용한 방법도 있다.

 

 

4. data.yaml파일 수정

 

roboflow에서 데이터를 받았다면 기본적으로 data.yaml파일이 있는데,

여기서 데이터셋 위치를 수정해줘야한다.

 

from glob import glob

# 이미지 경로 list로 넣기
train_img_list = glob('./train_bong/images/*.jpg') + glob('./train_bong/images/*.jpeg')
valid_img_list = glob('./valid_bong/images/*.jpg') + glob('./valid_bong/images/*.jpeg')


# txt 파일에 write
with open('./train.txt', 'w') as f:
	f.write('\n'.join(train_img_list) + '\n')
    
with open('./valid.txt', 'w') as f:
	f.write('\n'.join(valid_img_list) + '\n')

 

이건 튜토리얼에 안나와있어서 구글링했다...

기본 data.yaml 파일에는 

 

train: ../train/images
val: ../valid/images

 

이렇게 경로로만 적혀있고 roboflow 튜토리얼에도 해당 내용이 없어서 이미지 파일들이 있는 경로만 지정해줘도 되는건가? 싶었는데 안되는것같다... ;; 

 

train: train.txt 파일 경로
val: valid.txt 파일 경로

 

이렇게 바꾸어주면 된다.

 

 

아앗... 참고로 저번에 test셋도 만들었는데... 그냥 없애버렸다... 

 

 

5. 학습시키기

 

대망의 학습 타임... 

 

일단 실패한 사진 먼저 던지고 시작

 

 

ㅋ..........

 

 

!python /content/yolov5/train.py

로 train.py 파일을 실행시키면 되는데 몇가지 properties를 설정해주어야한다.

 

img: 입력 이미지 크기
batch: 배치 크기
epochs: 학습 epoch 수 (참고: 3000개 이상이 일반적으로 사용된다고 한다...)
data: data.yaml 파일 경로
cfg: 모델 구성 지정

weights: 가중치에 대한 사용자 정의 경로를 지정합니다(참고: Ultraalytics Google Drive 폴더에서 가중치를 다운로드할 수 있습니다).
name: 모델이 저장 될 폴더 이름
nosave: 최종 체크포인트만 저장
cache: 더 빠른 학습을 위해 이미지를 캐시

 

 

Yolo5의 튜토리얼에는 

# Train YOLOv5s on COCO128 for 3 epochs
!python train.py --img 640 --batch 16 --epochs 3 --data coco128.yaml --weights yolov5s.pt --cache

이렇게 나와있다.

 

!python /content/yolov5/train.py --img 320 --batch 16 --epochs 30 --data /content/datasets/data.yaml --weights yolov5x.pt --name result_bong --cfg /content/yolov5/models/yolov5x.yaml

이건 내가 쓴 코드. 1시간정도 걸렸다.

 

참고로 저 실패한 사진에는 yolo5s 모델을 사용했고 epoch도 낮게 설정했었다.

 

 

혹시 RuntimeError: CUDA out of memory가 뜬다면 batch 사이즈를 줄인다.

 

 

모델 학습이 완료되면

/content/yolov5/runs/train에 설정한 이름의 폴더가 생성된다.

 

 

 

 

 

6. 확인쓰~

 

!cp -r /content/yolov5/runs/train/result_bong /content/drive/MyDrive/models

학습시간은 보통... 엄청엄청 길고... 

n시간 걸린 내 소중한 모델.. 런타임 웅앵 하면서 날리면 눈물나니 내 구글드라이브로 복사해준다.

 

 

 

 

!python /content/yolov5/detect.py --weights /content/yolov5/runs/train/result_bong/weights/best.pt --img 320 --conf 0.5 --source /content/drive/MyDrive/img/test2.jpg

 

detect를 위해 detect.py 파일을 실행시켜주는데

weights에 방금 학습시켜서 나온 pt파일을 넣어주고 source에 test용 이미지를 넣어주었다.

 

 

이제 detect 작업이 끝난 이미지를 확인해보자.

import cv2
from google.colab.patches import cv2_imshow

img = cv2.imread('/content/yolov5/runs/detect/exp/test2.jpg')
cv2_imshow(img)

코랩에서는 cv2.imshow를 못쓴다... 는걸 알게됐다.

 

 

흑흑 뿌듯....

남의 고양이 사진도 한번 넣어봤다.

 

 

 

 

 

 

 

 

 

 

..........

 

개나소나 봉식이인 것이다.....

아마 봉식이 사진만 학습시켜서 고양이면 다 봉식이라고 하는 느낌....

 

 

다음 모델은 남의 고양이와 봉식이를 구분해보도록...

근데 될까? 약간 자신이 없어졌다 흑흑~