NVIDIA Jetson TensorRT Inference

Sep 27, 2023

NVIDIA Jetson TensorRT Inference

이 글에서는 pytorch를 사용해 만든 모델을 torch -> onnx -> trt 로 변환하여 최적화 후 배포하기위한 과정을 소개합니다.

Install torch

아래의 링크를 참고하여 설치하였습니다.

https://docs.nvidia.com/deeplearning/frameworks/install-pytorch-jetson-platform/index.html

Install dependencies

sudo apt-get -y update; 
sudo apt-get -y install autoconf bc build-essential g++-8 gcc-8 clang-8 lld-8 gettext-base gfortran-8 iputils-ping libbz2-dev libc++-dev libcgal-dev libffi-dev libfreetype6-dev libhdf5-dev libjpeg-dev liblzma-dev libncurses5-dev libncursesw5-dev libpng-dev libreadline-dev libssl-dev libsqlite3-dev libxml2-dev libxslt-dev locales moreutils openssl python-openssl rsync scons python3-pip libopenblas-dev;

Install Torch 2.0

export TORCH_INSTALL=https://developer.download.nvidia.cn/compute/redist/jp/v511/pytorch/torch-2.0.0+nv23.05-cp38-cp38-linux_aarch64.whl
python3 -m pip install --upgrade pip; python3 -m pip install aiohttp numpy=='1.19.4' scipy=='1.5.3' export "LD_LIBRARY_PATH=/usr/lib/llvm-8/lib:$LD_LIBRARY_PATH"; python3 -m pip install --upgrade protobuf; python3 -m pip install --no-cache $TORCH_INSTALL

저의 경우 LDLIBRARYPATH 환경변수 관련해서 에러 메세지가 나왔지만 정상적으로 설치되었습니다.

Verify Install

$ export LD_LIBRARY_PATH=/usr/lib/llvm-8/lib:$LD_LIBRARY_PATH
$ python3

>>> import torch

에러메세지 없이 torch가 import되고 torch.cuda.is_available() == True 면 정상적으로 설치된 것입니다.

Torch to Onnx

https://tutorials.pytorch.kr/advanced/super_resolution_with_onnxruntime.html

위의 글을 따라 쉽게 변환할 수 있습니다.

Onnx to TRT

onnx 변환을 jetson 이 아닌 다른 환경에서 했다면 rsync를 통해 옮겨줍니다.

rsync -rvz <onnx_name>.onnx <user_name>@<jetson_ip>:<dest_path>

앞선 과정처럼 환경을 설치했다면, tensorrt가 이미 깔려있기 때문에 아래의 명령어로 onnx를 trt로 빌드만 하면 됩니다.

/usr/src/tensorrt/bin/trtexec --onnx=<onnx_name>.onnx --saveEngine=<trt_name>.engine --int8 --verbose

빌드에 성공하면 아래와 같은 메세지 출력 후 파일이 생성됩니다.

&&&& PASSED TensorRT.trtexec [TensorRT v8502]
...

TRT Inference

https://docs.nvidia.com/deeplearning/tensorrt/api/index.html

위의 링크에서 trt의 전반적인 API를 확인할 수 있습니다.

pytorch로 TRT inference 하는 코드는 아래와 같이 작성할 수 있습니다.

몇몇 function은 추후 deprecated 될 수 있으니 참고바랍니다.

import argparse

import torch
import tensorrt as trt

def load_trt_engine(engine_path: str):
    with open(engine_path, "rb") as stream:
        serialized = stream.read()
    logger = trt.Logger(trt.Logger.ERROR)
    runtime = trt.Runtime(logger)
    engine = runtime.deserialize_cuda_engine(serialized)
    return engine

class TRTModel:
    def __init__(self, trt_path):
        self.model = load_trt_engine(trt_path)
        self.model_ctx = self.model.create_execution_context()

    @torch.no_grad()
    def run_engine(self, input_tensor):
        idx = self.model.get_binding_index("output_0")
        device = torch.device("cuda")
        dtype = torch.int8
        shape = tuple(self.model.get_binding_shape(idx))

        input_tensor = input_tensor.to(device=device, dtype=dtype)

        bindings = [None] * 2
        bindings[0] = input_tensor.contiguous().data_ptr()

        output = torch.empty(size=shape, dtype=dtype, device=device)
        bindings[idx] = output.data_ptr()

        self.model_ctx.execute_async_v2(
            bindings, torch.cuda.current_stream().cuda_stream
        )
        torch.cuda.synchronize()
        return output

if __name__ == "__main__":
    import time
    import numpy as np
    parser = argparse.ArgumentParser()
    parser.add_argument("--trt_path", type=str)
    args = parser.parse_args()
    trt_model = TRTModel(args.trt_path)
    total_time = []
    
    for i in range(100):
        input_tensor = torch.rand(1, 21, 1088, 1920)
        start = time.time()
        output = trt_model.run_engine(input_tensor)
        total_time.append(time.time() - start)
    print(np.mean(total_time[1:]))

output

>>> 0.24760057227780122