文章目录
0. 前言
-
之前浏览过Python API并输出了笔记,但在实际使用过程中,上次的笔记没有任何卵用……
-
所以,本文根据 API 提供的几个功能,分别介绍相关API以及实例,希望下次用到TensorRT的时候,可以直接在这里复制粘贴。
-
资料:
- 官方samples:
- 在官方Github都能看到,这里还列出了Python实例。
- Github: ensorrt-demo:TRT内容不多,有不少 TF 相关的
- Github: tensorrt-utils:有一系列TRT实例,包括 OSS/inference/int8/network/onnx/plugins/uff。
- Githbu: tensorrt-sample:提供了一个以TensorFlow PB作为输入的,转为UFF后,进行推理的实例。
- 官方samples:
-
目前进展
- 基本概念,介绍 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)
- 判断 binding 类型(是否是input类型):
- 另外,可通过
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
执行模型推理。
- 通过 Runtime 读取 engine 文件,创建
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()
结果。
- API 形式是
- 从Host到Device是通过
pycuda.driver.cuda.memcpy_htod
,API可以参考这里- API的形式是
memcpy_htod(dest, src)
,dest是mem_alloc
的结果,src 是numpy/pagelocked_empty
。
- API的形式是
- 从Device到Host是通过
pycuda.driver.cuda.memcpy_dtoh
,API可以参考这里- API的形式是
memcpy_dtoh(dest, src)
,dest是numpy/pagelocked_empty
,src是mem_alloc
。
- API的形式是
- 模型推理,即
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