[DeepSeg] YOLOv5 설치 및 MNIST 데이터셋으로 학습

이전 포스트에서 변환한 MNIST 데이터셋으로 YOLOv5 모델을 학습시킨다. 
[DeepSeg] MNIST 데이터셋 YOLO 형식으로 변환 — zerogod 코코딩딩

[DeepSeg] MNIST 데이터셋 YOLO 형식으로 변환

 

1. YOLOv5 설치

git clone https://github.com/ultralytics/yolov5.git
cd yolov5
pip install -r requirements.txt

프로젝트 디렉터리에서 YOLOv5를 설치한다. models 폴더를 생성하고, 다음 명령어를 입력해 가중치 파일을 다운로드한다.

 

wget -O models/yolov5s.pt https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5s.pt

다운로드 후, 다음과 같이 디렉터리를 구성해둔다.

 

data
├── mnist_data
│   └── MNIST
│       └── raw
└── yolo_data
    ├── images
    ├── labels
    ├── train
    └── val
    
models
└── yolov5s.pt
scripts
├── evaluate_yolo.py
├── export_model.py
├── split_data.py
└── train_yolo.py
src
├── main.cpp
└── main.py

yolo_data 폴더에 yolo_mnist.yaml을 생성한다.

 

train: ../data/yolo_data/train/images  # 학습 데이터 이미지 경로
val: ../data/yolo_data/val/images      # 검증 데이터 이미지 경로
test: ../data/yolo_data/test/images    # 테스트 데이터 이미지 경로

nc: 10  # MNIST 숫자는 0~9, 총 10개의 클래스
names: ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']

 

2. YOLOv5 모델 학습

아래 파이썬 코드들을 생성한다.

split_data.py

import os
import shutil
import random

# 데이터 경로 설정
ROOT_DIR = "../data/yolo_data"
IMAGE_DIR = os.path.join(ROOT_DIR, "images")
LABEL_DIR = os.path.join(ROOT_DIR, "labels")
TRAIN_DIR = os.path.join(ROOT_DIR, "train")
VAL_DIR = os.path.join(ROOT_DIR, "val")

# train/val 디렉터리 생성
for d in [TRAIN_DIR, VAL_DIR]:
    os.makedirs(os.path.join(d, "images"), exist_ok=True)
    os.makedirs(os.path.join(d, "labels"), exist_ok=True)

# 이미지 파일 리스트 로드
random.seed(42)  # 항상 같은 데이터 분할 유지
image_files = sorted(os.listdir(IMAGE_DIR))
random.shuffle(image_files)

# 80%:20% 비율로 분할
split_idx = int(len(image_files) * 0.8)
train_files = image_files[:split_idx]
val_files = image_files[split_idx:]

# train/val 디렉터리로 데이터 복사
for files, folder in [(train_files, TRAIN_DIR), (val_files, VAL_DIR)]:
    for file in files:
        src_img = os.path.join(IMAGE_DIR, file)
        dst_img = os.path.join(folder, "images", file)
        
        src_label = os.path.join(LABEL_DIR, file.replace(".jpg", ".txt"))
        dst_label = os.path.join(folder, "labels", file.replace(".jpg", ".txt"))

        # 파일이 존재하면 덮어쓰지 않고 삭제 후 복사
        if os.path.exists(dst_img):
            os.remove(dst_img)
        if os.path.exists(dst_label):
            os.remove(dst_label)

        shutil.copy(src_img, dst_img)
        shutil.copy(src_label, dst_label)

print("Dataset successfully split into train and val sets.")

split_data.py는 mnist 데이터를 70% 20% 10%로 분할해 각각 학습/검증/평가 데이터로 사용한다. 

train_yolo.py

 

import os
import sys
import shutil
import subprocess

BASE_DIR = os.path.dirname(os.path.abspath(__file__))

# 로컬 YOLOv5 경로
YOLOV5_DIR = os.path.abspath(os.path.join(BASE_DIR, "../yolov5"))
if not os.path.exists(YOLOV5_DIR):
    raise FileNotFoundError(f"YOLOv5 directory not found at: {YOLOV5_DIR}")

# 모델/체크포인트 관련 경로
weights_path = os.path.abspath(os.path.join(BASE_DIR, "../models/yolov5s.pt"))
output_dir = os.path.abspath(os.path.join(BASE_DIR, "../runs/train/exp"))
final_model_path = os.path.abspath(os.path.join(BASE_DIR, "../models/mnist.pt"))

