YOLO v4 custom데이터 훈련하기

이전에 YOLO v4를 설치해서 example실행 까지 잘 마쳤다면, 이제 본인의 데이터를 훈련시키고 응용하는 것까지 알아보도록 하자. YOLO v4 설치에 관해서는 이전글 참조.

2020/05/19 - [Computer Vision/Object detection] - Ubuntu 18.04 에서 Yolo v4 훈련하기

Ubuntu 18.04 에서 Yolo v4 설치하기

 

이제 데이터 훈련을 본격적으로 알아보도록 하자.

대부분의 내용은 YOLO v3의 훈련 내용과 같으므로 YOLO v3 글에서 긁어왔다. 참조로 링크를 달 수도 잇었지만,

굳이 두개의 글을 띄워 놓고 보기는 싫을테니 .. 

Custom 데이터 Yolo로 훈련시키기

custom데이터를 훈련시키기 위해서는 가장 먼저 cfg를 수정해주어야 한다. 제공된 cfg파일이 coco데이터셋을 기반으로 되어 있기 때문에 사용하고자 하는 데이터셋에 맞춰줄 필요가 있다. 

cfg수정

이전에 설치한 darknet 디렉토리로 이동하여 cfg파일을 수정한다.

gedit cfg/yolov4.cfg

cfg파일을 열면 첫 부분을 수정해준다 width와 height가 608 , 608로 되어있는데 이렇게 진행하면 gpu 메모리를 상당히 차지하므로 416,416로 수정해준다. 만약 본인의 메모리가 부족하다면 더 줄여도 상관없다. cuda memory에러시 batch size를 수정

여기서 이번에는 저번과 다르게 AlexeyAB github에 적힌 내용대로 진행하도록 하겠다.

subdivisions = 16 으로 수정
max_batches는 본인 데이터의 클래스 갯수 * 2000 을 적어주면 된다. 
필자의 클래스는 15개이므로 15*2000 = 30000 을 입력
steps는 max_batches의 80%의 수치와 90%의 수치를 작성해주면 된다.

(단, 이 것은 최소 추천 iter이므로 만약 본인이 훈련이 불충분하다고 느끼거나 iter를 더 진행하고 싶다면 max_batches에 원하는 iter의 수를 입력하고, step도 그에 맞게 수정해주면 된다. ex) max_batches = 100000 / steps = 80000,90000)

[net]
batch=64
subdivisions=16
# Training
#width=512
#height=512
width=416
height=416
channels=3
momentum=0.949
decay=0.0005
angle=0
saturation = 1.5
exposure = 1.5
hue=.1

learning_rate=0.00261
burn_in=1000
max_batches = 30000
policy=steps
steps=24000,27000
scales=.1,.1

 

첫 부분을 수정하였으면 Ctrl+f를 이용하여 yolo부분을 검색해서 수정해주면 된다. yolo 부분에 있는 classes와 yolo위쪽에 있는 filter의 크기를 수정해주면 된다. 

classes는 훈련시키고자 하는 데이터셋의 class갯수, filter의 갯수는 (classes+5) * 3 으로 계산하여 적어준다. 필자의 클래스는 15개이고, filter의 갯수는 (15+5)*3 = 60개 이다.

( [Gaussian_yolo] layer를 이용한다면 filters = (classes+9)*3 ) 

[convolutional]
size=1
stride=1
pad=1
filters=60
activation=linear


[yolo]
mask = 0,1,2
anchors = 12, 16, 19, 36, 40, 28, 36, 75, 76, 55, 72, 146, 142, 110, 192, 243, 459, 401
classes=15
num=9
jitter=.3
ignore_thresh = .7
truth_thresh = 1
scale_x_y = 1.2
iou_thresh=0.213
cls_normalizer=1.0
iou_normalizer=0.07
iou_loss=ciou
nms_kind=greedynms
beta_nms=0.6

 

yolo 부분은 총 3개가 있으므로 꼭 cfg파일에서 3개를 모두 찾아서 수정해줘야한다.

