YOLOv3 Custom Data Train on Ubuntu 24.04

Instruction

그동안 일과 논문 작성이 겹치면서 포스팅을 미루고 미루다 이제서야 한적해진 수요일 오전에 포스팅을 해본다.. (사실 한건 산더미이지만 OS 포팅을 기다리면서 이해했던 것도 정리할 겸 쓴다.)

실제로 YOLO 훈련 관련한 공부는 24년도 하반기부터 본격적으로 시작했지만 시스템 환경 문제, 각종 패키지 버전 문제, 코드 문제 등이 겹치면서 습득하는 데까지 오래 걸렸다. 머리가 안 좋아 삽질을 많이 한 케이스 같다.

사실 내가 사용한 소스코드는 Google Colab에서 사용하면 바로 동작이 가능한 코드이다. 하지만 PC로(RTX3080, i9-12900) 학습하면 학습 속도가 더 빠르다는 얘기를 듣고 시작하게 되었다.

Darknet Struct & NVIDIA Library 

우선 지금까지 했던 흔히 우리가 말하는 "삽질"이 어떤 것들이 있었는지 돌이켜보면 앞서 말한 시스템, 패키지 버전 문제가 가장 크다. 이러한 문제들을 알려면 YOLO의 구조까진 아니더라도 YOLO의 신경망을 구성하는 프레임워크인 Darknet을 이해하는 것이 조금은 도움이 된다.

Darknet은 YOLO의 신경망을 구성하는 프레임워크로서 백본 역할을 한다. (다른 말로 YOLO는 Darknet을 사용하여 구현된 객체 탐지 모델이다.) Darknet은 C와 CUDA로 작성되어 있기 때문에 컴파일을 통해 기기에 최적화된 실행 파일을 생성해야 비로소 내 PC에서 사용할 수 있다. 이 말인 즉슨 내 PC에 있는 그래픽 카드(CUDA)를 사용하여 추론, 학습을 진행하려면 Darknet을 컴파일해야 한다는 것이다.

이를 구현하기 위해 아래와 같이 크게 세 가지 NVIDIA 드라이버, 라이브러리가 필요하다.

• NVIDIA Driver: 운영 체제와 GPU 하드웨어 간의 통신을 가능하게 하는 기본 소프트웨어로 GPU 기능 활용에 필수 요소이다.
• CUDA Toolkit: Compute Unified Device Architecture의 약자로 NVIDIA가 개발한 병렬 컴퓨팅 플랫폼 및 API 모델이다. CUDA 런타임 및 드라이버 라이브러리, 컴파일러와 개발 도구, 수학 라이브러리에 사
• cuDNN: CUDA Deep Neural Network library의 약자로 딥 러닝 연산을 위해 GPU 가속화된 기본 함수들을 제공하는 라이브러리이다. 딥 러닝 프레임워크의 성능을 크게 향상시켜주는 라이브러리이다.

이 세 가지 구성요소들을 필수적으로 PC에 설치해 줘야 하는데 NVIDIA에서 친절하게 제공해 준다. 하지만 나는 한 번 버전을 잘못 설치하는 바람에 지우고 설치하고를 여러 번 반복했다. 

 

NVIDIA Driver는 자신의 GPU를 지원하는 드라이버를 설치하면 된다. 매우 쉬움 명령어 두 줄 정도니 찾아보면 된다.

 

CUDA Toolkit이 비교적 까다로운데 아래 따라하기를 만들어 놓았다.

자신의 GPU가 지원하는 것을 확인하면 된다.↓

https://ko.wikipedia.org/wiki/CUDA

CUDA - 위키백과, 우리 모두의 백과사전

 

내 그래픽 카드의 맞는 버전은 다음과 같다.

ex.) RTX 3080

 

아래 그림을 통해 내 그래픽 카드를 지원하는 CUDA 버전을 찾으면 된다.

나는 Ampere, 8.6 이상의 버전은 모두 지원하기 때문에 CUDA 11.1 이후의 버전에서는 모두 지원하는 것을 확인할 수 있다.

Ampere, 8.6 이상 버전 확인

 

마지막으로 cuDNN 설치는 아래 공식 링크에서 가능하고, 설치된 CUDA 버전을 지원하는지 확인하고 설치하면 된다.(이것도 비교적 간단)

https://developer.nvidia.com/cudnn-downloads

cuDNN 9.7.0 Downloads

 

OpenCV 설치는 센스(아래 C++ 컴파일로 설치하는 것보단 명령어로 딸깍해서 설치하는 게 잘 되어 있는 것 같으니 그런 것들을 설치하는 것을 추천!!)

https://forthe-future.tistory.com/135

Install OpenCV Lite 4.5 on Raspberry Pi

OpenCV를 설치하는 것까지 끝났다면 학습에 필요한 패키지, 드라이버 준비는 모두 끝났다.

 

Make

이제 darknet을 다운로드하고, Makefile을 수정한 뒤 make를 진행하면 된다.

자세한 명령어 설명은 생략

https://github.com/AlexeyAB/darknet

