[TensorRT 변환] ONNX2TensorRT(transformation and inference)

TensorRT serialize & deserialize code with using dynamic batch

 

파이썬으로 Onnx 파일을 불러와 engine파일로 직렬화하여 변환하여 저장하고, (다이나믹 배치 포함) 다시 그 파일을 불러와서 inferenc하는 코드이다.

 

TensorRT Version : 8.2.2.1

 

전에 pytorch모델을 onnx로 변환하는 방법은 아래 포스팅에 pytorch method 옵션과 함께 간략히 적어뒀다.

torch.onnx.export() option

 

공부하면서 자료가 많이 없어서 gpt의 도움을 많이 받았다. gpt가 한 번에 오류없는 코드를 만들어주진 않지만, trt버전 바꿔달라면 바꿔주고 오류가 발생하면 해당 오류에대한 해결책도 알려주었다.

만능은 아니지만 시간을 엄~청 단축해줬다.

 

직렬화 : Serialize 

import tensorrt as trt

trt logger, builder 생성 및 onnx weight load

# Create a builder object
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)

# Set the maximum batch size
max_batch_size = 32
builder.max_batch_size = max_batch_size

network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, logger) # network meaning : The network definition to which the parser will write.

# Parse ONNX model
#with open("weights/resnet_CIFAR10_epoch600_dynamic.onnx", "rb") as f:
#    onnx_model = f.read()
#success = parser.parse(onnx_model)
success = parser.parse_from_file('weights/resnet_CIFAR10_epoch600_dynamic.onnx')
if not success:
    raise Exception('ERROR: Failed to parse ONNX model')

config 생성

profile은 변환시 dynamic batch, image size를 적용할때 어떻게 최적화할지 명시해준다.

# Set config 
config = builder.create_builder_config()
# Set precision of model
config.set_flag(trt.BuilderFlag.FP16)  ################# FP16
# Set the workspace size
max_workspace_size = 1 << 30  # 1GB
config.max_workspace_size = max_workspace_size
# Create a profile for the engine
profile = builder.create_optimization_profile()
profile.set_shape("input", (1, 3, 32, 32), (max_batch_size, 3, 32, 32), (max_batch_size, 3, 32, 32)) ############### dynamic batch
config.add_optimization_profile(profile)

 

위에서 설정한 옵션들로 engine을 build한다. 주석친 부분으로 해도 작동하는데 경고문이 떠서 아래로 바꾸어주었다.

# Create an engine object from the parsed ONNX model
# engine = builder.build_engine(network=network, config=config) # type(engine) -> tensorrt.tensorrtt.ICudaEngine
# if engine is None:
#    print("ERROR: Failed to build TensorRT engine")
    
# Serialize the engine to a file
#serialized_engine = engine.serialize() # return IHostMemory
serialized_engine = builder.build_serialized_network(network=network, config=config) # type(serialized_engine) -> bytes
with open("modelfp16.engine", "wb") as f:
    f.write(serialized_engine)

 

역직렬화 deserialize & Inference

import tensorrt as trt
import numpy as np
import pycuda.driver as cuda
import pycuda.autoinit

추론은 gpu에서 하기때문에 pycuda가 필수이다.

pycuda 개념 및 사용 예시

모델 weight load & deserialize

# Create a CUDA stream
stream = cuda.Stream()

# Load the serialized engine from the file
with open("model.engine", "rb") as f:
    serialized_engine = f.read()
    
# Create a runtime object and engine object from the serialized engine
runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING))
engine = runtime.deserialize_cuda_engine(serialized_engine)

 

추론을 위한 dummy data 생성

# Get the binding indices for the input and output tensors
input_index = engine.get_binding_index(input_name)
output_index = engine.get_binding_index(output_name)

# Get the original input and output tensor shapes
input_shape = engine.get_binding_shape(input_index)
output_shape = engine.get_binding_shape(output_index)

# Set the desired dynamic batch size
batch_size = 8
input_shape = (batch_size, *input_shape[1:])
output_shape = (batch_size, *output_shape[1:])
# Create new numpy arrays with the desired dynamic batch size
input_data = np.random.rand(*input_shape).astype(np.float32)
output_data = np.empty(output_shape, dtype=np.float32)
print("Input shape:", input_data.shape)
print("Output shape:", output_data.shape)

batch size는 8로 지정했지만 max batch size를 32로 지정했기때문에 32 이하의 다른 숫자도 가능하다.

 

생성한 dummy 데이터를 gpu 메모리로 보낸다.

먼저 메모리에 공간을 할당하고 해당 위치로 보낸다.

# Allocate device memory for the input and output
input_buf = cuda.mem_alloc(input_data.nbytes)
output_buf = cuda.mem_alloc(output_data.nbytes)
# print("Input buffer loc :", int(input_buf))
# print("output buffer loc : ", int(output_buf))

# Transfer input data to device
cuda.memcpy_htod_async(input_buf, input_data, stream)

# Create a context object and set the input and output bindings
context = engine.create_execution_context()
context.set_binding_shape(input_index, input_shape)
context.set_binding_shape(output_index, output_shape)

bindings = [int(input_buf), int(output_buf)]

# Run inference
context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)
cuda.memcpy_dtoh_async(output_data, output_buf, stream)

# Synchronize the stream and print the results
stream.synchronize()
print(output_data)
idx = [np.argmax(output) for output in output_data]
print(idx)