TensorRT Python API笔记

基本概念
1 Logger
API
概述：为Builder/ICudaEngine/Runtime对象提供logger
创建所需参数：min_severity，参数就是 trt.Logger.INTERNAL_ERROR/WARNING/ERROR/VERBOSE 等
成员变量：无
成员函数：log(severity, msg)
其他作用：无

2 Builder
API
概述：通过 INetworkDefinition 对象创建 ICudaEngine 对象。
创建所需参数：Logger对象
成员变量：模型相关参数，如max_batch_size/max_workspace_size/int8_mode/fp16_mode等。
成员函数：
创建 INetworkDefinition 对象，如create_network(flags)
通过 INetworkDefinition 创建 ICudaEngine，如 build_cuda_engine(network)/build_engine(network, config)
其他作用：还有创建 builder_config以及optimization_profile。相关功能暂时用不到，所以只是了解一下。

3 Runtime
API
概述：反序列化Engine文件，换句话说，就是解析本地engine文件，创建 ICudaEngine 对象。
创建所需参数：Logger
成员变量：无
成员函数：deserialize_cuda_engine(serialized_engine, plugin_factory)，其中前者就是 open(filename, “rb”).read() 的结果。
其他作用：无

4 ICudaEngine
API
概述：就是一个 TensorRT Engine 对象了，可以理解为一个模型以及相关参数
创建所需参数：
可通过 Builder 的 build_cuda_engine/build_engine 创建。
可通过 Runtime 的 deserialize_cuda_engine 创建。
成员变量：模型相关参数，主要包括 num_bindings/max_batch_size/num_layers/max_workspace_size 等。
成员函数：
创建 IExecutionContext 对象，例如 create_execution_context()或 create_execution_context_without_device_memory()
序列化 Engine，serialize，大概用法就是 open(filename, “wb”).write(engine.serialize())
这里要单独介绍一下 binding 相关内容
概念：可理解为 端口，用于表示输入tensor与输出tensor。
对应类 pycuda.driver.cuda.mem_alloc
可通过id或name获取对应的binding
作用：在后续模型推理过程中，需要以 bindings 作为输入，其具体数值为内存地址，即 int(buffer)。
ICudaEngine 相关函数包括：
判断 binding 类型（是否是input类型）：binding_is_input(idx/name)
获取shape与dtype：get_binding_shape(idx/name)和get_binding_dtype(idx/name)
根据 id 获取 name，根据 name 获取 id，get_binding_shape/get_binding_name
其他不懂的方法
get_binding_bytes_per_component(idx)
get_binding_components_per_element(idx)
get_binding_format(idx)
get_binding_format_desc(idx)
get_binding_vectorized_dim(idx)
is_execution_binding(idx)
is_shape_binding(idx)
另外，可通过 for binding_name in engine: 遍历获取所有 binding_name

5 IExecutionContext
API
概述：模型推理上下文
创建所需参数：
通过 ICudaEngine.create_execution_context() 获取对象
成员变量：推理相关参数，如 profiler/engine/name 等。
成员函数：
主要就是执行推理的方法 execute/execute_v2/execute_async/execute_async_v2
不太清楚 v1 v2 有什么区别
官方sample中，v1的注解是 This function is generalized for multiple inputs/outputs.，v2的注解是 This function is generalized for multiple inputs/outputs for full dimension networks.，但不太懂
有 get_shape/get_binding_shape/set_shape_input/set_binding_shape 等方法
其他作用：我也不知道剩下函数干什么用的 get_strides/set_optimization_profile_async