*anchor size의 경우 kmeans를 이용하여 본인의 데이터셋에 맞게끔 수정해줄 수 있다. 하지만 훈련을 충분히 시킬 수 있다면, 그냥 default로 설정되어 있는 coco데이터셋의 anchor size를 이용해도 된다.

 

훈련 파일 관리하기

cfg파일까지 모두 수정하였다면 본인의 커스텀 데이터의 맞게끔 names,data,train.txt파일을 수정해준다.

먼저 train.txt파일을 확인해보면,

x64/Release/data/img2/001_aaaa.jpg
x64/Release/data/img2/001_aaac.jpg
x64/Release/data/img2/001_aaad.jpg
x64/Release/data/img2/001_aaae.jpg
x64/Release/data/img2/001_aaaf.jpg
x64/Release/data/img2/001_aaag.jpg
x64/Release/data/img2/001_aaah.jpg
x64/Release/data/img2/001_aaai.jpg

 

이런식으로 되어있는 것을 알 수 있다. Yolo mark에서 진행했기 때문에 상대경로로 적혀 있는데 darknet폴더로 데이터셋을 모두 옮겨와서 사용해도 되지만 필자는 Yolo mark폴더에 접근하기 위해서 파일내용을 절대경로로 바꿔주었다.

jupyter notebook을 활용하여 train.txt를 절대경로로 변환해 주었다. 

txt = open('/media/gopiz/extra/vision_project/darknet/custom/train.txt','r')
f = open('/media/gopiz/extra/vision_project/darknet/custom/train2.txt','w')
while True :
    line = txt.readline()
    if not line:
        break
    f.write('/media/gopiz/extra/vision_project/Yolo_mark/'+line)
        
txt.close()
f.close()

 

Output으로 나오는 train2.txt파일을 보면 절대경로로 변환된 것을 확인할 수 있다.

/media/gopiz/extra/vision_project/Yolo_mark/x64/Release/data/img2/001_aaaa.jpg
/media/gopiz/extra/vision_project/Yolo_mark/x64/Release/data/img2/001_aaac.jpg
/media/gopiz/extra/vision_project/Yolo_mark/x64/Release/data/img2/001_aaad.jpg
/media/gopiz/extra/vision_project/Yolo_mark/x64/Release/data/img2/001_aaae.jpg
/media/gopiz/extra/vision_project/Yolo_mark/x64/Release/data/img2/001_aaaf.jpg
/media/gopiz/extra/vision_project/Yolo_mark/x64/Release/data/img2/001_aaag.jpg
/media/gopiz/extra/vision_project/Yolo_mark/x64/Release/data/img2/001_aaah.jpg
/media/gopiz/extra/vision_project/Yolo_mark/x64/Release/data/img2/001_aaai.jpg

 

다음은 훈련을 위해 train set과 validation set을 나눠보자. 필자의 데이터셋은 동영상에서 추출한 이미지이기 때문에 연속성이 강해서 train set과 validation set을 의미있게 나누기 위해 데이터를 섞어 주었다.

역시 jupyter notebook을 통해 한번에 파일을 관리하였다.

import random

txt = open('/media/gopiz/extra/vision_project/darknet/custom/train2.txt','r')
f = open('/media/gopiz/extra/vision_project/darknet/custom/train_suffle.txt','w')

tmp = []

while True :
    line = txt.readline()
    if not line:
        break
        
    tmp.append(line)
    
random.shuffle(tmp)
        
for i in tmp :  
    f.write(i)

txt.close()
f.close()

 

train_suffle.txt를 확인해보면 순서가 섞여있는 것을 확인할 수 있다.

