Pytorch To TensorRT Engine

TensorRT란?

TensorRT는 학습된 딥러닝 모델을 최적화하여, NVIDIA GPU 상에서의 추론 속도를 수배 ~ 수십배까지 향싱시켜 딥러닝 서비스를 개선하는데 도움을 줄 수 있는 모델 최적화 엔진이다. 실시간으로 그래픽 연산을 수행해야 하거나, 빠른 연산속도를 요구하는 분야에 많이 사용되는 것 같다. 본 포스팅에서는 Pytorch 모델을 Onnx로 변환 후 TensorRT Engine을 생성하는 방법과, torch2trt 라이브러리를 사용하여 바로 TensorRT Engine으로 변환하는 방법을 소개한다.

Pytorch to Onnx to TensorRT Engine

About Onnx

Onnx란 Open Neural Network eXchange의 줄임말로서, Facebook과 MS가 개발한 신경망 모델의 호환 포맷이다. Onnx는 특정 프레임워크에서 작성한 모델을 다른 프레임 워크에서 사용할 수 있게하는 도구이며, pytorch, caffe, CNTK, MXNet등을 지원한다. 예를 들어 Pytorch 작성 -> Onnx 변환 -> Caffe에서 Imfort와 같은 작업이 가능해진다.

Create Onnx

Onnx 모델을 생성할 때는 Pytorch 모델에 입력되는 input shape 과 동일해야한다. shape 만 맞춰준다면 어떠한 랜덤 값이 들어가도 무방하다. torch.onnx.export 시 중요한 것은 파이토치 모델, 입력 값 만 있으면 Onnx 모델을 만들 수 있다. torch.onnx.export 함수는 기본적으로 scripting 이 아닌 tracing 을 사용하기 때문에 example input 을 넣어주어야 한다. 때에 따라 필요한 옵션을 추가하여 Onnx 모델을 생성 할 수 있다. 이 때 verbose=False 라는 옵션을 추가하면, 어떤 노드들이 Onnx 모델을 위해 인식되었는지 로그를 띄울 수 있다.

Onnx 모델 생성시에는 입력 데이터와 크기가 동일한 더미 데이터를 사용해야한다. 이 때, 더미 데이터에는 어떠한 값이 들어가도 상관 없으며 shape만 맞추어주면 된다. torch.randn을 사용하여도 되고, torch.empty를 사용해주어도 된다. 예를 들어 입력데이터가 (3,224,224)의 이미지라면, 배치 차원도 포함해서 (1, 3, 224, 224)의 더미 데이터를 만들어서 넣어주면 된다.

import numpy as npq

from torch import nn
import torch.utils.model_zoo as model_zoo
import torch.onnx

import torch.nn as nn
import torch.nn.init as init

batch_size = 1
model.load_state_dict(torch.load(PATH))
torch_model.eval()

dummy_data = torch.randn(batch_size, 3, 224, 224, requires_grad=True)

# 가장 기본 형태
torch.onnx.export(torch_model, dummy_data, 'super_resulotion.onnx')

# 모든 options
torch.onnx.export(torch_model,               # 실행될 모델
                  dummy_data,                # 모델 입력값 (튜플 또는 여러 입력값들도 가능)
                  "super_resolution.onnx",   # 모델 저장 경로 (파일 또는 파일과 유사한 객체 모두 가능)
                  export_params=True,        # 모델 파일 안에 학습된 모델 가중치를 저장할지의 여부
                  opset_version=10,          # 모델을 변환할 때 사용할 ONNX 버전
                  do_constant_folding=True,  # 최적하시 상수폴딩을 사용할지의 여부
                  input_names = ['input'],   # 모델의 입력값을 가리키는 이름
                  output_names = ['output'], # 모델의 출력값을 가리키는 이름
                  dynamic_axes={'input' : {0 : 'batch_size'},    # 가변적인 길이를 가진 차원
                                'output' : {0 : 'batch_size'}},
                verbose=False                # 어떤 노드들이 Onnx 모델을 위해 인식되었는지 로그를 띄울 수 있다. 
                )

Create TensorRT Engine by Onnx

onnx_file_path와 engine_file_path를 인자로 받는 get_engine 함수이다. 만약 TensorRT Engine 파일이 존재한다면 engine_file_path에 경로를 넣어주어서 확인이 가능한 것 같다. 그렇지 않으면 Onnx를 TensorRT Engine파일로 변환하여 return하는 방식이다.

def get_engine(onnx_file_path, engine_file_path=""):
    def build_engine():
        with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.OnnxParser(network, TRT_LOGGER) as parser:
            builder.max_workspace_size = 1 << 30 # 1GB
            builder.max_batch_size = 1

			if not os.path.exists(onnx_file_path):
                print('ONNX file {} not found, please run yolov3_to_onnx.py first to generate it.'.format(onnx_file_path))
                exit(0)
            print('Loading ONNX file from path {}...'.format(onnx_file_path))
            with open(onnx_file_path, 'rb') as model:
                print('Beginning ONNX file parsing')
                parser.parse(model.read())
            print('Completed parsing of ONNX file')
            print('Building an engine from file {}; this may take a while...'.format(onnx_file_path))
            engine = builder.build_cuda_engine(network)
            print("Completed creating Engine")
            
            with open(engine_file_path, "wb") as f:
                f.write(engine.serialize())
            return engine

    if os.path.exists(engine_file_path):
        # If a serialized engine exists, use it instead of building an engine.
        print("Reading engine from file {}".format(engine_file_path))
        with open(engine_file_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
            return runtime.deserialize_cuda_engine(f.read())
    else:
        return build_engine()