TensorRT开发者手册（3）使用TensorRT的PythonAPI
注意：原文中所有超链接均已更新，部分链接可能需要科学上网才能访问。部分单词我觉得翻译成中文总是缺少点意思，所以直接保留！由于本人实力有限，有错误之处希望指明，谢谢。–2020/11/25
Chapter 3. Using The Python API
该部分主要讲的是NVIDIA®TensorRT™用户目标和任务以及如何使用Python API来实现。
这些部分集中在使用Python API。 更多细节在Samples Support Guide。
假设你从一个已经训练好的模型开始。 本章将涵盖使用TensorRT的以下必要步骤：
- 从模型中创建TensorRT网络定义
- 调用TensorRT生成器从网络创建优化的运行时引擎
- 序列化和反序列化引擎，以便在运行时快速重新创建引擎
- 向引擎提供数据以执行推理
Python API vs C++ API
实际上，C++ API和Python API在支持您的需求方面应该接近相同。 Python API的主要好处是数据预处理和后处理很容易使用，因为您可以使用各种库，如NumPy和SciPy。
C++ API应该在安全很重要的情况下使用，例如在汽车中。 有关C++ API的更多信息，参见使用C++ API章节。
有关如何使用Python优化性能的更多信息，请参见来自TensorRT最佳实践指南的如何优化Python性能 。
3.1. Importing TensorRT Into Python
Procedure
- Import TensorRT:
import tensorrt as trt
- 实现一个日志接口，TensorRT通过该接口报告错误、警告和信息消息。 下面的代码展示了如何实现日志接口。 在这种情况下，我们抑制了信息消息，并且只报告警告和错误。 TensorRT Python bindings中包含了一个简单的记录器。
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
3.2. Creating A Network Definition In Python
使用TensorRT执行推理的第一步是从模型中创建一个TensorRT网络。
实现这一点的最简单方法是使用TensorRT解析器库导入模型(参见使用Python中的解析器导入模型、使用Python从Caffe导入模型、使用Python从TensorFlow导入模型和使用Python从ONNX导入模型)，它支持以下格式的序列化模型：
- Caffe (both BVLC and NVCaffe)
- Supports ONNX releases up to ONNX 1.6, and ONNX opsets 7 to 11, and
- UFF (used for TensorFlow)
另一种方法是直接使用TensorRT Network API定义模型(详见3.2.1Creating A Network Definition From Scratch Using The Python API)。 这需要您进行少量的API调用，以定义网络图中的每个层，并实现您自己对模型训练参数的导入机制。
注意：TensorRT PythonAPI不适用于所有平台。 有关更多信息，请参见TensorRT支持矩阵。
3.2.1. Creating A Network Definition From Scratch Using The Python API
在创建网络时，必须首先定义引擎并创建用于推理的生成器对象。 Python API用于从Network API创建网络和引擎。 网络定义引用用于向网络添加各种层。
About this task
有关使用PythonAPI创建网络和引擎的更多信息，请参见"Hello World" For TensorRT Using PyTorch And Python (network_api_pytorch_mnist)。
下面的代码说明了如何创建一个简单的网络与输入，卷积，池，全连接，激活和SoftMax层。
# Create the builder and network
with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network:
# Configure the network layers based on the weights provided. In this case, the weights are
imported from a pytorch model.
# Add an input layer. The name is a string, dtype is a TensorRT dtype, and the shape can be
provided as either a list or tuple.
input_tensor = network.add_input(name=INPUT_NAME, dtype=trt.float32, shape=INPUT_SHAPE)
# Add a convolution layer
conv1_w = weights['conv1.weight'].numpy()
conv1_b = weights['conv1.bias'].numpy()
conv1 = network.add_convolution(input=input_tensor, num_output_maps=20, kernel_shape=(5, 5),
kernel=conv1_w, bias=conv1_b)
conv1.stride = (1, 1)
pool1 = network.add_pooling(input=conv1.get_output(0), type=trt.PoolingType.MAX,
window_size=(2, 2))
pool1.stride = (2, 2)
conv2_w = weights['conv2.weight'].numpy()
conv2_b = weights['conv2.bias'].numpy()
conv2 = network.add_convolution(pool1.get_output(0), 50, (5, 5), conv2_w, conv2_b)
conv2.stride = (1, 1)
pool2 = network.add_pooling(conv2.get_output(0), trt.PoolingType.MAX, (2, 2))
pool2.stride = (2, 2)
fc1_w = weights['fc1.weight'].numpy()
fc1_b = weights['fc1.bias'].numpy()
fc1 = network.add_fully_connected(input=pool2.get_output(0), num_outputs=500, kernel=fc1_w,
bias=fc1_b)
relu1 = network.add_activation(fc1.get_output(0), trt.ActivationType.RELU)
fc2_w = weights['fc2.weight'].numpy()
fc2_b = weights['fc2.bias'].numpy()
fc2 = network.add_fully_connected(relu1.get_output(0), OUTPUT_SIZE, fc2_w, fc2_b)
fc2.get_output(0).name =OUTPUT_NAME
network.mark_output(fc2.get_output(0))
3.2.2. Importing A Model Using A Parser In Python
要使用解析器导入模型，需要执行以下高级步骤：
-
- 创建TensorRT builder和network。
-
- 为特定格式创建TensorRT parser。
-
- 使用解析器解析导入的模型并填充网络。
有关更加详细的步骤，请参见使用Python从Caffe导入(3.2.3)、使用Python从TensorFlow导入(3.2.4)和使用Python从ONNX导入(3.2.5)。
builder必须在network之前创建，因为它是网络的工厂。 不同的解析器有不同的机制来标记网络输出。 有关更多信息，请参见UFF Parser API、 Caffe Parser API和 ONNX Parser API。
3.2.3. Importing From Caffe Using Python
下面的步骤说明如何直接使用CaffeParser和PythonAPI导入Caffe模型。
About this task
有关更多信息，请参见Introduction To Importing Caffe, TensorFlow And ONNX Models
Into TensorRT Using Python (introductory_parser_samples)。
Procedure
- Import TensorRT.
import tensorrt as trt
- 定义数据类型。 在本例中，我们将使用float32。
datatype = trt.float32
- 此外，定义一些路径。 将路径改为您存放模型的位置：
deploy_file = 'data/mnist/mnist.prototxt'
model_file = 'data/mnist/mnist.caffemodel'
- 创建builder、network和parser：
with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network,
trt.CaffeParser() as parser:
model_tensors = parser.parse(deploy=deploy_file, model=model_file, network=network,
dtype=datatype)
解析器返回model_tensors，这是一个包含从张量名称到ITensor object的映射。
3.2.4. Importing From TensorFlow Using Python
下面的步骤说明如何直接使用UffParser和PythonAPI导入TensorFlow模型。
About this task
此示例可以在/tensorrt/samples/python/end_to_end_tensorflow_mnist目录中找到。 有关更多信息，请参见 “Hello World” For TensorRT Using TensorFlow And Python (end_to_end_tensorflow_mnist)。
Procedure
- Import TensorRT:
import tensorrt as trt
-
为tensorflow model创建一个frozen TensorFlow model。 将TensorFlow model 冻结到stream中的说明可以在 Freezing A TensorFlow Graph(14.1.2)中找到。
-
使用UFF转换器将冻结的tensorflow模型转换为UFF文件。 通常，这很简单：
convert-to-uff frozen_inference_graph.pb
取决于如何安装Tensor RT， convert-to-uff实用程序可能不会安装在系统路径中。 在本例中，直接调用底层Python脚本。 它应该位于UFF模块的bin目录中；例如~/.local/lib/python3.6/site-packages/uff/bin/convert_to_uff.py.
若要查找UFF模块的位置，请运行python-c "import uff; print(uff.path)"命令。
或者，您可以使用 UFF Parser API并直接转换TensorFlow GraphDef。
- 定义一些路径。 更改路径成您模型对应的位置：
model_file = '/data/mnist/mnist.uff'
- 创建生成器、网络和解析器：
with builder = trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network,
trt.UffParser() as parser:
parser.register_input("Placeholder", (1, 28, 28))
parser.register_output("fc2/Relu")
parser.parse(model_file, network)
3.2.5. Importing From ONNX Using Python
下面的步骤说明如何直接使用OnnxParser和PythonAPI导入ONNX模型。
About this task
有关更多信息，请参见 Introduction To Importing Caffe, TensorFlow And ONNX Models
Into TensorRT Using Python (introductory_parser_samples)。
注：
一般来说，新版本的Onnx Parser被设计为向后兼容，因此，遇到由早期版本的ONNX导出器生成的模型文件不应该引起问题。 当更改不向后兼容时，可能会有一些异常。 在这种情况下，将早期的ONNX模型文件转换为以后支持的版本。 有关此主题的更多信息，请参见ONNX Model Opset Version Converter。
用户模型也有可能是由支持后期操作集的导出工具生成的，而不是由Tensor RT附带的ONNX parser支持的。 在这种情况下，检查发布到GitHub的TensorRT的最新版本onnx-tensorrt是否支持所需的版本。 有关更多信息，请参阅Object Detection With The ONNX TensorRT Backend In Python (yolov3_onnx)。
支持的版本由oonnx_trt_backend.cpp中的BACKEND_OPSET_VERSION变量定义。 从GitHub下载并构建最新版本的ONNX TensorRT Parser。 构建说明可以在这里找到：TensorRT backend for ONNX。
在Tensor RT7.0中，ONNX解析器只支持full-dimensions mode，这意味着必须使用 explicitBatch flag set创建网络定义。 有关更多信息， Working With Dynamic Shapes.（chapter 7）
Procedure
- Import TensorRT:
import tensorrt as trt
- 创建 builder, network, and parser:
EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
with trt.Builder(TRT_LOGGER) as builder, builder.create_network(EXPLICIT_BATCH) as
network, trt.OnnxParser(network, TRT_LOGGER) as parser:
with open(model_path, 'rb') as model:
if not parser.parse(model.read()):
for error in range(parser.num_errors):
print(parser.get_error(error))
3.2.6. Importing From PyTorch And Other Frameworks
About this task
使用带有PyTorch的TensorRT(或具有NumPy兼容权重的任何其他框架)涉及使用TensorRT API复制 network architecture(参见 Creating A Network Definition From Scratch Using The Python API（3.2.1）)，然后从PyTorch复制权重。 有关更多信息，请参见 Working With PyTorch And Other Frameworks（14.2）。
要执行推理，请查看Performing Inference In Python.（3.5）。
3.3. Building An Engine In Python
builder的功能之一是搜索其CUDA内核目录，以获得最快的实现，因此有必要使用相同的GPU来构建，就像优化引擎将在其上运行一样。（这里的翻译得很拉跨，大概意思就是在某个固定型号的GPU上生成的engine，在其他型号上可能不能运行或者效果不好。）
About this task
IBuilder Config有许多属性，您可以设置这些属性来控制网络运行的精度，以及自校正参数，例如Tensor RT在确定哪个最快时应该给每个内核计时多少次（更多的迭代会导致更长的运行时间，但对噪声的敏感性较低。） 您还可以查询构建器以了解硬件本地支持哪些混合精度类型。
一个特别重要的属性是最大工作空间大小。
- 层算法往往需要临时工作区。 此参数限制网络中任何层都可以使用的最大大小。 如果提供了不足的划痕，TensorRT可能无法为给定的层找到实现。
有关在Python中构建引擎的更多信息，请参见 Introduction To Importing Caffe, TensorFlow And ONNX Models Into TensorRT Using Python (introductory_parser_samples)。
Procedure
- 使用builder对象构建引擎：
with trt.Builder(TRT_LOGGER) as builder, builder.create_builder_config() as config:
config.max_workspace_size = 1 << 20 # This determines the amount of memory available
to the builder when building an optimized engine and should generally be set as high as
possible.
with builder.build_engine(network, config) as engine:
# Do inference here.
当engine创建时，TensorRT将会复制weights。
- 进行推理。 要执行推理，请遵循Python中执行Inference中概述的说明（3.5）。
3.4. Serializing A Model In Python
从这里开始，您可以序列化引擎，也可以直接使用引擎进行推理。 在将模型用于推理之前，序列化和反序列化是一个可选步骤，如果需要，引擎对象可以直接用于推理。
About this task
当您序列化时，您正在将引擎转换为一种格式，以便在稍后的时间存储和使用以进行推理。 要用于推理，您只需将引擎反序列化即可。 序列化和反序列化是可选的。 由于从Network Definition中创建引擎可能很耗时，所以每次应用程序重新运行时，您可以避免通过序列化一次并在引用时反序列化来重建引擎。 因此，在构建引擎之后，用户通常希望将其序列化以供以后使用。
注意：序列化引擎不能跨平台或TensorRT版本移植。 引擎是特定的GPU模型，他们是建立在(除了平台和TensorRT版本)。
- 将模型序列化为modelstream：
serialized_engine = engine.serialize()
- 设计模型流来执行推理。 反序列化需要创建运行时对象：
with trt.Runtime(TRT_LOGGER) as runtime:
engine = runtime.deserialize_cuda_engine(serialized_engine)
还可以将序列化引擎保存到文件中，并从文件中读取：
- 序列化引擎并写入文件：
with open(“sample.engine”, “wb”) as f:
f.write(engine.serialize())
- 从文件中读取引擎并反序列化：
with open(“sample.engine”, “rb”) as f, trt.Runtime(TRT_LOGGER) as runtime:
engine = runtime.deserialize_cuda_engine(f.read())
3.5. Performing Inference In Python
下面的步骤说明了如何在Python中执行推理，现在您有了一个引擎。
Procedure
- 为输入和输出分配一些host 和device缓冲区。 此示例假设context.all_binding_dimensions == True，引擎在binding_index == 0处有一个输入，在binding_index==1处有一个输出：
# Determine dimensions and create page-locked memory buffers (i.e. won't be swapped to
# disk) to hold host inputs/outputs.
h_input = cuda.pagelocked_empty(trt.volume(context.get_binding_shape(0)),
dtype=np.float32)
h_output = cuda.pagelocked_empty(trt.volume(context.get_binding_shape(1)),
dtype=np.float32)
# Allocate device memory for inputs and outputs.
d_input = cuda.mem_alloc(h_input.nbytes)
d_output = cuda.mem_alloc(h_output.nbytes)
# Create a stream in which to copy inputs/outputs and run inference.
stream = cuda.Stream()
- 创建一些空间来存储中间激活值。 由于引擎保存网络定义和训练的参数，因此需要额外的空间。 这些都是在执行上下文中保存的：
with engine.create_execution_context() as context:
# Transfer input data to the GPU.
cuda.memcpy_htod_async(d_input, h_input, stream)
# Run inference.
context.execute_async_v2(bindings=[int(d_input), int(d_output)],
stream_handle=stream.handle)
# Transfer predictions back from the GPU.
cuda.memcpy_dtoh_async(h_output, d_output, stream)
# Synchronize the stream
stream.synchronize()
# Return the host output.
return h_output
引擎可以有多个执行上下文，允许一组权重用于多个重叠推理任务。 例如，您可以使用一个引擎和每个流的一个上下文来处理并行CUDA流中的图像。 每个上下文将在与引擎相同的GPU上创建。