[윈도우] Darknet으로 YOLO 커스텀 모델 학습시키기

윈도우

[윈도우] Darknet으로 YOLO 커스텀 모델 학습시키기

IntegerString 2020. 8. 15. 13:43

본인은 로컬에서 학습시키다가 너무 오래 걸려서 colab으로 학습시켰지만 방법은 같음

classes: 학습시킬 클래스 갯수

train: 학습용 이미지들의 경로가 적혀있는 텍스트 파일

valid: 검증용 이미지들의 경로가 적혀있는 텍스트 파일 (본인은 생략)

(이미지 폴더에는 각 이미지마다 YOLO 라벨 텍스트가 존재해야함)

names: 클래스 이름이 적힌 파일

backup: 가중치 파일들이 저장될 폴더

본인은 colab을 통해 학습했기 때문에 이미지들의 경로가 아래와 같음

아래 파이썬 코드를 이용하면 custom.txt를 쉽게 생성할 수 있음

import glob


# 대상 폴더
image_directory = "C:/Users/asdfr/Pictures/Saved Pictures/pose/"

# 확장자
extension = "*.jpg"

# 텍스트 파일이 저장될 경로
save_at = "C:/Users/asdfr/darknet/build/darknet/x64/pose/custom.txt"

# 대상 폴더에서 지정한 확장자를 가진 파일들의 경로를 리스트화
files = sorted(glob.glob(image_directory + extension))

# 파일들의 경로를 텍스트 파일에 추가 및 출력
for file in files:
    f = open(save_at, 'a')
    f.write(file + "\n")
    print(file)

cfg파일은 darknet에서 제공하는 yolov4-custom.cfg를 복사 후 수정하여 사용함

batch = 64

subdivisions = 16 (GPU 메모리에 따라 32, 64로 수정 가능)

width = 416

height = 416 (너비, 높이 또한 GPU 메모리에 따라 32의 배수로 수정 가능 ex. 224, 320, 608)

max_batches = 클래스 갯수 * 2000

steps = max_batches * 0.8, max_batches * 0.9

나머지는 건드리지 않음