NVIDIA-TensorRT-Python推理

1,前言

        NVIDIA TensorRT进行模型推理的Python实现。TensorRT是一个高性能的深度学习推理优化器和运行时，它能够为深度学习模型提供低延迟和高吞吐量的推理能力。(由于官方文档的使用还是比较简单，也可能自己很菜，参考了别人的文档和自己摸索，写出来这个可以使用的API)

2. Python-API推理

step1：导入基本库(环境自行配置)

# 导入TensorRT库
import tensorrt as trt
# 导入PyCUDA的驱动程序接口，用于与GPU通信和管理CUDA资源
import pycuda.driver as cuda
# 导入PyCUDA的自动初始化模块，它会在脚本开始时自动初始化CUDA环境
import pycuda.autoinit

step2：加载模型

# 创建TensorRT日志对象，并设置其日志级别为WARNING，只记录警告和更严重级别的信息
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
# 创建TensorRT运行时环境，它将用于加载序列化的引擎文件
runtime = trt.Runtime(TRT_LOGGER)
# 以二进制读取模式打开预编译的TensorRT引擎文件，使用runtime对象反序列化这个文件到一个engine对象
with open(engine_path, 'rb') as f:
    engine = runtime.deserialize_cuda_engine(f.read())

step3.1：获取输入和输出并推理

# 使用engine创建执行上下文context，这是执行推理所必需的      
with engine.create_execution_context() as context:
    # 初始化列表bindings，绑定输入和输出的CUDA内存地址，循环遍历engine的所有绑定
    bindings = []
    for binding in engine:
        # binding是模型的input和output节点名，本例是input.1和437
        # print(binding)
        # 获取当前绑定的索引
        binding_idx = engine.get_binding_index(binding)
        # 获取当前绑定的形状
        size = trt.volume(context.get_binding_shape(binding_idx))
        # 获取当前绑定的数据类型
        dtype = trt.nptype(engine.get_binding_dtype(binding))
        # 若绑定是input
        if engine.binding_is_input(binding):
            # 分配input数据的CPU锁页内存和GPU显存
            input_buffer = np.ascontiguousarray(image_data)
            # 分配输入数据的cuda显存
            input_memory = cuda.mem_alloc(image_data.nbytes)
            # 分配的地址添加到bindings列表
            bindings.append(int(input_memory))
        # 若绑定是output
        else:
            # 分配output数据的CPU锁页内存和GPU显存
            output_buffer = cuda.pagelocked_empty(size, dtype)
            # 分配输出数据的cuda显存
            output_memory = cuda.mem_alloc(output_buffer.nbytes)
            # 分配的地址添加到bindings列表
            bindings.append(int(output_memory))
    # 创建cuda流
    stream = cuda.Stream()
    # 将输入数据转入cuda
    cuda.memcpy_htod_async(input_memory, input_buffer, stream)
    # 执行推理
    context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)
    # 从GPU中将输出数据取出（output_buffer）
    cuda.memcpy_dtoh_async(output_buffer, output_memory, stream)
    # 同步流
    stream.synchronize()
# 输出buffer长度是(n*c*h*w),需要reshape
output = np.reshape(output_buffer, (1, 2, 584, 565))

step3.2 ：多输出

# 使用engine创建执行上下文context，这是执行推理所必需的      
with engine.create_execution_context() as context:
    # 初始化列表bindings，绑定输入和输出的CUDA内存地址，循环遍历engine的所有绑定
    bindings = []
    #缓冲区和内存地址添加到output_buffers和output_memories中
    output_memories = []
    output_buffers = []
    for binding in engine:
        # binding是模型的input和output节点名，本例是input.1和437
        # print(binding)
        # 获取当前绑定的索引
        binding_idx = engine.get_binding_index(binding)
        # 获取当前绑定的形状
        size = trt.volume(context.get_binding_shape(binding_idx))
        # 获取当前绑定的数据类型
        dtype = trt.nptype(engine.get_binding_dtype(binding))
        # 若绑定是input
        if engine.binding_is_input(binding):
            # 分配input数据的CPU锁页内存和GPU显存
            input_buffer = np.ascontiguousarray(image_data)
            # 分配输入数据的cuda显存
            input_memory = cuda.mem_alloc(image_data.nbytes)
            # 分配的地址添加到bindings列表
            bindings.append(int(input_memory))
        # 若绑定是output
        else:
            # 分配output数据的CPU锁页内存和GPU显存
            output_buffer = cuda.pagelocked_empty(size, dtype)
            # 分配输出数据的cuda显存
            output_memory = cuda.mem_alloc(output_buffer.nbytes)
            # 分配的地址添加到bindings列表
            bindings.append(int(output_memory))
            output_buffers.append(output_buffer)
            output_memories.append(output_memory)
    # 创建cuda流
    stream = cuda.Stream()
    # 将输入数据转入cuda
    cuda.memcpy_htod_async(input_memory, input_buffer, stream)
    # 执行推理
    context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)
    # 从GPU中将多输出数据取出
    for output_buffer, output_memory in zip(output_buffers, output_memories):
        cuda.memcpy_dtoh_async(output_buffer, output_memory, stream)
    # 同步流
    stream.synchronize()
# 输出buffer长度是(n*c*h*w),需要reshape
output_0 = np.reshape(output_buffers[0], (1, 100, 92))
output_1 = np.reshape(output_buffers[1], (1, 100, 4))

3. 注意

        (1) 测试使用FP32的trt模型，所以在输入数据时需要变成float32数据（减去均值和除方差)

        (2) 单输出也可以用多输出的代码,就是获取输出的时候不同

        (3) 后续使用INT8看看，再继续琢磨！！！