yolov8实战第三天——yolov8TensorRT部署（python推理）（保姆教学）

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
 
def load_plan(engine_path):
    # 加载TensorRT引擎文件
    with open(engine_path, "rb") as f, trt.Runtime(trt.Logger()) as runtime:
        engine = runtime.deserialize_cuda_engine(f.read())
    return engine
 
def infer(engine_path, input_image):
    # 加载TensorRT引擎
    engine = load_plan(engine_path)
    context = engine.create_execution_context()
 
    # 设置输入数据
    # 假设已经有了一个预处理过的input_image，并且有输入维度信息
    # 输入维度信息需要根据实际模型进行设置
    input_dims = (1, 3, 640, 640)  # 假设的输入维度 (batch_size, channels, height, width)
    input_img = input_image.ravel()  # 假设input_image是一个numpy数组，并且已经预处理过
    input_host, input_device = cuda.mem_alloc(1228800), cuda.In(1228800)  # 为输入数据分配GPU内存
    cuda.memcpy_htod(input_device, np.ascontiguousarray(input_img, dtype=np.float32))
 
    # 获取输出维度信息
    output_dims = engine.get_binding_shape(1)  # 假设输出维度是绑定在索引1的
 
    # 执行推理
    context.execute_v2([int(input_device), cuda.InOut(output_dims)])
    cuda.memcpy_dtoh(input_host, output_device)
 
    # 处理输出
    output_array = np.reshape(input_host, output_dims)
    # 输出结果处理可以根据具体的模型和需求进行
    # ...
 
    return output_array
 
# 使用方法
# engine_path 是 TensorRT 引擎文件的路径
# input_image 是需要进行推理的图片或者数据
output_array = infer(engine_path, input_image)

这段代码提供了一个简化的模板，用于加载TensorRT引擎，设置输入数据，执行推理，并处理输出结果。注意，这里的输入维度和输出维度信息都是假设的，需要根据实际的模型进行设置和调整。同时，输出结果的处理也需要根据具体的模型和需求进行相应的操作。