【TensorRT python API Inference pipeline】

提示：TensorRT python API Inference pipeline 记录自用

---

前言

提示： Inference pipeline 推理管道（框架）：

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

# TODO: 在这里导入模型的预处理和后处理函数
# from your_module import preprocess, postprocess

class MyLogger(trt.Logger):
    def log(self, severity, msg):
        if severity >= self.verbosity:
            print("[TRT Log]: {}".format(msg))

logger = MyLogger(trt.Logger.Severity.VERBOSE)

def load_engine(engine_file_path):
    assert os.path.exists(engine_file_path), "Engine file does not exist."
    with open(engine_file_path, "rb") as f, trt.Runtime(logger) as runtime:
        return runtime.deserialize_cuda_engine(f.read())

def infer(engine, input_img):
    # TODO: 应用模型特定的预处理
    # processed_img = preprocess(input_img)

    with engine.create_execution_context() as context:
        # 分配缓冲区
        input_shape = (1, 3, 48, 160)
        output_shape = (1, 20, 78)
        dtype = np.float32
        
        # 分配输入和输出的内存
        input_size = int(np.prod(input_shape) * 4)
        output_size = int(np.prod(output_shape) * 4)