if not os.path.exists(weights_path):
    raise FileNotFoundError(f"Pretrained YOLOv5 model does not exist: {weights_path}")

# 데이터 yaml
data_yaml = os.path.abspath(os.path.join(BASE_DIR, "../data/yolo_data/yolo_mnist.yaml"))
project_dir = os.path.abspath(os.path.join(BASE_DIR, "../runs/train"))

# 체크포인트 경로
resume_checkpoint = os.path.join(output_dir, "weights", "last.pt")
resume_flag = os.path.exists(resume_checkpoint) and os.path.isfile(resume_checkpoint)
if resume_flag:
    print(f"Checkpoint detected: {resume_checkpoint}. Resuming training.")
else:
    print("No valid checkpoint detected. Starting training from scratch.")

# base train command
train_command = [
    sys.executable,
    os.path.join(YOLOV5_DIR, "train.py"),
    "--img", "640",
    "--batch", "8",
    "--epochs", "20",
    "--data", data_yaml,
    "--device", "0",
    "--project", project_dir,
    "--name", "exp",
]

# resume 플래그 처리
if resume_flag:
    # resume 시, last.pt 경로를 직접 넘김
    # => YOLOv5가 이 경로를 정확히 체크포인트로 인식하여 기존 상태를 불러옴
    train_command.append(f"--resume")
    train_command.append(os.path.abspath(resume_checkpoint))
else:
    # 처음 학습 시, 사전 학습 weights를 지정
    train_command.extend(["--weights", weights_path])

print("\nStarting YOLOv5 training...\n" + "=" * 50)
print("DEBUG: train_command =", train_command)
sys.stdout.flush()

# YOLOv5 디렉터리를 working directory로
process = subprocess.Popen(
    train_command,
    stdout=subprocess.PIPE,
    stderr=subprocess.STDOUT,
    text=True,
    bufsize=1,
    cwd=YOLOV5_DIR
)

output_lines = []
for line in iter(process.stdout.readline, ''):
    sys.stdout.write(line)
    sys.stdout.flush()
    output_lines.append(line)

process.wait()

if process.returncode != 0:
    full_output = "".join(output_lines)
    print(f"\nError during training. Full log output:\n{full_output}\n")
    sys.exit(1)

print("\nTraining completed successfully.")

# best.pt -> 최종 모델 저장
best_model_path = os.path.join(output_dir, "weights", "best.pt")
if os.path.exists(best_model_path):
    shutil.copy(best_model_path, final_model_path)
    print(f"\nTraining completed. Model saved as {final_model_path}\n")
else:
    print("\nWarning: best.pt not found. Please check if training completed successfully.\n")

학습 과정에서 에러가 발생할 경우, 에러 확인 및 수정하고 다시 학습을 시작하면 자동으로 last.pt를 불러와 학습을 이어서 재개할 수 있다. 에러가 나지 않더라도 이를 이용해 epoch를 늘려가면서 학습 정도를 확인할 수 있다. 유의할 점은 resume에서는 epoch를 늘릴수만 있고, 체크포인트의 epoch보다 작은 epoch로 재설정하는 것은 불가능하다.

 

 

3. 모델 평가 및 onnx 변환

evaluate_yolo.py

import os
import sys
import subprocess

BASE_DIR = os.path.dirname(os.path.abspath(__file__))

# YOLOv5 디렉터리
YOLOV5_DIR = os.path.abspath(os.path.join(BASE_DIR, "../yolov5"))
if not os.path.exists(YOLOV5_DIR):
    raise FileNotFoundError(f"YOLOv5 directory not found at: {YOLOV5_DIR}")

# 모델(학습 완료된) 경로
model_path = os.path.abspath(os.path.join(BASE_DIR, "../models/mnist.pt"))
if not os.path.exists(model_path):
    raise FileNotFoundError(f"Trained model does not exist: {model_path}")

# data yaml
data_yaml = os.path.abspath(os.path.join(BASE_DIR, "../data/yolo_data/yolo_mnist.yaml"))
if not os.path.exists(data_yaml):
    raise FileNotFoundError(f"Data config does not exist: {data_yaml}")

