[Python] 구글 Colab을 활용한 YOLO 모델 학습

구글 Colab의 경우 클라우드 기반의 주피터 노트북 개발환경으로

별도의 파이썬 설치가 필요 없으며

데이터 분석에 용이한 패키지들이 기본적으로 내장되어 있다.

 

또한 구글 드라이브 마운트 시

자신의 구글 드라이브, 깃과 연동하여 파일을 클라우드 형태의

저장소에 보관하며 협업 및 관리할 수 있다.

 

가장 큰 장점이라면 자신이 저사양 노트북이나 PC 환경이더라도

공간 제약을 받지 않으며 Google Colab에서 제공하는 GPU를

제한적이지만 무료로 사용 가능하다는 것이다.

 

 

우선 YOLO 모델을 Colab 가상 환경에서 학습하기 위해

Colab에 접속한다.

https://colab.google/

colab.google

 

 

Colab 환경을 사용하기 위해 로그인 후 좌측 상단의 파일>새 노트를 클릭하여

새로운 주피터 노트북을 열어준다.

런타임에 연결되면 아래와 같이 sample_data 폴더를 볼 수 있다.

 

해당 가상환경으로 Colab을 사용하게 되면 이후 런타임 연결 해제 시

업로드 했던 파일이나 실행 결과들이 없어진다

(별도로 드라이브에 저장하거나 파일로 저장하지 않는 경우)

 

이 때문에 드라이브 마운트 기능을 사용하여 구글 드라이브와 연동하여

파일을 관리해주도록 한다.

좌측 상단의 드라이브 폴더 아이콘 클릭 하여 드라이브 마운트 기능을 실행한다.

 

해당 작업들이 순차적으로 완료가 되면 좌측 하단의 드라이브 폴더

아이콘에 빗금이 생기고 자신의 구글 드라이브 폴더와 연동된 것을 볼 수 있다.

 

이제 Colab 노트 작성 후 해당 파일을 드라이브에 저장하고

프로젝트 파일 또한 동시에 관리할 수 있다.

좌측 상단의 노트 파일명을 변경해주고 저장하면

drive/MyDrive/Colab Notebooks 경로에 자동 저장되는 것을 확인할 수 있다.

추후 파일 경로 설정 시 공백 문자 처리를 따로 해줘야 하므로

본인은 미리 ColabNotebooks로 폴더명을 변경해주었다.

 

 

이제 본격적으로 YOLO 학습을 위한 환경을 설정해 보겠다.

우선 모델 학습 시 GPU를 통해 학습을 진행할 것이므로

런타임 유형을 CPU -> GPU로 변경 해주어야 한다.

 

좌측 상단의 런타임 > 런타임 유형 변경 버튼을 클릭하여 변경한다.

런타임 유형 변경 시 자동으로 런타임이 재시작 되고

조금 기다리면 구글 드라이브 마운트도 자동으로 된다.

 

이제 아래의 코드를 작성하여 정상적으로 GPU가 할당되어 있는지 확인한다.

(현재 GPU 성능과 CUDA 버전을 확인할 수 있다)

 

본인이 사용하고자 하는 YOLO 버전의 git을 cloning 해준다.

이 때 각 YOLO 버전에 따라 해당 라이브러리 유의점 등이 다를 수 있으므로

각 내용들은 원 제작자의 깃허브를 참고한다.

나는 YOLOv7을 클론하여 진행했다.

클론 후 필요 라이브러리들을 설치 한다.

 

이제 훈련시킬 데이터 셋이 필요하다.

해당 페이지에서는 roboflow 패키지를 통해 자동으로 데이터셋을

다운로드 하는 방법을 안내하고 있으니 참고하자.

 

https://colab.research.google.com/drive/1X9A8odmK4k6l26NDviiT6dd6TgR-piOa?ref=blog.roboflow.com 

Training YOLOv7 on Custom Data

 

나는 이전에 포스팅 했던 방법을 통해 준비한

이미지 데이터 셋을 사용할 것이다.

---

[Python] YOLO 모델 학습을 위한 오픈 데이터 셋 준비 (1)

[Python] YOLO 모델 학습을 위한 오픈 데이터 셋 준비 (2)

---

해당 방법들을 통해 이미지 데이터 셋이 준비되었다면

YOLOv7에서 제공하는 여러 모델의 가중치 파일을 선택하여

yolov7 폴더에 다운로드 받는다.

(전 YOLOv7x 모델을 사용했습니다.)

정상적으로 파이토치 형태의 가중치 파일을 다운로드 받은 것을 확인할 수 있습니다.

 

해당 이미지 데이터셋을 이제 학습 모델에 맞춰 훈련 하기 위해

몇 가지 파일들을 수정 및 생성해야 해준다.

 

수정 및 필요한 파일 목록은 아래와 같습니다.

- test_data.yaml [사용자 훈련 데이터의 정보를 나타내는 파일]

- yolov7x-custom.yaml [--cfg 파라미터의 yolov7x.yaml 파일의 사용자 수정본]

- hyp.scratch.custom.yaml [모델 하이퍼파라미터 설정 파일]

 

우선 훈련 데이터의 정보를 담을 .yaml 형식의 파일을 생성

필자의 경우 test_data.yaml 파일로 작성 했습니다.

해당 파일의 위치는 ./yolov7/data/test_model.yaml에 위치합니다.

# test_data.yaml

# train 이미지 폴더와 valid 폴더 경로 설정
train: /content/drive/MyDrive/ColabNotebooks/test_dataset/train/images
val: /content/drive/MyDrive/ColabNotebooks/test_dataset/train/images

# parameters (클래스 갯수)
nc : 1

# class names
names : ['Jellyfish']

 

두 번째로 yolov7/cfg/training/ 경로에 위치한 yolo7x.yaml 파일의 내용을

전부 복사하여 yolo7x-custom.yaml 파일을 생성해 붙여 넣어줍니다.

자신의 훈련 데이터 클래스 갯수에 맞게 nc 부분만 우선 변경해줍니다.

(필요시 해당 사항을 맞춰 수정하시면 됩니다)

 

전 따로 하이퍼 파라미터 최적화를 하지 않았기 때문에 hyp.scrath.custom.yaml 파일을

그대로 하이퍼 파라미터 모델 파일로 사용했습니다.

 

이제 아래의 코드를 통해 학습을 시작합니다.

!python train.py \
--device 0 \
--batch 16 \	# batch size
--epochs 40 \	# epochs size 
--data ./data/test_data.yaml \
--weights ./yolov7x.pt \
--cfg ./cfg/training/yolov7x-custom.yaml \
--hyp ./data/hyp.scratch.custom.yaml \
--project runs/train/test

약 25분 정도 소요된 것 같습니다.

 

 

PR_curve 파일과 results.png 파일을 통해

훈련 결과를 한 눈에 파악할 수 있습니다.

(좌) 학습 정답 라벨 데이터    (우) 예측 결과 데이터

 

Google Colab 무료 제공 GPU의 경우 사용량 제한이 있으니

참고하시길 바랍니다.