TensorRT, Semantic Segmentation, ONNX 에서 TensorRT 모델 변환, C++ 에서 모델 실행까지

엔비디아(NVIDIA)에서 만든 TensorRT 라는 SDK(Software Development Kit)는 학습된 딥러닝 모델을 최적화하여 높은 추론 성능을 가능하게 하게 한다.

TensorRT를 활용한 딥러닝 모델 개발 과정

 

TensorRT 공식 깃허브 저장소에서 간단한 예제를 통해 자세히 알아보자. 먼저, 도커(Docker)가 설치되어 있다고 가정한다.

$ git clone https://github.com/NVIDIA/TensorRT.git
$ cd TensorRT/quickstart

 

 

도커 컨테이너(Container)를 실행한다. -v는 마운트 옵션으로, 현재 폴더와 컨테이너 내부를 연결해줄 것이다.

 

# In CMD, use %cd%
$ docker run --rm -it --gpus all -p 8888:8888 -v %cd%:/workspace -w /workspace/SemanticSegmentation nvcr.io/nvidia/pytorch:20.12-py3 bash

 

다소 요란한 로그들과 함께 컨테이너에 연결(Attach)되었다.

저작권만 11줄이다.

 

오늘 실험하게 될 영상이다.

해당 영상에 Semantic Segmentation을 통해 개체들(사람, 양)을 다른 색상으로 구분하고자 한다.

https://pytorch.org/assets/images/deeplab1.png

 

앞서 생성한 컨테이너에서 아래의 명령어를 입력한다.

 

$ python export.py

 

이를 통해 torch hub 에서 fcn-resnet101 모델을 다운받고, 이를 ONNX 포맷으로 변환한다.

토치 모델 파일(*.pth)을 받아 ONNX 로 변환하여 저장

 

그리고 TensorRT 모델로 변환하기 위해 trtexec 모듈을 활용한다.

$ trtexec --onnx=fcn-resnet101.onnx --fp16 --workspace=64 --minShapes=input:1x3x256x256 --optShapes=input:1x3x1026x1282 --maxShapes=input:1x3x1440x2560 --buildOnly --saveEngine=fcn-resnet101.engine

 

실행 이후 로그에서 &&&& PASSED 이하 메세지가 마지막으로 있어야 변환에 성공한 것이다.

변환에 실패하는 경우는 여러가지이나 보통은 용량 부족이니, workspace 매개변수를 더 늘려보자.

 
그럼 이제 C++ 에서 실행하기 위해 make 명령어를 통해 빌드하고 이를 실행한다.

마찬가지로 컨테이너 안에서 실행하는 것이다.

$ make
$ ./bin/segmentation_tutorial

 
실행 결과로 output.ppm 파일이 생성되었을 것이다.

ppm 은 Portable Pixmap Format의 약자이다. (1980년 후반 다른 플랫폼 간 이미지 공유를 위해 만들었다고 한다.)

여러 뷰어들이 있지만 gimp를 추천한다. 

output.ppm

 

주피터 노트북(Jupyter Notebook)으로 실행하면서 보고 싶다면 pycuda를 설치해야한다.

그리고 주피터 노트북을 실행하자.

$ pip install pycuda
$ jupyter notebook --port=8888 --no-browser --ip=0.0.0.0 --allow-root

 
앞서 도커 컨테이너 생성시 8888 포트를 열었다. 따라서, localhost:8888 로 접속하면 주피터 환경을 사용할 수 있다.

상단의 토큰(혹은 비밀번호)은 아래와 같이 확인이 가능하다.

 

주피터 노트북에서도 동일한 결과를 확인할 수 있다.

 

 

참고자료 : 
https://docs.nvidia.com/deeplearning/tensorrt/quick-start-guide/index.html#runtime

Quick Start Guide :: NVIDIA Deep Learning TensorRT Documentation