eval_results_path = os.path.abspath(os.path.join(BASE_DIR, "../models/evaluation_results.txt"))

# YOLOv5 val.py 실행 명령어
val_command = [
    sys.executable,               # 파이썬 실행기
    "val.py",                     # yolov5/val.py 스크립트
    "--data", data_yaml,
    "--weights", model_path,
    "--batch", "8",
    "--img", "640",
    "--task", "test"              # optional: val/test
]

print("\nStarting YOLOv5 evaluation...\n" + "=" * 50)

# YOLOv5 디렉터리에서 val.py 실행
process = subprocess.Popen(
    val_command,
    stdout=subprocess.PIPE,
    stderr=subprocess.STDOUT,
    text=True,
    bufsize=1,
    cwd=YOLOV5_DIR  # cwd를 YOLOv5 폴더로 설정
)

output_lines = []
for line in iter(process.stdout.readline, ''):
    sys.stdout.write(line)
    sys.stdout.flush()
    output_lines.append(line)

process.wait()

if process.returncode != 0:
    full_output = "".join(output_lines)
    print(f"\nError during evaluation. Full log output:\n{full_output}\n")
    sys.exit(1)

print("\nEvaluation completed successfully.")
print(f"Results are saved in YOLOv5 default path (runs/val/*).")

# 필요하면 결과를 저장(추가 기능)
with open(eval_results_path, "w") as f:
    f.write("".join(output_lines))

print(f"\nEvaluation log saved to {eval_results_path}.\n")

export_model.py

import os
from yolov5.export import run

model_path = "../models/mnist.pt"
onnx_output = "../models/mnist.onnx"

if not os.path.exists(model_path):
    raise FileNotFoundError(f"Trained model does not exist: {model_path}")

# ONNX 변환 실행
run(weights=model_path, simplify=True, device="cpu")

# 변환 확인
if os.path.exists(onnx_output):
    print(f"ONNX conversion completed: {onnx_output}")

else:
    print("ONNX conversion failed.")

 

위 파이썬 코드들은 각각 데이터셋 분할(학습, 검증용), 학습, 평가, 변환을 수행한다. 다음 main.py에 script 리스트를 작성하면 순서대로 실행하여 결과가 저장된다. 

main.py

import os
import subprocess

BASE_DIR = os.path.dirname(os.path.abspath(__file__))  # 현재 main.py가 위치한 디렉터리
SCRIPTS_DIR = os.path.join(BASE_DIR, "../scripts")      # scripts 디렉터리 절대 경로

# 실행할 script 리스트
scripts = [
    "split_data.py",      # 데이터 분할
    "train_yolo.py",      # YOLOv5 학습
    "evaluate_yolo.py",   # 모델 평가
    "export_model.py"     # ONNX 변환
]

# 모든 script 실행
for script in scripts:
    script_path = os.path.join(SCRIPTS_DIR, script)
    print(f"\nRunning: {script_path}\n" + "=" * 50)
    
    # stderr와 stdout을 합쳐서 처리
    process = subprocess.Popen(
        ["python3", script_path],
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True
    )
    
    # 표준 출력을 실시간으로 출력
    for line in process.stdout:
        print(line, end="")

    process.wait()

    if process.returncode != 0:
        print(f"\nError occurred while running: {script_path}\n")
        break  # 오류 발생 시 중단

print("\nFull pipeline execution completed.\n")

만약 수행 중 특정 단계에서 에러가 발생한다면, 에러가 발생한 script 이전 scipt까지 주석 처리한 후 다시 실행하면 된다.

 

학습 로그. 학습이 완료되는 데에는 배치사이즈에 따라 다르지만 4시간 가량 걸렸다.

좌 학습 우 검증

학습 완료

 

