深入理解 TensorRT (1) TensorRT Python API 详解

毛声

2023-12-01

0. 前言

之前浏览过Python API并输出了笔记，但在实际使用过程中，上次的笔记没有任何卵用……
所以，本文根据 API 提供的几个功能，分别介绍相关API以及实例，希望下次用到TensorRT的时候，可以直接在这里复制粘贴。
资料：
- 官方samples：
  - 在官方Github都能看到，这里还列出了Python实例。
- Github: ensorrt-demo：TRT内容不多，有不少 TF 相关的
  - Engine 作为输入进行infer的实例
  - 通过uff将tf转tensorrt的实例
- Github: tensorrt-utils：有一系列TRT实例，包括 OSS/inference/int8/network/onnx/plugins/uff。
- Githbu: tensorrt-sample：提供了一个以TensorFlow PB作为输入的，转为UFF后，进行推理的实例。
目前进展
- 基本概念，介绍 Builder/Runtime/Logger/ICudaEngine/ICudaExecution 的基本概念、API、使用
- 推理相关 API 详解以及实例
- ONNX 模型转换
- Dynamic Shape
- 插件

1. 基本概念

下面分别介绍基本对象。
- 概述（包括基本功能）
- 创建所需参数
- 成员变量
- 成员函数
- 其他作用（作为别的基本对象的输入）

1.1 Logger

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

1.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。相关功能暂时用不到，所以只是了解一下。

1.3. Runtime

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

1.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

1.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

2. 推理

2.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() 对象

2.2 实例

读取 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]

3. ONNX 模型转换

TODO

4. Dynamic Shape

TODO

5. 插件

TODO

深入理解 TensorRT (1) TensorRT Python API 详解

0. 前言

1. 基本概念

1.1 Logger

1.2 Builder

1.3. Runtime

1.4 ICudaEngine

1.5 IExecutionContext

2. 推理

2.1 相关API详解

2.2 实例

3. ONNX 模型转换

4. Dynamic Shape

5. 插件

相关阅读

相关文章

相关问答

相关文档