YOLOv5 커스텀 데이터 학습 및 RK3588 보드 배포를 위한 RKNN 변환 가이드

1. YOLOv5 소스 코드 확보 및 데이터셋 준비

먼저 Ultralytics의 공식 YOLOv5 저장소에서 소스 코드를 다운로드합니다. 커스텀 객체 인식을 위해 LabelImg와 같은 도구를 사용하여 이미지의 객체 위치와 클래스를 지정하는 어노테이션 작업을 수행합니다.

데이터셋은 학습(train)과 검증(val)으로 나누어 다음과 같은 디렉토리 구조로 구성해야 합니다.

/workspace/custom_dataset
    /images
        /train
            img_001.jpg
            ...
        /val
            img_101.jpg
            ...
    /labels
        /train
            img_001.txt
            ...
        /val
            img_101.txt
            ...

2. 학습 설정 및 모델 구성

data 폴더 내에 데이터셋 경로와 클래스 정보를 담은 custom_data.yaml 파일을 생성합니다.

# custom_data.yaml
train: /workspace/custom_dataset/images/train
val: /workspace/custom_dataset/images/val

nc: 2  # 클래스 개수
names: ['target_a', 'target_b']  # 클래스 이름 리스트

이후 models/yolov5s.yaml 파일의 nc 파라미터를 자신의 클래스 개수에 맞게 수정합니다.

3. 모델 학습 실행

준비된 환경에서 다음 명령어를 통해 학습을 시작합니다. 하드웨어 사양에 따라 배치 크기와 이미지 크기를 조절할 수 있습니다.

# GPU 사용 시
python train.py --img 640 --batch 16 --epochs 100 --data data/custom_data.yaml --weights yolov5s.pt --device 0

# CPU 사용 시
python train.py --img 640 --batch 16 --epochs 100 --data data/custom_data.yaml --weights yolov5s.pt --device cpu

학습 완료 후 생성된 best.pt 파일을 사용하여 detect.py로 추론 결과를 확인합니다.

4. RKNN 배포를 위한 모델 변환

RKNN(Rockchip Neural Network) 환경에서 최적의 성능을 내기 위해서는 YOLOv5의 출력 레이어를 수정해야 합니다. models/yolo.py 내 Detect 클래스의 forward 함수를 다음과 같이 단순화하여 수정합니다.

def forward(self, x):
    # RKNN 배포를 위해 후처리를 제외한 원시 출력만 반환
    outputs = []
    for i in range(self.nl):
        outputs.append(torch.sigmoid(self.m[i](x[i])))
    return outputs

수정이 완료되면 PyTorch 모델을 ONNX 형식으로 변환합니다.

python export.py --weights best.pt --include onnx

5. ONNX 모델을 RKNN으로 변환

RKNN-Toolkit2를 사용하여 ONNX 파일을 RKNN 파일로 변환합니다. Docker 환경에서 수행하는 것을 권장합니다.

# Docker 컨테이너 실행 예시
docker run -it --name rknn_dev -v $(pwd):/app rknn-toolkit2:latest /bin/bash

# 변환 스크립트 실행 (test.py 등을 수정하여 사용)
python convert_to_rknn.py

변환 시 데이터셋의 앵커(Anchor) 값과 Sigmoid 함수 중복 적용 여부를 반드시 확인해야 합니다. 훈련 시 자동으로 계산된 앵커 값이 기본 COCO 앵커와 다를 경우, 변환 스크립트의 앵커 설정도 동일하게 맞춰야 정확한 탐지가 가능합니다.

6. RK3588 보드에서 추론 실행

RK3588의 NPU를 활용하기 위해 C++ 데모 소스 코드를 수정합니다. 주요 수정 사항은 다음과 같습니다.

include/postprocess.h: OBJ_CLASS_NUM을 학습시킨 클래스 개수로 변경
model/coco_80_labels_list.txt: 자신의 클래스 이름 목록으로 교체
src/postprocess.cc: 모델 학습 시 생성된 커스텀 앵커 값 반영

빌드 및 실행 과정은 다음과 같습니다.

# 빌드 스크립트 실행
bash ./build-linux_RK3588.sh

# 생성된 실행 파일 경로로 이동
cd install/rknn_yolov5_demo_Linux

# 추론 실행
./rknn_yolov5_demo ./model/RK3588/my_model.rknn ./model/test_image.jpg

7. 주요 문제 해결 가이드

메모리 부족 (OOM) 발생

학습 시 cuda: out of memory 오류가 발생하면 --batch 크기를 줄여 메모리 부하를 낮춥니다.

앵커 불일치 문제

검색 결과 상자가 너무 작거나 위치가 부정확하다면 앵커 설정이 원인일 가능성이 높습니다. 아래 코드로 best.pt의 실제 앵커를 확인하고 배포 코드에 반영하십시오.

from models.experimental import attempt_load
net = attempt_load('best.pt')
m = net.module.model[-1] if hasattr(net, 'module') else net.model[-1]
print(m.anchor_grid)

OpenCV 라이브러리 참조 오류

C++ 빌드 시 VideoCapture 관련 오류가 발생하면 CMakeLists.txt의 OpenCV_DIR 경로를 보드 내 실제 설치된 경로로 명시해 주어야 합니다.

8. 외부 카메라 연동

RK3588에 연결된 USB 카메라의 인덱스를 확인하려면 다음 명령어를 사용합니다.

v4l2-ctl --list-devices

드라이버 명이 uvcvideo인 장치의 /dev/videoX 번호를 확인한 후, OpenCV의 VideoCapture(X) 함수에 해당 번호를 입력하여 실시간 스트리밍 추론을 구현할 수 있습니다.

태그: YOLOv5 rk3588 RKNN DeepLearning ObjectDetection

8월 2일 20:08에 게시됨

괴물 클럽