/media/gopiz/extra/vision_project/Yolo_mark/x64/Release/data/img2/001_abgr.jpg
/media/gopiz/extra/vision_project/Yolo_mark/x64/Release/data/img2/001_acqb.jpg
/media/gopiz/extra/vision_project/Yolo_mark/x64/Release/data/img2/001_acgr.jpg
/media/gopiz/extra/vision_project/Yolo_mark/x64/Release/data/img2/001_abqb.jpg
/media/gopiz/extra/vision_project/Yolo_mark/x64/Release/data/img2/001_aanh.jpg
/media/gopiz/extra/vision_project/Yolo_mark/x64/Release/data/img2/001_acur.jpg
/media/gopiz/extra/vision_project/Yolo_mark/x64/Release/data/img2/001_acge.jpg
/media/gopiz/extra/vision_project/Yolo_mark/x64/Release/data/img2/001_abes.jpg

 

이제 train set과 validation set으로 나눠보자.

count = 0
length = 2135 #total line

txt = open('/media/gopiz/extra/vision_project/darknet/custom/train_suffle.txt','r')

i = 0

f = open('/media/gopiz/extra/vision_project/darknet/custom/train.txt','w')
f2 = open('/media/gopiz/extra/vision_project/darknet/custom/validation.txt','w')

while True :
    if i == 0 :
        line = txt.readline()
        if not line :
            break
        count +=1
        if count < int(length/10)*2 :
            f2.write(line)
        else :
            f.write(line)

txt.close()
f.close()
f2.close()

 

모든 작업이 끝났으면 train : validation이 8:2의 비율로 나눠진 것을 확인할 수 있다.

custom.names 파일 생성

yolo는 class의 index로 훈련을 하지만 index에 따른 class의 이름을 같이 표시해줄 수 있게 만들어 졌다. 그러므로 names파일을 넣어줘야 class의 이름을 확인할 수 있다. gedit명령어를 통해 custom.names파일을 생성한다. 

Dough
Tomato
Onion
Pineapple
Mushroom
Sweet_potato
Wedge_potato
Garlic_chip
Black_olive
Bulgogi
Pepperoni
Bacon
Corn
Fire_chicken
SP_cube

 

이런식으로 본인의 데이터셋에 맞는 Class이름을 index순서대로 쭉 작성해주면 된다.

custom.data 파일 생성

gedit명령어를 통해 custom.data파일을 생성한다. data파일은 yolo 훈련의 좌표가 되는 파일이라고 볼 수 있다.

클래스의 갯수, train set 경로, validation set경로, names파일의 경로, weight파일이 출력될 경로를 입력하여 준다.

classes= 15
train  = custom/train.txt
valid  = custom/validation.txt
names = custom/custom.names
backup = backup/

 

 

Pretrain 모델 다운로드

darknet에서 제공하는 pretrain 모델을 이용하면 좀 더 빠르고 정확한 학습을 진행할 수 있다. AlexeyAB github에서 제공하는 yolov4의 pretrain모델을 다운로드 받아보자

drive.google.com/open?id=1JKF-bdIklxOOVy-2Cr5qdvjgGpmGfcbp

yolov4.conv.137

다운로드가 완료되었다면 본격적으로 훈련을 진행한다.

YOLO v4 훈련하기

모든 준비가 끝났다면 훈련 명령어를 통해 훈련을 진행해보자. 하지만 yolo는 훈련 진행중 과정을 확인하기 위해 AlexeyAB의 darknet에서 제공하는 map태그를 통해 훈련 과정을 확인하면서 훈련을 진행할 수 있다.

./darknet detector train custom/topping.data custom/topping_yolov4.cfg yolov4.conv.137 -map

 

훈련을 진행하게 되면

 

yolo v4 훈련 터미널

 

 

이런 터미널 화면과 함께 Map와 loss의 그래프를 그려주는 화면이 나타난다.

 

map 와 loss 그래프

 

훈련이 완료되고 나면 1000iter 마다 생성된 가중치와 함께 훈련 iter중 마지막 가중치를 담고 있는 last weight 파일 훈련 과정중 가장 평가 지표상 높게 나타났던 best weight파일과 함께 훈련 모두 끝난 최종 파일인 final weight 파일이 떨어져 나오게 된다.

가중치 까지 모두 얻었다면 다음 post를 통해 응용하는 법을 알아보자!