import tensorrt as trt

# 创建TensorRT的推理引擎
def create_trt_engine(model_file):
    with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.OnnxParser(network, TRT_LOGGER) as parser:
        builder.max_batch_size = 1
        builder.max_workspace_size= 1 << 30  # 1GB
        builder.fp16_mode = True  # 使用半精度浮点数加速推理
        builder.strict_type_constraints = True

        # 解析ONNX模型
        with open(model_file, 'rb') as model:
            parser.parse(model.read())

        # 构建并返回推理引擎
        return builder.build_cuda_engine(network)

# 加载TensorRT引擎
def load_trt_engine(engine_file):
    with open(engine_file, 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime:
        return runtime.deserialize_cuda_engine(f.read())

# 执行推理
def infer(engine, input_data):
    # 分配GPU内存
    inputs, outputs, bindings, stream = allocate_buffers(engine)

    # 将输入数据拷贝到GPU内存
    np.copyto(inputs[0].host, input_data.ravel())

    # 执行推理
    with engine.create_execution_context() as context:
        context.execute_v2(bindings)

    # 将输出数据从GPU内存拷贝回主机内存
    output_data = outputs[0].host

    return output_data

# 示例使用
# 创建或加载TensorRT引擎
if not os.path.exists(engine_file):
    engine = create_trt_engine(model_file)
    with open(engine_file, 'wb') as f:
        f.write(engine.serialize())
else:
    engine = load_trt_engine(engine_file)

# 执行推理
output = infer(engine, input_data)