Wego LIMO 딥러닝 모델 학습과 TensorRT 최적화 방법 #2

5. YOLO 객체 인식 모델 학습

YOLO 모델은 표지판, 신호등, 보행자처럼 이미지 안에서 위치를 특정해야 하는 객체 인식 작업에 사용합니다.

Wego LIMO에서 수집한 YOLO 데이터셋은 이미지와 라벨이 학습용과 검증용으로 나뉘어 있어야 합니다.

권장 폴더 구조는 다음과 같습니다.

dataset/
└── yolo_images/
    ├── images/
    │   ├── train/
    │   └── val/
    └── labels/
        ├── train/
        └── val/

각 이미지 파일에는 동일한 이름의 라벨 파일이 필요합니다.

images/train/frame_001.jpg
labels/train/frame_001.txt

YOLO 라벨 파일은 일반적으로 다음 형식을 사용합니다.

class_id center_x center_y width height

좌표값은 이미지 크기를 기준으로 0에서 1 사이로 정규화되어야 합니다.

1) YOLO 학습 환경 준비하기

Google Drive 연결과 데이터셋 압축 해제 과정은 TensorFlow 학습과 동일합니다.

from google.colab import drive

drive.mount('/content/drive')

!unzip -o /content/drive/MyDrive/dataset.zip -d /content

Ultralytics 패키지를 설치합니다.

!pip install -U ultralytics

설치된 버전과 GPU 상태를 확인합니다.

!yolo checks

Python 코드에서도 확인할 수 있습니다.

import torch
import ultralytics

print("Ultralytics:", ultralytics.__version__)
print("PyTorch:", torch.__version__)
print("CUDA 사용 가능:", torch.cuda.is_available())

2) YOLO 데이터 설정 파일 작성하기

YOLO 학습에는 이미지 경로와 클래스 이름을 정의한 YAML 파일이 필요합니다.

%%writefile /content/wego_limo.yaml
path: /content/dataset/yolo_images

train: images/train
val: images/val
test:

nc: 7

names:
  0: Speed_up
  1: Slow_down
  2: Turn_right
  3: Crosswalk
  4: Red_light
  5: Green_light
  6: Yellow_light
  7: Crossing_gate

클래스 번호는 라벨 파일에 기록된 클래스 ID와 정확히 일치해야 합니다.

예를 들어 라벨 파일의 첫 번째 숫자가 4이라면 해당 객체는 red_light로 처리됩니다.

클래스 순서를 변경하면 같은 라벨 파일을 사용하더라도 잘못된 객체 이름이 표시될 수 있습니다.

특히 yellow_light와 같은 클래스 이름의 철자를 통일해야 이후 ROS 2 노드에서 문자열을 비교할 때 오류가 발생하지 않습니다.

3) YOLO 데이터셋 검사하기

학습 전에 이미지와 라벨 개수를 확인합니다.

import glob

train_images = glob.glob(
    '/content/dataset/yolo_images/images/train/*'
)
train_labels = glob.glob(
    '/content/dataset/yolo_images/labels/train/*.txt'
)
val_images = glob.glob(
    '/content/dataset/yolo_images/images/val/*'
)
val_labels = glob.glob(
    '/content/dataset/yolo_images/labels/val/*.txt'
)

print("학습 이미지:", len(train_images))
print("학습 라벨:", len(train_labels))
print("검증 이미지:", len(val_images))
print("검증 라벨:", len(val_labels))

이미지 수와 라벨 수가 반드시 완전히 같을 필요는 없습니다. 객체가 없는 배경 이미지는 빈 라벨 파일을 사용하거나 라벨 파일 없이 구성할 수도 있습니다.

다만 라벨 누락이 의도된 것인지, 파일 복사 과정에서 발생한 오류인지는 구분해야 합니다.

4) YOLO11n 모델 학습하기

Wego LIMO처럼 제한된 연산 성능에서 실시간 추론이 필요한 경우에는 비교적 가벼운 yolo11n.pt 모델로 시작하는 것이 적합합니다.

