python中的rt_3. Working With TensorRT Using The Python API

以下部分将重点介绍使用 Python API 可以执行 TensorRT 用户的目标和任务。这些部分主要讨论在没有任何框架的情况下使用 Python API。示例部分提供了进一步的详细信息，并在适当情况下链接到下面。

假设你从一个训练过的模型开始。本章将介绍以下使用TensorRT的必要步骤:

从模型中创建 TensorRT 网络定义

调用 TensorRT builder 从网络创建优化的运行时引擎

序列化和反序列化引擎，以便在运行时快速重新创建引擎

向引擎提供数据以执行推理

Python API vs C++ API

从本质上说，C++ API 和 Python API 在支持您的需求方面应该接近相同。Python API 的主要优点是数据预处理和后处理很容易使用，因为您可以使用各种库，如 NumPy 和 SciPy。

C++ API 应该用于安全非常重要的场合，例如在汽车中。有关 C++ API 的更多信息，请参见使用 C++ API 处理 TensorRT。

有关如何使用 Python 优化性能的更多信息，请参见如何优化我的 Python 性能?来自最佳实践指南。

导入 TensorRT:

import tensorrt as trt

实现一个日志接口，通过该接口 TensorRT 报告错误、警告和信息消息。下面的代码展示了如何实现日志记录接口。在这种情况下，我们抑制了信息消息，只报告警告和错误。TensorRT Python 绑定中包含了一个简单的日志记录器。

TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

使用 TensorRT 执行推理的第一步是从您的模型创建一个 TensorRT 网络。最简单的方法是使用 TensorRT 解析器库导入模型(请参阅使用 Python 中的解析器导入模型、使用 Python 从 Caffe 导入模型、使用 Python 从 TensorFlow 导入模型和使用 Python 从 ONNX 导入模型)，该库支持以下格式的序列化模型:

Caffe (both BVLC and NVCaffe)

ONNX 1.0 and 1.1, and

UFF (used for TensorFlow)

另一种选择是直接使用 TensorRT 网络 API 定义模型(请参阅使用Python API从头创建网络定义)。这需要您进行少量的 API 调用来定义网络图中的每一层，并为模型的训练参数实现您自己的导入机制。

注意: TensorRT Python API 仅适用于 x86_64 平台。更多信息请参见深度学习 SDK 文档- TensorRT 工作流。

在创建网络时，必须首先定义引擎并创建用于推理的构建器对象。Python API 用于从网络 API 创建网络和引擎。网络定义引用用于向网络添加各种层。有关使用 Python API 创建网络和引擎的更多信息，请参见 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))

要使用解析器导入模型，需要执行以下高级步骤:

创建TensorRTbuilder和网络。

为特定格式创建TensorRT解析器。

使用解析器解析导入的模型并填充网络。

有关这些步骤和示例代码的示例，请参见使用 Python 从 Caffe 导入、使用 Python 从 TensorFlow 导入和使用 Python 从 ONNX 导入。

构建器必须在网络之前创建，因为它是网络的工厂。不同的解析器有不同的网络输出标记机制。有关更多信息，请参见 UFF 解析器 API、Caffe 解析器 API 和 ONNX 解析器 API。

下面的步骤说明了如何使用 Caffe 解析器和 Python API 直接导入 Caffe 模型。有关更多信息，请参阅 introductory_parser_samples 示例。

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 对象的映射。

下面的步骤说明了如何使用 UffParser 和 Python API 直接导入 TensorFlow 模型。这个示例可以在 /tensorrt/samples/python/end_to_end_tensorflow_mnist 目录中找到。有关更多信息，请参见 end_to_end_tensorflow_mnist Python 示例。

Import TensorRT:

import tensorrt as trt

为TensorFlow模型创建一个冻结的TensorFlow模型。 关于将TensorFlow模型冻结到流中的说明可以在 Freezing A TensorFlow Graph 中找到。

使用 UFF 转换器将冻结的 tensorflow 模型转换为 UFF 文件。通常，这很简单:

convert-to-uff frozen_inference_graph.pb

根据您如何安装 TensorRT，转换到 uff 实用程序可能不会安装在您的系统路径中。在这种情况下，直接调用底层 Python 脚本。它应该位于 UFF 模块的 bin 目录中；例如, ~/.local/lib/python2.7/site-packages/uff/bin/convert_to_uff.py 。

要找到 UFF 模块的位置，运行如下命令 python -c "import uff; print(uff.__path__)"

或者，您可以使用 UFF Paser API 并直接转换 TensorFlow GraphDef。

定义一些路径。更改以下路径，以反映您将示例中包含的模型放置在何处:

