[CV] YOLOv8 + MLFlow For Beginner

🌷 학습 중 정리한 내용으로 문제가 있다면 메일(jnw__@네이버)로 연락 부탁드립니다. :D

Reference

• filot님이 공유해주신 YOLOv8 코드
• Ultralytics Document

00. 개요

Reference
Baek Kyun Shin님의 논문 리뷰 - YOLO(You Only Look Once) 톺아보기

You Only Look Once라는 이름의 YOLO
한 번의 Convolution Network로 이미지의 bounding box와 그 확률을 계산하여 가장 확률이 높은 bounding box를 추론하는 모델

객체 탐지에 좋은 성능을 보이면서도 실시간 탐지 속도도 빠르다.

최근 Ultralytics에서 YOLOv8를 발표했고, backbone과 neck architecture를 강화하여 개선된 성능을 보여준다.

아래 구조는 YOLOv4 논문에서 가져온 것

Docs에 잘 정리된 내용들을 실습해보면서 MLFlow로 간단하게 학습 관리를 해보았고, 그 내용을 기록하고자 한다. :D

테스트 환경

• Linux 18.04
• Ultralytics YOLOv8.0.110
• Python-3.8.16
• torch-2.0.1+cu117 CUDA:0
  • (Tesla V100-SXM2-32GB, 32510MiB)

패키지 설치
pip install ultralytics
pip install opencv-python

01. Simple Test

🍒 빠르고 간단한 사용성

이미 학습된 모델을 가져와서 간단히 테스트 해볼 수 있다.

from ultralytics import YOLO
import cv2

model = YOLO("yolov8n.pt")
# accepts all formats - image/dir/Path/URL/video/PIL/ndarray. 0 for webcam

# from ndarray
im2 = cv2.imread("bbang2.jpg")
results = model.predict(source=im2, save=True, save_txt=True)  # save predictions as labels

아래와 같이 예측 결과가 뜨고 저장된 이미지를 확인할 수 있다.

640x480 8 persons, 7.0ms
Speed: 2.5ms preprocess, 7.0ms inference, 1.5ms postprocess per image at shape (1, 3, 640, 640)

02. Custom Dataset 구성하기

🍒 Custom Dataset 으로 학습/예측을 해보자.

• 이미지와 파일 이름이 같은 텍스트 쌍이 필요
• 텍스트는 객체의 class와 center_x, center_y, width, height로 구성되어야 한다.

🍒 Ultralytics는 yaml 파일에 저장된 경로를 기반으로 하여 데이터를 참조

• 데이터 경로 오류가 나면 아래 파일에서 datasets_dir을 수정하면 됨
  • .config/Ultralytics/settings.yaml

Reference
Ultralytics Document Dataset 구성 방법

YOLO 모델에 사용할 수 있게 데이터 변환

• .txt 데이터는 image 데이터와 이름이 같아야 하고,
  object_class center_x center_y width height 로 쓰여야 한다.

변환 코드

• image_width, image_height 로 나눠 주는것은 bounding box의 크키를 일반화하기 위함

x_min, y_min = float(min(line[5], line[7])), float(min(line[6], line[8]))
x_max, y_max = float(max(line[1], line[3])), float(max(line[2], line[4]))
x, y = float(((x_min + x_max) / 2) / image_width), float(((y_min + y_max) / 2) / image_height)
w, h = abs(x_max - x_min) / image_width, abs(y_max - y_min) / image_height
yolo_labels.append(f"{class_name} {x} {y} {w} {h}")            

02. Yaml 파일 작성 후 Train하기

🍒 YOLO 모델은 yaml 파일을 기준으로 데이터와 분류 클래스를 참조

yaml 파일 정의

yaml_data = {
              "names": classes,
              "nc": len(classes),
              "path": "data/yolo/",
              "train": "train",
              "val": "valid",
              "test": "test"
            }

train

• last.pt 와 best.pt가 저장됨
• last.pt를 불러와 이어 학습하기 가능

# model = YOLO(f"{MODEL}/train/weights/last.pt")
model = YOLO("yolov8x")

results = model.train(
    **opt
    )

predict

model = YOLO("v2/train/weights/best.pt")
test_image_paths = glob("./data/yolo/test/*.png")
for i, image in tqdm(enumerate(get_test_image_paths(test_image_paths)), total=int(len(test_image_paths)/BATCH_SIZE)):
    model.predict(image, imgsz=(1024, 1024), iou=0.2, conf=0.5, save_conf=True, save=False, save_txt=True, project=f"{MODEL}", name="predict",
                  exist_ok=True, device=0, augment=True, verbose=False)
    if i % 5 == 0:
        clear_output(wait=True)

result

03. MLFlow에 학습 결과 저장하기 (Databricks)

🍒 더 좋은 방법이 있다면 꼭 알려주세요 !

참고
Databricks 무료 버전은 아래 버튼으로 가입하기

라이브러리 설치

conda install mlflow 
conda install databricks-cli

계정 설정

• 아래 명령어 입력 후 Username(메일 주소)과 password 입력

databricks configure --host https://community.cloud.databricks.com/

call back function 선언

• 콜백 함수를 지원 문서 보기

def on_fit_epoch_end(trainer):
    if mlflow:
        metrics_dict = {f"{re.sub('[()]', '', k)}": float(v) for k, v in trainer.metrics.items()}
        mlflow.log_metrics(metrics=metrics_dict, step=trainer.epoch)

import mlflow

mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Users/{user-id}/{project-name}")

# 실험 세션 생성
with mlflow.start_run():

	model.add_callback("on_fit_epoch_end",on_fit_epoch_end)
	results = model.train(
    **opt
    )

04. 파라미터 최적화

🍒 8기가 이상의 램을 할당할 것을 제안

• 오류가 엄청 발생하는데, 실습용으로 epoch를 낮춰서 실행

패키지 설치

pip install -U ultralytics "ray[tune]"  # install and/or update
pip install wandb  # optional

from ray import tune

model = YOLO(f"{MODEL}/train/weights/last.pt")

result = model.tune(
    data="/opt/ml/yujin/DataAnalysisPractice/Dacon/03.데이콘 합성데이터 기반 객체 탐지 AI 경진대회/data/yolo/custom.yaml",
    space={"lr0": tune.uniform(1e-5, 1e-1)},
    train_args={"epochs": 10}
)