使用Python部署TensorRT

本文介绍在没有任何框架的情况下使用 Python API进行TensorRT7.2.3的任务。TensorRT7.2.3样例支持指南中提供了更多详细信息，并在适当的情况下链接到以下内容。

假设你从训练好的模型开始。本文将介绍使用TensorRT的以下必要步骤：

• 从您的模型创建TensorRT网络定义

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

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

• 为引擎提供数据以执行推理

Python API 与 C++ API

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

C++ API 应该用于安全性和性能很重要的情况，例如在汽车中。有关C++ API的更多信息，请参阅使用C++部署TensorRT。

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

1. 将TensorRT导入 Python

程序

1. 导入TensorRT：

   import tensorrt as trt
   

2. 实现一个logging接口，TensorRT通过它报告errprs、warnings和info。以下代码展示了如何实现logging接口。在这种情况下，我们抑制了info，只报告errors和warnings。TensorRT Python绑定中包含一个简单的logger。

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

尽管可以无需创建 CUDA 上下文（将创建默认上下文），但这是不可取的。 建议在创建运行时runtime或builder对象之前创建和配置 CUDA 上下文。

将使用与创建线程关联的 GPU 上下文创建构建器或运行时。 虽然如果默认上下文尚不存在，但会创建它，但建议在创建运行时或构建器对象之前创建和配置 CUDA 上下文

import pycuda.driver as cuda

ctx = cuda.Device(0).make_context()
ctx.push()
...  # 申请空间，传数据，推理执行，回传
ctx.pop()

2. 在Python中创建网络

使用 TensorRT 执行推理的第一步是从您的模型创建一个 TensorRT 网络。实现这一目标有两种方式：

a. 最简单方法是使用TensorRT解析库导入模型（见本文2.2，2.3，2.4，2.5），它支持以下格式的序列化模型：

• Caffe（BVLC 和 NVCaffe）
• 支持最高ONNX1.6发行版，7至11级ONNX算子集
• UFF（用于TensorFlow）

b. 直接使用TensorRT Network API定义模型（请参阅本文2.1)。这需要您进行少量API调用来定义网络图中的每一层，并为模型的训练参数实现您自己的导入机制。

注： 该TensorRT Python的API不是适用于所有平台。有关更多信息，请参阅TensorRT支持矩阵

2.1. 使用Python API从头开始创建网络定义

创建网络时，必须首先定义engine并创建用于推理的builder对象。Python API用于从Network APIs创建网络和引擎。网络定义参考用于向网络添加各种层。

关于这个方式

有关使用 Python API创建网络和引擎的更多信息，请参阅“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))

2.2. 在 Python 中使用解析器导入模型

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

1. 创建TensorRT的Builder和Network

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

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

对于一步步的说明，请参阅本文2.3，2.4，2.5。

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

2.3. 使用Python从Caffe导入

以下步骤说明了如何使用CaffeParser和Python API 直接导入Caffe模型。

关于这个方式

有关更多信息，请参阅使用Python将Caffe、TensorFlow和ONNX模型导入TensorRT简介(introductory_parser_samples)示例。

程序

1. 导入TensorRT。

   import tensorrt as trt
   

2. 定义数据类型。在这个例子中，我们将使用float32。

   datatype = trt.float32
   

3. 另外，定义一些路径。更改以下路径以反映你放置示例中包含的模型的位置：

   deploy_file = 'data/mnist/mnist.prototxt'
   model_file = 'data/mnist/mnist.caffemodel'
   

4. 创建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_tensor，这是一个包含从张量名称到ITtensor对象的表。

2.4. 使用 Python从TensorFlow导入

以下步骤说明了如何使用UffParser和Python API直接导入TensorFlow模型。

关于这个方式

这个样例在/tensorrt/samples/python/end_to_end_tensorflow_mnist目录。有关详细信息，请参阅"Hello World" For TensorRT Using TensorFlow And Python(end_to_end_tensorflow_mnist)示例。

程序

1. 导入TensorRT：

   import tensorrt as trt
   

2. 创建一个TensorFlow模型。有关将TensorFlow模型冻结为流的说明，请参见Freezing A TensorFlow Graph。

3. 使用UFF转换器转换frozen tensorflow 模型到UFF文件。通常，这很简单：
   convert-to-uff frozen_inference_graph.pb

根据TensorRT的安装方式，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。

4. 定义一些路径。更改以下路径以反映你放置示例中包含的模型的位置：

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