다음 명령은 이미지 크기 640픽셀, 최대 300에포크 조건으로 학습합니다.

!yolo detect train \
    model=yolo11n.pt \
    data=/content/wego_limo.yaml \
    epochs=300 \
    imgsz=640 \
    batch=16 \
    device=0 \
    project=/content/drive/MyDrive/wego_limo_yolo \
    name=yolo11n_sign_detection

GPU 사용이 제한될 경우 device=0를 device=cpu로 수정한 후 재실행하시기 바랍니다.

GPU 메모리가 부족하면 배치 크기를 줄입니다.

batch=16 → batch=8 또는 batch=4

실행 시 아래와 같은 오류가 발생할 경우 GPU가 설정되지 않는 경우입니다. 런타임 유형변경에서 GPU를 추가해야 합니다.

설정이 완료되면 다시 실행합니다.

학습 결과는 다음과 같은 경로에 저장됩니다.

/content/drive/MyDrive/wego_limo_yolo/yolo11n_sign_detection-4

가장 좋은 검증 성능을 기록한 모델은 일반적으로 다음 파일입니다.

weights/best.pt

마지막 에포크의 모델은 다음 파일에 저장됩니다.

weights/last.pt

실제 주행에는 보통 best.pt를 사용하는 것이 좋습니다.

5) YOLO 학습 결과 검증하기

학습된 모델의 성능을 검증합니다.

!yolo detect val \
    model=/content/drive/MyDrive/wego_limo_yolo/yolo11n_sign_detection-4/weights/best.pt \
    data=/content/wego_limo.yaml \
    imgsz=640 \
    device=0

검증 결과에서는 다음 지표를 확인할 수 있습니다.

Precision
모델이 객체라고 판단한 결과 중 실제 객체가 맞는 비율입니다.
Recall
실제 객체 중 모델이 정상적으로 찾아낸 비율입니다.
mAP50
IoU 기준값 0.5에서 계산한 평균 정밀도입니다.
mAP50-95
여러 IoU 기준에서 종합적으로 계산한 더 엄격한 성능 지표입니다.

전체 mAP가 높더라도 특정 클래스의 성능이 낮을 수 있습니다. 예를 들어 적색 신호등 데이터는 충분하지만 황색 신호등 데이터가 적다면 yellow_light 클래스만 인식률이 낮아질 수 있습니다.

6) YOLO 이미지 추론 테스트

학습된 모델을 검증 이미지에 적용합니다.

!yolo detect predict \
    model=/content/drive/MyDrive/wego_limo_yolo/yolo11n_sign_detection-4/weights/best.pt \
    source=/content/dataset/yolo_images/images/val \
    imgsz=640 \
    conf=0.4 \
    save=True

conf는 객체를 표시할 최소 신뢰도입니다.

값이 너무 낮으면 잘못된 검출이 증가하고, 너무 높으면 실제 객체를 놓칠 수 있습니다. 실제 Wego LIMO 주행 환경에서 테스트한 뒤 적절한 값으로 조정해야 합니다.

7) 학습 모델을 Wego LIMO로 복사하기

Google Colab에서 학습한 파일을 다운로드한 뒤 Wego LIMO로 복사합니다.

TensorFlow 모델 파일 예시는 다음과 같습니다.

dl_drive.h5

YOLO 모델 파일 예시는 다음과 같습니다.

best.pt

SCP를 사용할 경우 다음과 같이 전송할 수 있습니다.

scp dl_drive.h5 wego@<LIMO_IP_ADDRESS>:/home/wego/
scp best.pt wego@<LIMO_IP_ADDRESS>:/home/wego/

<LIMO_IP_ADDRESS>는 실제 Wego LIMO의 IP 주소로 변경합니다.

6. 딥러닝 모델 최적화가 필요한 이유

TensorFlow의 .h5 모델이나 YOLO의 .pt 모델을 그대로 실행할 수도 있지만, 임베디드 환경에서는 추론 시간이 길어질 수 있습니다.

