[TensorRT 튜토리얼] 3-1. ONNX/TRT engine conversion

이전 2. TensorRT ecosystem 글에서 설명한 바와 같이 TensorRT는 conversion/deployment로 나뉘며 그 중 더 중요하다고 볼 수 있는 conversion to trt engine 과정은 1) TF-TRT 2) ONNX 총 2가지 방식으로 나뉜다. 1) TF-TRT 방식의 경우, pytorch 사용자가 더 많으니 스킵하고, ONNX를 이용한 conversion/deployment 과정만 document를 따라가보려고 한다. 

 

ONNX가 tensorflow도 지원을 하니 굳이 TF-TRT부터 배울 필요는 더욱 없는 것 같으니, 이 글에서는 conversion만 정리해본다.

 

Prerequisite

pip install onnx # necessary
pip install tf2onnx # optional, only for tensorflow

 

onnx는 pytorch/tensorflow 모델을 변환할 때 공통적으로 쓰이니 필수적으로 설치해야 하고 tf2onnx는 tensorflow-to-onnx변환할 때만 설치하면 된다.

 

tensorflow를 안 써서 tensorflow도 설치해야 했는데 과정은 다음과 같다. 

pip install tensorflow[and-cuda]
python3 -c "import tensorflow as tf; print(tf.config.list_physical_devices('GPU'))"

 

[and-cuda]가 버전명을 넣어야 하는 variable이 아니다. 문장 그대로 입력하면 된다. 장치 목록이 제대로 출력되었다면 설치가 잘 된 것이다.

 

CUDA 12 이상은 아직 안된다고 하니 버전에 유의하자. (20240115)

 

 

TensorFlow-to-ONNX

 

튜토리얼 차 설치한 TensorRT는 CUDA 12.1로 환경을 구성했기 때문에 Tensorflow가 아직 CUDA 12 이상은 지원하지 않는 문제가 있었다. 그렇다고 CUDA/cuDNN/TensorRT를 다시 깔 수는 없으니 그냥 진행해보았다. CPU 모델 변환한 것이라 보면 된다. 

 

import os
import onnx
import time
import tf2onnx
import numpy as np
import tensorflow as tf
from tensorflow.keras.applications import ResNet50

### preparations
batch_size= 8
batch_img = np.ones([batch_size, 224, 224, 3], dtype=np.float32)
model = ResNet50(weights="imagenet")
model.save("./resnet")
predictions = model.predict(batch_img) # warm up

 

아주 기본적인 tensorflow resnet50 모델 준비, 더미 입력도 준비

### tensorflow
start = time.time()
result = model.predict_on_batch(batch_img)
end = time.time()

print(f"RAW Tensorflow spped : {end-start}sec in CPU")
# 0.7622sec

 

간단히 동작 확인 까지 마쳤다.

### convert to onnx
input_signature = [tf.TensorSpec([batch_size, 224, 224, 3], tf.float32, name='x')]
onnx_model, _ = tf2onnx.convert.from_keras(model, input_signature)
onnx.save(onnx_model, "./resnet.onnx")

 

첫번째 방식은 tf2onnx를 직접 import 해서 코드 상에서 변환하는 것이다. from_function 류 기능이 지원되므로 간단한 형태의 모델들은 호출로 변환이 끝난다.

 

### convert to onnx
model.save("./resnet")
os.system("python3 -m tf2onnx.convert --saved-model ./resnet --output temp.onnx")
onnx_model = onnx.load_model('temp.onnx')

 

두번째 방식은 커맨드 라인에서 돌리는 방식이다. os.system을 이용하면 커맨드라인 명령어도 코드 상에서 돌릴 수 있으므로 뭐 큰 차이는 없겠다. 근데 개인적으로 온갖 메뉴얼들과 글들은 커맨드 라인을 선호하는 것 같아서 두번째 방식을 따르는게 좋을 것 같다.

 

saved-model, output 에는 절대/상대 경로 자유롭게 넣어주면 된다.

 

PyTorch-to-ONNX

 

훨씬 간단한 느낌이다. torch 자체에 onnx가 들어있는 형태로 정리해두었기 때문에 코드 상에서 작업이 끝난다. 

import torch
import torch.onnx
import torchvision.models as models

### preparations
batch_size = 8
dummy_img = torch.randn(batch_size, 3, 224, 224)
model = models.resnet50(pretrained=True, progress=False).eval()

 

준비 과정은 동일하고 입력만 tensor라는 것이 다르다.

### convert to onnx
torch.onnx.export(model, dummy_img, "./resnet_torch.onnx", verbose=True)

 

호출만으로 변환 과정이 끝난다. 

 

ONNX-to-TRT

Using trtexec tool

 

원래는 변환 과정을 TensorRT API를 활용해서 직접 onnx to trt engine 코드를 구현해야 한다. 하지만 이 구현이 생각보다 옵션들이 많고 복잡하기에 어느 정도 수준까지는 코드 구현없이 자동으로 해줄 수 있는 tool이 존재한다. 그게 trtexec인데 특별한 모델이 아닌 경우 단순히 tool을 이용해 trt engine으로 빌드할 수 있다. 

trtexec --onnx=./resnet.onnx --saveEngine=resnet_engine.trt  --explicitBatch --inputIOFormats=fp16:chw --outputIOFormats=fp16:chw --fp16

 

fp16으로 precision낮추는 걸 포함하고 싶으면 위와 같이 관련 argument를 추가해주면 되며,

trtexec --onnx=resnet50_pytorch.onnx --saveEngine=resnet_engine_pytorch.trt  --explicitBatch

 

precision은 그대로 유지하고 싶을 경우에는 위와 같이 간단한 형태로 입력하면 된다. 

 

Issue

 

trtexec: command not found 오류가 뜰 수 있다.

 

만약 TensorRT 설치 시 tar file로 설치했다면 bashrc에 제대로  LD_LIBRARY_PATH export 안해줘서 그런 것이고, deb file로 설치했다면 /usr/src/tensorrt/bin 안에 trtexec 있는데 경로를 못찾아서 그렇다.

 

후자의 경우, 절대 경로로 입력하면 실행되는 것을 볼 수 있음

trtexec --help # fail
usr/src/tensorrt/bin/trtexec --help # success

 

이 경우 간단하게 alias로 trtexec만 입력해도 절대 경로를 입력해준 것과 같이 매칭해주면 해결된다. 

alias trtexec="/usr/src/tensorrt/bin/trtexec"

 

alias로 매칭을 해두어도 코드 상에서 os.system()을 활용하여 trtexec 커맨드를 입력하면 여전히 못 찾는다. 따라서 os.system()을 활용해서 작업할 경우에는 full path를 같이 적어주어야 된다.  

 

Using TensorRT API

 

생략한다. 이건 튜토리얼에서 커버할 범위가 아니다. trtexec toolbox를 이용해서 간단히 작업하는 것 외에는 정형화된 틀이 없고 모델 구조나 팩토링 방식에 따라 달라지는 것이기 때문에 실제 코드 리뷰를 통해 경험으로 API를 조합해서 구현해야 하는 영역이다.