Yolov5의 model을 Tx2 engine을 만들어봅시다 (Onnx, netron)

안녕하세요

Yolov5는 pytorch 기반의, 객체 검출을 쉽게 해 줄수 있는 프레임워크입니다.

하지만 yolov5에서 만든 모델을 다른곳에서 사용하려면, 그에 맞게 모델을 변환시켜 주어야 합니다.

그럴때 사용할 수 있는 기술이  Onnx 입니다.

 

onnx소개 : https://github.com/onnx/onnx, https://www.megazone.com/techblog_180816_aws-deep-learning-amis/#:~:text=ONNX%20%EB%9E%80%20%EB%AC%B4%EC%97%87%EC%9D%B8%EA%B0%80%3F,%EC%97%B0%EC%82%B0%EC%9E%90%20%EB%AA%A9%EB%A1%9D%EC%9D%84%20%EC%A0%95%EC%9D%98%ED%95%A9%EB%8B%88%EB%8B%A4.,

 

detect.py로 제대로 잘 잡히는지 본 이후, onnx file로 export 합니다.

 

python3 models/export.py --wieghts best.pt --batch 1 --img-size [width height]

onnx를 시각화할 수 있는 사이트가 있습니다.

https://netron.app/

 

 

그리고 나서 tensorRT를 이용해서 jetsonTx2에서 처리해야 합니다. 

그 이유는 Tx2 의 cpu type이 intel type이 아니기 때문입니다.

txtexec --onnx=[onnx 이름] --explicitBatch --saveEngine=[바꾸고자하는 engine파일] --fp16 --workspace=4096

trtexec는 /usr/src/tensorrt/bin/trtexec에 있을 수 있습니다!

find . -name trtexec로 검색 가능!

 

커스텀 Parser 만들기

 

onnx output에 맞춰 /opt/nvidia/deepstream/deepstream-5.0/sources/libs/nvdsinfer_customparser 안에 있는

nvdsinfer_custombboxparser.cpp를 수정 후 아래와 같이 입력합니다.

 

cd /opt/nvidia/deepstream/deepstream-5.0/sources/libs/nvdsinfer_customparser
sudo make

 

이렇게 하면 libnvds_infercustomparser.so가 생성됩니다. 

 

deepstream_python_apps 아래 폴더 안에 

dtest_pgie_config.txt에 이렇게 세팅 해야 합니다.

 

model-engine-file=/home/dev/yolov5.engine # engine 파일 위치 설정
parse-bbox-func-name=NvDsInferParseCustomYOLOv5 # 커스텀 추출기 함수명
custom-lib-path=/opt/nvidia/deepstream/deepstream-5.0/sources/libs/nvdsinfer_customparser/libnvds_infercustomparser.so
# 커스텀 추출기 파일명

custom parser를 쓰겠다는 내용입니다. 

이렇게 맞추고 deepstream app을 실행하면 사용할 수 있습니다.

 

ref : 한기대 이선훈 연구원