Jetson Nano에서, Custom Data YOLO 구동하기 (with TensorRT)

일본 교환학생 중, 3학년 프로젝트로 Jetson Nano를 이용한 팀 프로젝트를 진행하였습니다. 

Yolo를 이용한 Object detection을 구현하려고 했는데, Yolo를 그대로 이용하면, Jetson Nano의 낮은 하드웨어 스펙 때문에, 성능이 매우 저열했습니다. (nano 모델인데도, 모델 로딩에 5분, FPS는 약 3~4 정도 나옴.)

따라서 TensorRT를 이용하여, Jetson Nano에 최적화 하는 과정을 거쳤습니다.

 

What is TensorRT?  

TensorRT란, NVIDIA CUDA® 병렬 프로그래밍 모델을 기반으로 구축된 SDK로, 모델을 NVIDA GPU에 맞게 최적화 시켜, 추론 속도를 수십배 가량 향상시켜주는 엔진입니다.

Python을 API 단계에서 지원하므로, 저처럼 관련 지식이 전무한 사람도 쉽게 사용이 가능했습니다. 

 

필자의 경우, 다른 환경(Google Colab)에서 학습시킨 Yolov5 가중치 파일을, Jetson Nano상에서TensorRT 엔진으로 변환시켜 이용하였습니다. 

 

Jetson Nano 환경

Jetson Nano 2GB

OS : JetPack 4.7

Python : version 3.6.9

CUDA : version 10.2

CuDnn : version 7.6.5

Memory Swap : 5GB

 

OpenCV는, CUDA 가속 활성화 해주어야 원활하게 구동될 것이라고 생각합니다. 

 

커스텀 데이터 학습

YOLOv5 ver 6.0을 기반으로, 대략 이미지 3000장 정도를 라벨링하여 YOLOv5n 모델을 이용하여 학습했습니다.

선택 가능한 사전 훈련 모델은 위와 같은데, 필자는 Nano 모델을 사용하였습니다.

 

커스텀 학습의 방법은, 다음 YOLOv5의 공식 문서를 참고해주세요. 한국어를 지원하므로, 쉽게 따라할 수 있습니다. 

https://docs.ultralytics.com/ko/yolov5/tutorials/train_custom_data/

사용자 지정 데이터 훈련

학습이 완료되면, pt 파일이 생성됩니다. 

 

Jetson Nano 환경 설정

이제 Jetson Nano로 돌아옵니다. TensoRT를 이용하기 위해, 다음 깃허브를 이용하였습니다. 

https://github.com/mailrocketsystems/JetsonYolov5.git

GitHub - mailrocketsystems/JetsonYolov5

git clone https://github.com/mailrocketsystems/JetsonYolov5.git

해당 Repo를 Jetson Nano에 다운로드합니다.

 

sudo apt-get update
sudo apt-get install -y liblapack-dev libblas-dev gfortran libfreetype6-dev libopenblas-base libopenmpi-dev libjpeg-dev zlib1g-dev
sudo apt-get install -y python3-pip

필요 라이브러리를 설치합니다. 

 

numpy==1.19.0
pandas==0.22.0
Pillow==8.4.0
PyYAML==3.12
scipy==1.5.4
psutil
tqdm==4.64.1
imutils

파이썬 패키지들을 설치합니다. Numpy의 경우, JetPack 순정 상태에서 이미 설치되어있는데, 삭제 후 해당 버전으로 설치하였습니다.

 

export PATH=/usr/local/cuda-10.2/bin${PATH:+:${PATH}}
export LD_LIBRARY_PATH=/usr/local/cuda-10.2/lib64:$LD_LIBRARY_PATH
python3 -m pip install pycuda --user

PyCUDA를 설치합니다. CUDA 아키텍쳐를 Python환경에서 이용할 수 있게 도와주는 패키지입니다. 

처음에 PyCUDA를 설치할 때, 계속 오류가 발생하여 골머리를 썩었는데, 위 Repo처럼 PATH 설정 후 설치 시 문제 없이 설치되었습니다.

다른 패키지와 다르게 설치 시 꽤 시간이 걸립니다. 

 

wget https://nvidia.box.com/shared/static/fjtbno0vpo676a25cgvuqc1wty0fkkg6.whl -O torch-1.10.0-cp36-cp36m-linux_aarch64.whl
pip3 install torch-1.10.0-cp36-cp36m-linux_aarch64.whl
git clone --branch v0.11.1 https://github.com/pytorch/vision torchvision
cd torchvision
sudo python3 setup.py install

Torch 및 Torchvison 설치 과정인데, 필자는 예전에 이미 설치 해 둔 상태여서 건너뛰었습니다.

 

sudo apt install cmake

Cmake를 설치합니다. 

 

YOLOv5 모델 TensorRT 엔진으로 변환 

먼저, 커스텀 학습으로 생성된 pt 파일을, 방금 받은 Repo 에 이동시켜주어야 합니다. 

$ cd JetsonYoloV5
$ python3 gen_wts.py -w CUSTOM_MODEL.pt -o CUSTOM_MODEL.wts

gen_wts.py를 통해, pt 파일을 wts 파일로 변환시킵니다. 

 

Make

yolov5 내부에 빌드 디렉토리를 만듭니다. 생성된 wts 파일을 빌드 디렉토리에 복사하여 붙여넣고 아래 명령을 실행합니다. 커스텀 모델을 사용하였기 때문에, yolov5/src/config.h에서 kNumClass를 업데이트해야 합니다. 자신의 모델의 class 수에 맞춰주시면 됩니다.

cd yolov5/
mkdir build
cd build
cp ../../CUSTOM_MODEL.wts .
cmake ..
make

 

engine 파일 빌드

./yolov5_det -s CUSTOM_MODEL.wts CUSTOM_MODEL.engine n

 

wts 파일을 통해 엔진 파일이 생성됩니다.

 

해당 파일을 이용하여 Jetson Nano에서 TensorRT를 이용할 수 있습니다.