推理
1 相关API详解
推理的整体流程是
模型解析与优化，即将 ONNX/UFF/CAFFE 等形式转换为 Engine。本节不考虑这个。
模型推理，即以 engine 文件作为输入，实现模型推理的基本流程。
相关 API 主要包括：
通过 Runtime 读取 engine 文件，创建 tensorrt.ICudaEngine 对象。
为输入与输出分配内存与显存，通过 pycuda 实现
通过 tensorrt.ICudaEngine 对象构建模型推理所需的 tensorrt.IExecutionContext 对象。
通过 tensorrt.IExecutionContext 执行模型推理。
tensorrt.ICudaEngine 对象的创建
可通过 tensorrt.Builder 实现，主要就是通过 INetworkDefinition 对象实现。
后面不会介绍这种形式，具体详情可参考 Builder API
可通过 tensorrt.Runtime 实现，具体就是通过一个 buffer 作为输入（如本地文件）。
具体方法就是 deserialize_cuda_engine，具体可参考文档。
tensorrt.IExecutionContext 对象的创建
主要就是通过 tensorrt.ICudaEngine 对象的 create_execution_context 方法。
内存分配
主要功能就是：将内存中（Host）输入数据转存到显存中（Device），将显存中（Device）的推理结果保存到内存（Host）中。
对于每一个输入张量与输出变量，都需要分配两块资源，分别是内存（Host）中的的资源以及显存（Device）中的资源。
在内存（Host）中分配空间通过 pycuda.driver.cuda.pagelocked_empty，API可以参考这里
也可以直接通过 numpy 实现
API 形式是 pagelocked_empty(shape, dtype)，主要参数就是shape和dtype。
shape 一般通过 trt.volume(engine.get_binding_shape(id))实现，可以理解为元素数量（而不是内存大小）
dtype就是数据类型，可以通过 np.float32 或 trt.float32 的形式。
在显存（Device）中分配空间通过 pycuda.driver.cuda.mem_alloc，API可以参考这里
API 形式是 mem_alloc(buffer.nbytes)，其中 buffer 可以是ndarray，也可以是前面的 pagelocked_empty() 结果。
从Host到Device是通过 pycuda.driver.cuda.memcpy_htod，API可以参考这里
API的形式是 memcpy_htod(dest, src)，dest是 mem_alloc 的结果，src 是 numpy/pagelocked_empty。
从Device到Host是通过 pycuda.driver.cuda.memcpy_dtoh，API可以参考这里
API的形式是memcpy_dtoh(dest, src)，dest是numpy/pagelocked_empty，src是mem_alloc。
模型推理，即 IExecutionContext对象的 execute系列方法
有四个方法 execute/execute_v2/execute_async/execute_async_v2
四个方法都有 batch_size, bindings 两个参数。异步方法还有 stream_handle/input_consumed 两个参数
bindings 是一个数组，包含所有input/outpu buffer（也就是device）的地址。获取方式就是直接通过 int(buffer)，其中 buffer 就是 mem_alloc 的结果。
stream_handle 是 cuda.Stream() 对象

实例
读取 Engine

# 输入 Engine 本地文件构建 ICudaEngine 对象
ENGINE_PATH = '/path/to/model.trt'
trt_logger = trt.Logger(trt.Logger.INFO)
runtime = trt.Runtime(trt_logger)
with open(ENGINE_PATH, "rb") as f:
    engine = runtime.deserialize_cuda_engine(f.read())
# 输入 ONNX/UFF/CAFFE 获取 ICudaEngine 对象

模型推理准备工作（构建Context、各种Buffer以及Bindings）

# 构建 context
context = engine.create_execution_context()

# 构建 buffer 方式一：如果确定只有一个输入一个输出
# 参考 https://github.com/dkorobchenko-nv/tensorrt-demo/blob/master/trt_infer.py
INPUT_DATA_TYPE = np.float32
stream = cuda.Stream()
host_in = cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(0)), dtype=INPUT_DATA_TYPE)
host_out = cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(1)), dtype=INPUT_DATA_TYPE)
devide_in = cuda.mem_alloc(host_in.nbytes)
devide_out = cuda.mem_alloc(host_out.nbytes)
bindings = [int(devide_in), int(devide_out)]

# 构建 buffer 的方式二：如果不知道有多少输入多少输出
# 参考 https://github.com/NVIDIA/TensorRT/blob/master/samples/python/common.py
class HostDeviceMem(object):
    def __init__(self, host_mem, device_mem):
        self.host = host_mem
        self.device = device_mem

    def __str__(self):
        return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device)

    def __repr__(self):
        return self.__str__()
inputs = []
outputs = []
bindings = []
stream = cuda.Stream()
for binding in engine:
    # 注意，上面循环得到的是 binding_name
    size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size
    dtype = trt.nptype(engine.get_binding_dtype(binding))
    # Allocate host and device buffers
    host_mem = cuda.pagelocked_empty(size, dtype)
    device_mem = cuda.mem_alloc(host_mem.nbytes)
    # Append the device buffer to device bindings.
    bindings.append(int(device_mem))
    # Append to the appropriate list.
    if engine.binding_is_input(binding):
        inputs.append(HostDeviceMem(host_mem, device_mem))
        else:
            outputs.append(HostDeviceMem(host_mem, device_mem))

模型推理

# 如果输入输出已经确定
np.copyto(host_in, img.ravel())
cuda.memcpy_htod_async(devide_in, host_in, stream)
context.execute_async(bindings=bindings, stream_handle=stream.handle)
cuda.memcpy_dtoh_async(host_out, devide_out, stream)
stream.synchronize()

# 如果输入输出数量不一定
# 参考 https://github.com/NVIDIA/TensorRT/blob/master/samples/python/common.py
# Transfer input data to the GPU.
[cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]
# Run inference.
context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)
# Transfer predictions back from the GPU.
[cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]
# Synchronize the stream
stream.synchronize()
# Return only the host outputs.
return [out.host for out in outputs]