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

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