原文地址:
https://developer.nvidia.com/zh-cn/blog/tensorrt-python-interface-cn/
(UPDATE: 我又发现上面这个原文其实是官方文档的翻译版…)
原作者:
作者主页:
https://developer.nvidia.com/zh-cn/blog/author/ken-he/
本文只转载其《4. TensorRT 的 Python 接口解析》,建议诸位去原文查看
在上一篇博客:
TensorRT windows10 安装过程记录
已经安装了TensorRT,对最后的代码只是有了感性的认识,本篇转载的博客,将介绍一些常用的简单的TensorRT API.
本文做了两件事:
- onnx 模型文件转为 trt 序列化文件
- 使用 trt 序列化模型进行推理
(前排提醒,本文阅读可能需要一些CUDA C编程的知识)
Python API官方文档:
https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/index.html
本章说明 Python API 的基本用法,假设您从 ONNX 模型开始.
Python API 可以通过tensorrt模块访问:
import tensorrt as trt
4.1. 构建过程(The Build Phase)
(额,读了依据,感觉不是正在语句顺序,好吧,大体意思懂了就行)
(我按照我的理解翻译成正常语序,诸位凑合看)
要创建构建器Builder
,需要首先创建一个记录器Logger
。
以下这行代码是一个简单的记录器Logger
实现,它将高于特定严重性的所有消息记录到stdout
。
logger = trt.Logger(trt.Logger.WARNING)
或者,可以通过从ILogger
类派生来定义您自己的记录器:
class MyLogger(trt.ILogger):
def __init__(self):
trt.ILogger.__init__(self)
def log(self, severity, msg):
pass # Your custom logging implementation here
logger = MyLogger()
然后,您可以创建一个构建器Builder
:
builder = trt.Builder(logger)
4.1.1. Creating a Network Definition in Python
4.1.1. 用Python创建一个网络定义
创建构建器Builder
后,优化模型的第一步是创建网络定义:
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
为了使用 ONNX 解析器导入模型,需要EXPLICIT_BATCH
标志。有关详细信息,请参阅显式与隐式批处理部分。
以下是 显式与隐式批处理 部分的中文版翻译链接:
4.1.2. Importing a Model using the ONNX Parser
4.1.2. 用ONNX解析器导入一个model
现在,需要从 ONNX 表示中填充网络定义。您可以创建一个 ONNX 解析器来填充网络,如下所示:
parser = trt.OnnxParser(network, logger)
然后,读取模型文件并处理可能存在的错误:
success = parser.parse_from_file(model_path)
for idx in range(parser.num_errors):
print(parser.get_error(idx))
if not success:
pass # Error handling code here
4.1.3. Building an Engine
4.1.3. 构建一个引擎(Engine)
下一步是创建一个构建配置,指定 TensorRT 应该如何优化模型:
config = builder.create_builder_config()
config 对象有很多属性,可以设置这些属性来控制 TensorRT 如何优化网络。
一个重要的属性是最大工作空间大小。网络中每一层的实现通常需要一个临时工作空间,此参数限制了网络中任何层可以使用的最大空间。如果提供的工作空间不足,TensorRT 可能无法进行该层的实现:(我猜是这个意思)
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 20) # 1 MiB
指定配置后,可以使用以下命令构建并将引擎序列化:
serialized_engine = builder.build_serialized_network(network, config)
将引擎保存到文件可以这样做:
with open("sample.engine", "wb") as f:
f.write(serialized_engine)
之前的步骤其实是将 onnx 模型导出为 trt 的序列化模型,接下来才是推理过程
4.2. 模型反序列化读取(Deserializing a Plan)
要执行推理,您首先需要使用Runtime接口反序列化引擎。与构建器Builder
一样,运行时需要记录器Logger
的实例。
runtime = trt.Runtime(logger)
然后,从内存缓冲区反序列化引擎(上一小节的变量):
engine = runtime.deserialize_cuda_engine(serialized_engine)
如果您需要首先从trt/engine
文件加载引擎,请运行:
with open("sample.engine", "rb") as f:
serialized_engine = f.read()
engine = runtime.deserialize_cuda_engine(serialized_engine)
4.3. Performing Inference
引擎已经拥有了优化的模型,但要执行推理需要额外的中间激活状态。
(我猜这句话的意思是,中间的推理过程时由 context 对象来执行的,我猜的啊,不一定对)
这是通过IExecutionContext接口完成的:
context = engine.create_execution_context()
一个引擎可以有多个执行上下文,允许一组权重用于多个重叠的推理任务。
(当目前有一个例外:当使用动态形状时,每个优化配置文件只能有一个执行上下文。)
要执行推理,必须为输入和输出传递 TensorRT 缓冲区,TensorRT 要求在 GPU 指针列表中指定。
您可以使用为输入和输出张量的名称 在引擎中 查询,以在数组中找到正确的位置:
input_idx = engine[input_name]
output_idx = engine[output_name]
使用这些索引,为每个输入和输出设置 GPU 缓冲区。多个 Python 包允许您在 GPU 上分配内存,包括但不限于 PyTorch、Polygraphy CUDA 包装器和 PyCUDA。
然后,创建一个 GPU 指针列表。例如,对于 PyTorch CUDA 张量,您可以使用data_ptr()方法访问 GPU 指针;对于 Polygraphy DeviceArray ,使用ptr属性:
buffers = [None] * 2 # Assuming 1 input and 1 output
buffers[input_idx] = input_ptr
buffers[output_idx] = output_ptr
(NOTICE: 也就是说,可以直接用PyTorch进行一些数据预处理)
插入一个小demo:
>>> import torch
>>> x = torch.randn((3, 3,))
>>> x
tensor([[-0.1134, -0.7201, 1.9786],
[ 0.2945, 1.4753, -0.2461],
[-0.3113, -0.9930, 0.4400]])
>>> x.data_ptr # Returns the address of the first element of self tensor.
<function Tensor.data_ptr>
>>> x.data_ptr()
2156338399168
填充输入缓冲区后,您可以调用 TensorRT 的execute_async
方法以使用 CUDA 流异步启动推理。
首先,创建 CUDA 流。如果您已经有 CUDA 流,则可以使用指向现有流的指针。例如,对于 PyTorch CUDA 流,即torch.cuda.Stream()
,您可以使用cuda_stream
属性访问指针;
插入 PyTorch 的 CUDA 流 demo:
>>> tr = torch.cuda.Stream()
>>> tr
<torch.cuda.Stream device=cuda:0 cuda_stream=0x1f61034aa10>
>>> tr.cuda_stream
2156345469456
对于 Polygraphy CUDA 流,使用ptr
属性。
接下来,开始推理:
context.execute_async_v2(buffers, stream_ptr)
通常在内核 (猜测是 CUDA 核程序?) 之前和之后将异步memcpy()
排入队列以从 GPU 中移动数据(如果数据尚不存在)。
要确定内核(可能还有memcpy() )何时完成,请使用标准 CUDA 同步机制,例如事件或等待流。
例如,对于 Polygraphy 和 PyTorch,使用:
stream.synchronize()
如果您更喜欢同步推理,请使用execute_v2
方法而不是execute_async_v2
。