[YOLOv5] 커스텀 데이터 학습 - Google Colab (1)

Google Colab

Google Colab이란 클라우드 기반 머신러닝 및 데이터 사이언스 툴로, 파이썬을 사용한다. 쉽게 말해 머신러닝 모델을 학습시킬 수 있는 환경을 구글이 제공해 주는 서비스이다. 따라서 구글 계정이 있다면 누구나 colab을 사용하여 머신러닝 모델을 만들 수 있다. 사용하기 전 cpu와 gpu 중 하나를 런타임 유형으로 선택하여 환경을 조성할 수 있다. gpu를 선택하면 학습을 보다 빠르게 할 수 있지만 무료버전은 gpu 사용이 하루 최대 12시간, runtime 90분으로 제한되어 있다.

---

YOLOv5 학습시키기

https://www.youtube.com/watch?v=T0DO1C8uYP8&ab_channel=%EB%B9%B5%ED%98%95%EC%9D%98%EA%B0%9C%EB%B0%9C%EB%8F%84%EC%83%81%EA%B5%AD 

유튜버 [빵형의 개발도상국]님의 "YOLO v5 커스텀 학습 튜토리얼" 영상을 참고하면 자세하게 설명되어 있다.

 

1. YOLOv5 환경설정

먼저 Google Colab에 YOLOv5를 다운로드해 준다. YOLOv5 링크: https://github.com/ultralytics/yolov5

%cd /content
!git clone https://github.com/ultralytics/yolov5.git

다음으로 YOLOv5을 위한 패키지를 설치해 준다.

%cd /content/yolov5/
!pip install -r requirements.txt

 

2. Dataset 준비하기

Dataset에는 여러 가지 종류가 있는데 roboflow의 권총 dataset으로 실습을 하였다.

https://public.roboflow.com/object-detection/pistols

Pistols Object Detection Dataset

다운로드 코드를 받아서 Dataset을 Google Colab에 받아오면 된다.

 

3. data.yaml 경로 재설정 및 train, valid set 나누기

아래 코드를 실행하면 data.yaml 파일 안의 train 경로와 valid 경로가 각각 train: ../train/images, valid: ../valid/images로 설정되어 있는 것을 확인할 수 있다. 경로가 정확하지 않기 때문에 경로를 바꿔준다.

%cat /content/dataset/data.yaml

먼저 export파일 안에 있는 이미지들을 list로 만들어준다.

%cd /
from glob import glob
img_list = glob('/content/dataset/export/images/*.jpg')
print(len(img_list))

다음에 img_list를 train set과 valid set으로 나눠준다.

from sklearn.model_selection import train_test_split
train_img_list, val_img_list = train_test_split(img_list, test_size = 0.2, random_state = 2000)
print(len(train_img_list), len(val_img_list))

나누어 주었던 train set과 valid set을 각각 txt파일로 만들어준다. 그러면 이미지의 주소들이 txt파일에 적히게 된다.

with open('/content/dataset/train.txt', 'w') as f:
  f.write('\n'.join(train_img_list) + '\n')

with open('/content/dataset/val.txt', 'w') as f:
  f.write('\n'.join(valid_img_list) + '\n')

경로를 재설정해주면 끝이 난다.

import yaml

with open('/content/dataset/data.yaml', 'r') as f:
  data = yaml.safe_load(f)

print(data)

data['train'] = '/content/dataset/train.txt'
data['val'] = '/content/dataset/val.txt'

with open('/content/dataset/data.yaml', 'w') as f:
  yaml.dump(data, f)

print(data)

 

4. YOLOv5 학습

대망의 학습 시간. train.py를 실행시켜 YOLOv5를 학습시켜 준다.

%cd /content/yolov5/

!python train.py --img 416 --batch 16 --epochs 50 --data /content/dataset/data.yaml --cfg ./models/yolov5s.yaml --weights yolov5s.pt --name results

다음과 같은 화면이 나오면 성공적으로 학습을 시킬 수 있다. Google Colab 런타임 유형을 gpu로 설정해 주었다면 사진과 같이 "CUDA: 0"이라고 적힌 화면을 볼 수 있다 (cpu로 설정하였다면 cpu라고 나온다).

총 50개의 Epoch 후에 학습이 완료된다.

 

5. 테스트하기

물체를 인식하기 위해서는 사진이나 동영상이 필요하다. 나는 영상을 이용하여 권총을 인식해 보았다. 파일에 mp4파일을 업로드해 준다.

그다음 --source 부분에 업로드한 동영상 위치와 파일명을 적어준다. 나의 경우 "../gun.mp4"였다.

!python detect.py --weights /content/yolov5/runs/train/results/weights/best.pt --img 416 --conf 0.5 --source ../gun.mp4

결과물을 보면 정상적으로 권총을 인식하는 걸 볼 수 있다.

---

Google Colab으로 YOLOv5 커스텀 학습을 해보았다. 클라우드이기 때문에 컴퓨터 사양이 맞춰지지 않을 때 사용하기 좋을 것 같다. 다음 포스팅에서는 데이터셋을 받아오는 것이 아니라 직접 만들어 보고 로컬 환경에서 모델을 가동하여 웹캠을 사용해 볼 예정이다. 여담으로 /content/yolov5/runs/train/results/weights/best.pt 경로에 있는 best.pt 파일은 다운로드하여 보관하는 것이 좋다. 이유는 best.pt가 학습된 모델이고 Google Colab은 런타임 세션이 종료되면 모든 파일이 없어지기 때문이다. best.pt 파일만 있다면 언제든지 다시 YOLOv5를 다운로드하여서 모델을 실행시킬 수 있다.