GitHub - AlexeyAB/darknet: YOLOv4 / Scaled-YOLOv4 / YOLO - Neural Networks for Object Detection (Windows and Linux version of Da

# darkent download
git clone https://github.com/AlexeyAB/darknet.git

# Makefile edit
cd darkent
sed -i 's/OPENCV=0/OPENCV=1/' Makefile
sed -i 's/GPU=0/GPU=1/' Makefile
sed -i 's/CUDNN=0/CUDNN=1/' Makefile
make clean
make
# It takes a 1~2min

 

Config Files Setting

다음은 학습시킬 때 필요한 cfg 파일인데 다음 코드는 기존 coco 데이터셋(80가지 물체 분류)에서 내가 원하는 객체(1가지 물체 분류)로 80 → 1로 수정하는 부분이다. 코드 설명은 안 해도 충분히 알아볼 수 있을 것이다.

cp cfg/yolov3.cfg cfg/yolov3_custom_training.cfg

sed -i 's/batch=1/batch=2/' cfg/yolov3_custom_training.cfg
sed -i 's/subdivisions=1/subdivisions=8/' cfg/yolov3_custom_training.cfg
sed -i 's/max_batches = 500200/max_batches = 6000/' cfg/yolov3_custom_training.cfg
sed -i '610 s@classes=80@classes=1@' cfg/yolov3_custom_training.cfg
sed -i '696 s@classes=80@classes=1@' cfg/yolov3_custom_training.cfg
sed -i '783 s@classes=80@classes=1@' cfg/yolov3_custom_training.cfg
sed -i '603 s@filters=255@filters=18@' cfg/yolov3_custom_training.cfg
sed -i '689 s@filters=255@filters=18@' cfg/yolov3_custom_training.cfg
sed -i '776 s@filters=255@filters=18@' cfg/yolov3_custom_training.cfg

# For training
sed -i '3 s@batch=1@# batch=1@' cfg/yolov3_custom_training.cfg
sed -i '4 s@subdivisions=1@# subdivisions=1@' cfg/yolov3_custom_training.cfg
sed -i '6 s@# batch=64@batch=64@' cfg/yolov3_custom_training.cfg
sed -i '7 s@# subdivisions=16@subdivisions=16@' cfg/yolov3_custom_training.cfg

 

Pre-Learning Settings

이 과정부터는 학습할 데이터가 필요하다 하지만 포스팅 순서가 꼬이는 바람에 바로 이후에 학습 데이터 수집에 대해 작성할 예정이다. 다음 파이썬 코드도 어렵지 않으므로 자세한 설명은 생략하겠다.

images_dir = '{your path}/output_frames'

# starting counter for classes
c = 0

# Creating file classes.names from classes.txt

with open(images_dir + '/' + 'classes.names', 'w') as classes, \
     open(images_dir + '/' + 'classes.txt', 'r') as txt:

    for line in txt:
        classes.write(line)

        c += 1
        
# Creating file path_data.data

with open(images_dir + '/' + 'path_data.data', 'w') as content:

    content.write('classes = ' + str(c) + '\n')

    # Location of the train.txt file
    content.write('train = ' + images_dir + '/' + 'train.txt' + '\n')

    # Location of the test.txt file
    content.write('valid = ' + images_dir + '/' + 'test.txt' + '\n')

    # Location of the classes.names file
    content.write('names = ' + images_dir + '/' + 'classes.names' + '\n')

    # Location where to save weights
    content.write('backup = backup')
    
import os

# Creating file train.txt and writing 85% of lines in it

p = [os.path.join(images_dir,f) for f in os.listdir(images_dir) if f.endswith('.jpg')]

os.chdir(images_dir)

# Slicing first 15% of elements for test
p_test = p[:int(len(p) * 0.15)]

# Deleting from initial list first 15% of elements
p = p[int(len(p) * 0.15):]
p = '\n'.join(p)

# Creating file train.txt and writing 85% of lines in it
with open('train.txt', 'w') as train_txt:
    # Going through all elements of the list
    for e in p:
        # Writing current path at the end of the file
        train_txt.write(e)

# Creating file test.txt and writing 15% of lines in it
with open('test.txt', 'w') as test_txt:
    # Going through all elements of the list
    for e in p_test:
        # Writing current path at the end of the file
        test_txt.write(e)

 

Train

이제 학습 시킬 일만 남았다.

# 중간 학습 데이터를 저장하기 위해 백업 폴더 생성
mkdir backup

# yolov3 학습 가중치 값 다운로드
wget https://pjreddie.com/media/files/darknet53.conv.74

# 명령 실행 위치는 darknet, outputs_frames, backup가 있는 프로젝트 가장 상위 폴더
time darknet/darknet detector train output_frames/path_data.data darknet/cfg/yolov3_custom_training.cfg darknet53.conv.74 -dont_show

 

나는 로컬 환경에서 50개 학습 데이터가 30분 정도 소요됐다.

아래 이미지는 학습이 정상적으로 종료되면 나타나는 출력이다.

31분 경과된 모습

학습 과정 그래프

Conclusion

여기까지가 우분투 로컬 환경에서 YOLOv3 Custom Train를 하는 방법을 소개하였다.

포스팅 순서를 데이터 수집부터 했어야 순서가 맞는데.. 중간에 파먹은 부분이 생겨버렸다...

 

다음 포스팅에서는 데이터 수집과 학습된 모델을 평가하는 부분을 포스팅 할 예정이다.