[Python, Deep Learning] Tensorflow2 Classification 모델, .onnx & .trt 변환 및 Inference 방법

안녕하세요, 오늘은 Tensorflow 프레임워크에서 Classification 모델을 .onnx 파일과 .trt 파일로 변환하고 해당 파일을 활용하여 Inference 하는 방법에 대해서 알아보겠습니다.

 

이를 위해 TEST를 진행한 패키지 버전 및 환경은 다음과 같습니다.

 

- OS : Ubuntu 22.04
- CUDA : 11.8
- GPU : NVIDIA GeForce RTX 4070 Laptop GPU

- Python : 3.10.12

- tensorflow : 2.14.0
- tf2onnx : 1.16.1
- onnxruntime-gpu : 1.18.1
- nvidia-tensorrt : 8.4.3.1
- pycuda : 2024.1

 

먼저, tensorflow 2 모델을 onnx 파일로 변환하는 코드에 대해서 알아보겠습니다.

 

import tensorflow as tf

def convert_tf_to_onnx(tf_model_path, onnx_model_path):
    model = tf.keras.models.load_model(tf_model_path)
    spec = (tf.TensorSpec((1, 1024, 1024, 3), tf.float32, name="input"),)
    model_proto, _ = tf2onnx.convert.from_keras(model, input_signature=spec, opset=13, output_path=onnx_model_path)
    return onnx_model_path

 

위의 모델은 Tensorflow2 모델의 Input Shape : (1024, 1024, 3)인 상황입니다.

 

.onnx 파일에 대한 Inference 코드는 다음과 같습니다.

 

import onnxruntime as ort
import numpy as np

def do_inference_onnx(session, input_data):
    input_name = session.get_inputs()[0].name
    output_name = session.get_outputs()[0].name
    input_data = {input_name: input_data}
    output = session.run([output_name], input_data)
    return output
    
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
ort_session = ort.InferenceSession(onnx_model_path, providers=providers)
input_data = np.random.random((1, 1024, 1024, 3)).astype(np.float32)

for _ in range(100):
    onnx_outputs = do_inference_onnx(ort_session, input_data)

---

 

이어서, 변환된 .onnx 파일을 .trt 파일로 변환하는 코드에 대해서 알아보겠습니다.

 

imort os

def convert_onnx_to_trt(onnx_model_path, trt_engine_path):
    command = f"trtexec --onnx={onnx_model_path} --saveEngine={trt_engine_path} --verbose"
    os.system(command)
    return trt_engine_path

 

import pycuda.driver as cuda
import tensorrt as trt

def allocate_buffers(engine):
    inputs = []
    outputs = []
    bindings = []
    stream = cuda.Stream()

    for binding in engine:
        size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size
        dtype = trt.nptype(engine.get_binding_dtype(binding))
        host_mem = cuda.pagelocked_empty(size, dtype)
        device_mem = cuda.mem_alloc(host_mem.nbytes)
        bindings.append(int(device_mem))

        if engine.binding_is_input(binding):
            inputs.append((host_mem, device_mem))
        else:
            outputs.append((host_mem, device_mem))
    return inputs, outputs, bindings, stream
    
def load_trt_engine(trt_engine_path):
    with open(trt_engine_path, 'rb') as f:
        engine_data = f.read()
    runtime = trt.Runtime(TRT_LOGGER)
    if runtime is None:
        print("Failed to create TensorRT runtime")
    engine = runtime.deserialize_cuda_engine(engine_data)
    if engine is None:
        print("Failed to load TensorRT engine")
    return engine

# GPU device 설정
cuda.init()
device = cuda.Device(0) # gpu_id = 0 (int)
context = device.make_context()

""" TensorRT Engine 로드 """
engine = load_trt_engine(trt_engine_path)
context = engine.create_execution_context()

inputs, outputs, bindings, stream = allocate_buffers(engine)

""" Input Data를 설정 """
input_data = np.random.random_sample(inputs[0][0].shape).astype(np.float32)
np.copyto(inputs[0][0], input_data.ravel())

""" Inference """
for _ in range(1000):
    trt_outputs = do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)
    
cuda.Context.pop()

 

---

지금까지, Tensorflow 2 모델을 .onnx 파일 및 .trt 파일로 변환하고 Inference 하는 방법에 대해서 알아보았습니다.

감사드립니다.