Wego LIMO에서 실시간에 가까운 속도로 동작시키려면 NVIDIA TensorRT 엔진으로 변환하는 것이 좋습니다.

TensorRT 최적화를 적용하면 다음과 같은 효과를 기대할 수 있습니다.

GPU 추론 속도 향상
모델 실행 지연 감소
메모리 사용 최적화
반복 추론 성능 향상
ROS 2 카메라 처리 주기 안정화

다만 TensorRT 엔진 파일은 일반적으로 생성한 장치의 GPU 구조, TensorRT 버전, CUDA 버전에 영향을 받습니다. 다른 Jetson 또는 다른 소프트웨어 환경에서 생성한 엔진을 그대로 복사하면 실행되지 않을 수 있습니다.

가장 안전한 방법은 실제로 모델을 실행할 Wego LIMO 장치에서 엔진 파일을 생성하는 것입니다.

1) TensorFlow 주행 모델 최적화

먼저 Wego LIMO에서 딥러닝용 Conda 환경을 활성화합니다.

conda activate dl_env

또는 터머널 실행 시 2번을 선택합니다.

ROS 2 작업 공간으로 이동합니다.

cd ~/wego_ws

빌드된 ROS 2 환경을 적용합니다.

source install/setup.bash

TensorFlow .h5 모델을 TensorRT 엔진으로 변환합니다.

ros2 run dl_ros2_application tf_to_trt <h5_file_path>

예를 들어 모델이 /home/wego/dl_drive.h5에 있다면 다음과 같이 실행합니다.

ros2 run dl_ros2_application tf_to_trt /home/wego/dl_drive.h5

변환 프로그램은 일반적으로 다음 과정을 수행합니다.

H5 모델 로드
중간 모델 형식으로 변환
TensorRT 네트워크 생성
최적화 프로파일 구성
엔진 빌드
.engine 파일 저장

모델 구조와 장치 성능에 따라 엔진 생성 과정에서 GPU 메모리를 많이 사용할 수 있습니다. 다른 딥러닝 노드가 실행 중이라면 종료한 뒤 변환하는 것이 좋습니다.

2) YOLO 객체 인식 모델 최적화

YOLO 모델도 동일하게 딥러닝 환경을 활성화합니다.

conda activate dl_env
cd ~/wego_ws
source install/setup.bash

YOLO .pt 모델을 TensorRT 엔진으로 변환합니다.

ros2 run dl_ros2_application yolo_to_trt <pt_file_path>

예를 들어 학습된 모델이 /home/wego/best.pt에 있다면 다음 명령을 사용합니다.

ros2 run dl_ros2_application yolo_to_trt /home/wego/best.pt

변환이 정상적으로 완료되면 .engine 확장자의 TensorRT 모델이 생성됩니다.

생성된 파일 이름과 저장 위치는 패키지 구현에 따라 달라질 수 있으므로 변환 로그를 확인해야 합니다.

find /home/wego -name "*.engine" 2>/dev/null

생성한 가중치 파일들을 아래의 위치로 복사합니다.

/home/wego/wego_ws/src/dl_ros2_application/weight/

재 빌드를 실행합니다.

3) TensorRT 엔진 생성 전 확인 사항

엔진 변환 전에 다음 항목을 확인합니다.

nvidia-smi

Jetson 환경에서는 nvidia-smi가 지원되지 않을 수 있으므로 다음 명령도 사용할 수 있습니다.

tegrastats

CUDA와 TensorRT 설치 상태를 확인합니다.

nvcc --version

dpkg -l | grep -i tensorrt

ROS 2 패키지가 정상적으로 인식되는지도 확인합니다.

ros2 pkg list | grep dl_ros2_application

실행 파일 목록은 다음 명령으로 확인할 수 있습니다.

ros2 pkg executables dl_ros2_application

tf_to_trt 또는 yolo_to_trt가 표시되지 않는다면 작업 공간 빌드 상태를 점검해야 합니다.

cd ~/wego_ws
colcon build --symlink-install
source install/setup.bash