YOLOv5 커스텀 학습 튜토리얼

파일 받아오기

from urllib import request 
url = "https://public.roboflow.com/ds/FbPSJa257p?key=2mKwetv7pn"
savename = "roboflow.zip" 
request.urlretrieve(url,savename) 
print("저장되었습니다.")

 

---

압축파일 풀기

import zipfile
         
roboflow_zip = zipfile.ZipFile('/home/soon5770/roboflow.zip')
roboflow_zip.extractall('/home/soon5770/')
 
roboflow_zip.close()

---

YOLOv5 git clone

# YOLOv5 github clone
%cd /home/soon5770
!git clone https://github.com/ultralytics/yolov5.git

---

YOLOv5 패키지 설치하기

# YOLOv5 를 위한 패키지 설치하기
%cd /home/soon5770/yolov5/
!pip install -r requirements.txt

---

yaml 파일 확인해보기

%cat /home/soon5770/dataset/data.yaml
# training 경로, validation 경로, nc : 클래스의 개수, names : 클래스의 이름

train 데이터와 validation 데이터의 경로를 바꾸어 주어야 한다.

---

%cd / 
from glob import glob # 이미지 싹다 긁어오기
img_list = glob('/home/soon5770/dataset/export/images/*.jpg')
print(len(img_list))

---

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))

---

training image, validation image 경로를 txt 파일로 다 저장하기

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

train.txt에다가 'w' 쓰는걸로 하고 

리스트로 되어 있으니 join으로 붙여주고 뒤에 엔터하나 쳐주고

val.txt도 똑같이 해준다.

그러면 dataset 폴더에 train.txt랑 val.txt 파일이 생긴다.

 

이제 training set과 validation set을 모두 나눠놨다.

---

이제 yaml파일을 수정하자.

여기에 train에다가 train.txt를 넣고

val에다가 val.txt를 넣자.

 

import yaml

with open('/home/soon5770/dataset/data.yaml','r')as f:
    data = yaml.load(f)
    
print(data)

data['train'] = '/home/soon5770/dataset/train.txt'
data['val'] = '/home/soon5770/dataset/val.txt'

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

print(data)

바뀐 것을 확인할 수 있다.

---

이제 학습시켜보자.

train.py 트레이닝을 시키는 파일을 실행시키자.

!python /home/soon5770/yolov5/train.py --img 416 --batch 16 --epochs 50 --data /home/soon5770/dataset/data.yaml --cfg /home/soon5770/yolov5/models/yolov5s.yaml --weights yolov5s.pt --name gun_yolov5s_results

--img 416 : 416사이즈

--batch 16

--epochs 50

--data /home/soon5770/dataset/data.yaml

--cfg /home/soon5770/yolov5s.yaml  : cfg 파일 - 모델의 구조 , 큰거 l, 중간거 m, 작은거 s

--weights yolov5s.pt : pretrained 된 모델을 쓸 것이다.

--name gun_yolov5s_results 

 

 

아래와 같은 에러가 떴고

RuntimeError: DataLoader worker (pid 1105) is killed by signal: Bus error. It is possible that dataloader's workers are out of shared memory. Please try to raise your shared memory limit.

 

!df -h 를 통해 메모리를 확인해보았을 때

shm 메모리가 64m 밖에 되지 않는 것을 확인했다.