학습이 완료된 이후 results.csv

               epoch,      train/box_loss,      train/obj_loss,      train/cls_loss,   metrics/precision,      metrics/recall,     metrics/mAP_0.5,metrics/mAP_0.5:0.95,        val/box_loss,        val/obj_loss,        val/cls_loss,               x/lr0,               x/lr1,               x/lr2
                   0,             0.02454,            0.014166,            0.036785,             0.89466,             0.87568,             0.94293,             0.82035,           0.0045999,           0.0016709,           0.0050996,            0.070005,           0.0033328,           0.0033328
                   1,            0.017273,            0.010344,            0.022659,             0.96396,             0.96127,             0.98822,             0.88367,           0.0034047,           0.0014933,           0.0019635,            0.039675,           0.0063361,           0.0063361
                   2,            0.015615,            0.009903,            0.020341,             0.96674,              0.9685,             0.98924,             0.89413,            0.003119,           0.0014988,           0.0016345,           0.0090151,           0.0090095,           0.0090095
                   3,            0.013721,           0.0094438,            0.018682,             0.97678,             0.96806,             0.99035,             0.93666,           0.0023929,           0.0012859,           0.0012933,            0.008515,            0.008515,            0.008515
                   4,            0.012365,           0.0089515,             0.01723,             0.98325,             0.97837,             0.99288,             0.95134,           0.0020823,           0.0012355,          0.00093813,            0.008515,            0.008515,            0.008515
                   5,            0.011394,           0.0085409,              0.0163,             0.98556,             0.97894,             0.99282,             0.96209,           0.0018873,            0.001175,          0.00089005,             0.00802,             0.00802,             0.00802
                   6,            0.010717,           0.0083432,            0.015516,             0.98323,             0.98336,             0.99353,             0.96325,             0.00181,           0.0011492,          0.00079767,            0.007525,            0.007525,            0.007525
                   7,             0.01029,           0.0081321,            0.015228,             0.98573,             0.98557,             0.99358,             0.97173,           0.0016585,           0.0011094,          0.00069697,             0.00703,             0.00703,             0.00703
                   8,            0.010062,            0.007963,            0.014919,             0.98976,             0.98494,             0.99408,             0.97524,            0.001573,           0.0010675,           0.0006226,            0.006535,            0.006535,            0.006535
                   9,           0.0095195,           0.0078317,            0.014369,             0.98743,             0.98789,             0.99393,              0.9794,           0.0014574,           0.0010328,          0.00057654,             0.00604,             0.00604,             0.00604
                  10,           0.0093978,           0.0076915,            0.014249,             0.98931,               0.987,             0.99408,             0.98079,            0.001431,           0.0010094,          0.00053787,            0.005545,            0.005545,            0.005545
                  11,           0.0089916,           0.0075712,             0.01383,             0.98897,             0.98725,             0.99398,             0.98022,           0.0014786,          0.00099321,          0.00054507,             0.00505,             0.00505,             0.00505
                  12,           0.0086958,           0.0074411,            0.013503,             0.98994,             0.98758,             0.99396,             0.98422,           0.0013222,          0.00096922,          0.00049742,            0.004555,            0.004555,            0.004555
                  13,           0.0084506,           0.0072285,            0.013165,              0.9888,             0.98937,             0.99402,             0.98537,           0.0012766,          0.00093808,          0.00048256,             0.00406,             0.00406,             0.00406
                  14,           0.0082058,           0.0071584,            0.012965,             0.98984,             0.98838,              0.9941,             0.98619,           0.0012481,          0.00093727,          0.00046001,            0.003565,            0.003565,            0.003565
                  15,           0.0079196,           0.0070333,            0.012799,             0.99062,             0.98759,             0.99412,             0.98629,           0.0012415,          0.00093543,          0.00045288,             0.00307,             0.00307,             0.00307
                  16,           0.0076483,           0.0069446,            0.012644,             0.98701,             0.99137,             0.99414,             0.98748,           0.0012169,          0.00091489,           0.0004456,            0.002575,            0.002575,            0.002575
                  17,           0.0074007,            0.006782,            0.012214,             0.98926,             0.99045,             0.99419,             0.98836,           0.0011751,          0.00090387,          0.00042609,             0.00208,             0.00208,             0.00208
                  18,           0.0070743,           0.0066269,            0.011996,             0.98989,             0.99025,             0.99421,             0.98853,           0.0011525,          0.00089901,           0.0004115,            0.001585,            0.001585,            0.001585
                  19,           0.0067405,           0.0065368,            0.011843,             0.98982,             0.99058,             0.99423,             0.98899,           0.0011364,          0.00089191,          0.00040374,             0.00109,             0.00109,             0.00109

 

모델 평가 및 변환 로그

 

이제 mnist 데이터셋을 이용해 숫자 객체를 추적하는 mnist.onnx모델이 완성되었다. 

 

zero0zerogod/YOLOv5-MNIST