5. 创建builder, network和parser：

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

2.5. 使用 Python从ONNX导入

以下步骤说明了如何使用OnnxParser和Python API直接导入ONNX模型。

关于这个方式

有关更多信息，请参阅使用使用Python将Caffe、TensorFlow和ONNX模型导入TensorRT(introductory_parser_samples)示例。

注意：
一般来说，较新版本的OnnxParser被设计为向后兼容，因此，遇到由较早版本的ONNX导出器生成的模型文件应该不会导致问题。当更改不向后兼容时，可能会有一些例外。在这种情况下，请将较早的ONNX模型文件转换为较晚支持的版本。有关此更多信息，请参阅ONNX模型算子集版本转换器。

用户模型也可能是由导出工具生成的，该工具支持的算子集比TensorRT附带的ONNX解析器支持的算子集要晚。在这种情况下，检查是否有最新版本的TensorRT发布到GitHub，onnx-tensorrt, 支持所需版本。有关更多信息，请参阅Python中使用ONNX TensorRT后端实现目标检测(yolov3_onnx)示例。

支持的版本由onnx_trt_backend.cpp中的BACKEND_OPSET_VERSION变量决定. 从 GitHub下载并构建最新版本的ONNXTensorRT Parser。可以在此处找到构建说明：ONNX的 TensorRT后端。

在 TensorRT 7中，ONNX 解析器仅支持全维模式，这意味着你的网络定义必须使用显式批处理explicitBatch标志设置。有关更多信息，请参阅使用动态shapes。

程序

1. 导入TensorRT：

   import tensorrt as trt
   

2. 创建builder,network和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))
   

   注： 你转化的ONNX模型可以使用onnx-simplifier简化：

   $ pip install onnx-simplifier
   $ python -m onnxsim {your_onnx_model}.onnx {your_sim_onnx_file}.onnx
   

2.6. 从PyTorch和其他框架导入

关于这个方式

使用带有Pytorch（或任何其他具有 NumPy兼容权重的框架）的TensorRT包含使用API复制网络架构（请参阅本文2.1），然后从PyTorch复制权重。有关更多信息，请参阅使用PyTorch和其他框架。

对于PyTorch来说，一般采用PyTorch ——> ONNX ——> TensorRT的转换路径构建engine。

要执行推理，请按照本文第5节中概述的说明进行操作。

3. 用 Python构建引擎

builder的功能之一是在其 CUDA 内核目录中搜索可用的最快实现，因此有必要使用与优化引擎将运行的相同GPU来构建。

关于这个步骤

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

一个特别重要的特性是最大工作空间大小。

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

有关在Python中构建引擎的更多信息，请参阅使用Python将Caffe、TensorFlow和ONNX模型导入TensorRT(introductory_parser_samples)示例。

程序

1. 使用builder对象构建引擎：

   with trt.Builder(TRT_LOGGER) as builder, builder.create_builder_config() as config:
       # max_workspace_size determines the amount of memory available to the builder when building an optimized engine
       config.max_workspace_size = 1 << 20
       with builder.build_engine(network, config) as engine:
           # Do inference here.
   

   构建引擎时，TensorRT会复制权重。

2. 进行推理。要执行推理，请按照本文第5节中概述的说明进行操作。

4. 在Python中序列化模型

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

关于这个步骤

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

注意： 序列化引擎不可跨平台或TensorRT版本移植。引擎特定于它们构建的具体GPU模型（包括平台和TensorRT版本）。

1. 将模型序列化为模型流：

serialized_engine = engine.serialize()

2. 反序列化模型流以执行推理。反序列化需要创建一个运行时对象：

with trt.Runtime(TRT_LOGGER) as runtime:
    engine = runtime.deserialize_cuda_engine(serialized_engine)

也可以将序列化引擎保存到文件中，然后从文件中读回:

1. 序列化引擎并写入文件：

with open(“sample.engine”, “wb”) as f:
    f.write(engine.serialize())

2. 从文件中读取引擎并反序列化：

with open(“sample.engine”, “rb”) as f, trt.Runtime(TRT_LOGGER) as runtime:
    engine = runtime.deserialize_cuda_engine(f.read())

5. 在 Python 中执行推理

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

程序

1. 为输入和输出分配一些主机和设备缓冲区。这个例子假设 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()

2. 创建一些空间来存储中间激活值。由于引擎保存网络定义和训练参数，因此需要额外的空间。这些保存在执行上下文中：

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上创建。