model_file = '/data/mnist/minist.uff'

创建 builder, network, 和 parser:

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)

限制:由于 ONNX 格式发展很快，您可能会遇到模型版本和解析器版本之间的版本不匹配。TensorRT 5.0.0 附带的 ONNX 解析器支持 ONNX IR (Intermediate Representation)版本 0.0.3，opset 版本 7。

一般来说，ONNX 解析器的新版本是向后兼容的，因此，遇到由早期版本的 ONNX 导出器生成的模型文件不会造成问题。当更改不向后兼容时，可能会有一些例外。在这种情况下，将早期的 ONNX 模型文件转换为稍后支持的版本。有关此主题的更多信息，请参见 ONNX Model Opset Version Converter。

用户模型也有可能是由一个导出工具生成的，该工具支持比 TensorRT 附带的 ONNX 解析器支持的更晚的 opset。在这种情况下，检查发布到 GitHub onnx-tensorrt 的最新版本是否支持所需的版本。有关更多信息，请参见 yolov3_onnx。

支持的版本由 onnx_trt_backend.cpp 中的 BACKEND_OPSET_VERSION 变量定义。从 GitHub 下载并构建 ONNX TensorRT 解析器的最新版本。构建的说明可以在这里找到: TensorRT backend for ONNX.

下面的步骤说明如何使用 OnnxParser 和 Python API 直接导入 ONNX 模型。有关更多信息，请参见 introductory_parser_samples Python示例。

import TensorRT:

import tensorrt as trt

创建 build, network, and parser:

with builder = trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.OnnxParser(network, TRT_LOGGER) as parser:

with open(model_path, 'rb') as model:

parser.parse(model.read())

注意:在 ubuntu14.04 和 CentOS 上，同时加载 torch 模块和 TensorRT 可能会导致 segmentation faults。

构建器的一个功能是搜索其 CUDA 内核目录以获得最快的可用实现，因此必须使用相同的 GPU 来构建优化引擎将运行的 GPU。

构建器具有许多属性，您可以设置这些属性以控制网络应运行的精度，以及自动调整参数，例如TensorRT在确定哪个最快时(多次迭代会导致更长的运行时间,但是对噪声的敏感性较小)应该为每个内核计时多少次。您还可以查询构建器以找出硬件本机支持的混合精度类型。

两个特别重要的属性是最大批处理大小(maximum batch size )和最大工作区大小(maximum workspace size)。

最大批大小指定 TensorRT 要优化的批大小。在运行时，可以选择较小的批处理大小。

层算法通常需要临时工作区。此参数限制网络中任何层可以使用的最大大小。如果提供的划痕不够，TensorRT 可能无法找到给定层的实现。

使用构建对象构建 engine：

builder.max_batch_size = max_batch_size

builder.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 trt.Builder(TRT_LOGGER) as builder:

with builder.build_cuda_engine(network) as engine:

# Do inference here.

当您序列化时，您将引擎转换为一种格式，以便稍后存储并用于推理。要用于推理，只需对引擎进行反序列化即可。序列化和反序列化是可选的。由于从网络定义创建引擎可能会耗费大量时间，因此您可以避免在每次应用程序重新运行时重新构建引擎，方法是对引擎进行一次序列化，并在进行推断时对其进行反序列化。因此，在构建引擎之后，用户通常希望对其进行序列化，以便以后使用。

从这里开始，您可以序列化引擎，也可以直接使用引擎进行推理。在将模型用于推理之前，序列化和反序列化是一个可选步骤——如果需要，引擎对象可以直接用于推理。

注意:序列化引擎不能跨平台或 TensorRT 版本移植。引擎特定于它们所构建的 GPU 模型(除了平台和 TensorRT 版本之外)。

将模型序列化为 modelstream:

serialized_engine = engine.serialize()

反序列化 modelstream 以执行推理。反序列化需要创建一个运行时对象:

with trt.Runtime(TRT_LOGGER) as runtime:

engine = runtime.deserialize_cuda_engine(serialized_engine)

最后一个参数是使用定制层应用程序的插件层工厂，它是可选的。可以在 Extending TensorRT With Custom Layers 中找到更多细节。

也可以将序列化引擎保存到文件中，并从文件中读回：

序列化引擎并写入文件:

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())

下面的步骤说明了如何在Python中执行推断，现在您已经有了一个引擎。

为输入和输出分配一些主机和设备缓冲区:

# 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(engine.get_binding_shape(0).volume(), dtype=np.float32)

h_output = cuda.pagelocked_empty(engine.get_binding_shape(1).volume(), 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(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 上创建。