3 WORKING WITH TENSORRT USING THE PYTHON API
以下部分重点介绍了可以使用Python API执行的TensorRT用户目标和任务。这些部分专注于使用Python API而不使用任何框架。进一步的细节在Samples部分提供,并在适当的地方链接到下面。
假设您从训练好的模型开始。本章将介绍使用TensorRT的以下必要步骤:
- 从模型中创建TensorRT网络定义;
- 调用TensorRT builder以从网络创建优化的运行时引擎;
- 序列化和反序列化引擎,以便在运行时快速重新创建;
- 向引擎填充数据,执行推理。
Python API vs C++ API
从本质上讲,C ++ API和Python API在支持您的需求方面应该完全相同。 Python API的主要优点是数据预处理和后处理易于使用,因为您可以使用各种库,如NumPy和SciPy。
C ++ API应该用于安全性很重要的情况,例如汽车行业。有关C ++ API的更多信息,请参阅第二章使用TensorRT的C ++ API。
有关如何使用Python优化性能的更多信息,请参阅优化Python性能最佳实践指南。
3.1 在python中导入TensorRT
1.导入TensorRT
import tensorrt as trt
2.实现日志记录接口,TensorRT通过该接口报告错误、警告和信息性消息。以下代码显示了如何实现日志记录接口。在这种情况下,我们已经抑制了信息性消息,并仅报告警告和错误。TensorRT Python绑定中包含一个简单的记录器。
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
3.2 在Python中创建网络定义
使用TensorRT进行推理的第一步是从您的模型创建TensorRT网络。实现此目的的最简单方法是使用TensorRT解析器库导入模型(参考3.2.2,3.2.3,3.2.4,3.2.5节),支持以下格式的序列化模型:
- Caffe (both BVLC and NVCaffe);
- ONNX 1.0 and 1.1;
- UFF (used for TensorFlow)。
另一种方法是使用TensorRT Network API直接定义模型,请参考3.2.1节。这要求您进行少量API调用以定义网络图中的每个层,并为模型的训练参数实现自己的导入机制。
TensorRT Python API仅适用于x86_64平台。有关详细信息,请参阅Deep Learning SDK文档 - TensorRT工作流程。
3.2.1 使用python API从头定义网络
创建网络时,必须首先定义引擎并为推理创建builder对象。Python API用于从Network API创建网络和引擎。网络定义参考用于向网络添加各种层。有关使用Python API创建网络和引擎的更多信息,请参阅network_api_pytorch_mnist示例。
以下代码说明了如何使用Input,Convolution,Pooling,FullyConnected,Activation和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 使用Python中的Parser导入模型
要使用解析器导入模型,您需要执行以下高级步骤:
- 创建TensorRT builder和network;
- 为特定格式创建TensorRT解析器;
- 使用解析器解析导入的模型并填充网络。
必须在创建network之前创建builder,因为它充当网络的工厂。不同的解析器具有用于标记网络输出的不同机制。具体可参考 UFF Parser API ,Caffe Parser API, and ONNX Parser API.
3.2.3 使用Python导入caffe模型
以下步骤说明了如何使用CaffeParser和Python API直接导入Caffe模型。有关更多信息,请参阅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_tensors,它是一个包含从张量名称到ITensor对象的映射的表。
3.2.4 使用Python接口导入Tensorflow模型
以下步骤说明了如何使用UffParser和Python API直接导入TensorFlow模型。此示例可以在 / tensorrt / samples / python / end_to_end_tensorflow_mnist目录中找到。有关更多信息,请参阅end_to_end_tensorflow_mnist Python示例。
1.导入TensorRT
import tensorrt as trt
2.为Tensorflow模型创建冻结的TensorFlow模型。关于将TensorFlow模型冻结到流中的说明可以在Freezing A TensorFlow Graph中找到。
3.使用UFF转换器将冻结的tensorflow模型转换为UFF文件。
通常,这很简单:
convert-to-uff frozen_inference_graph.pb
根据您安装TensorRT的方式,可能无法在系统路径中安装convert-to-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 Parser API并直接转换TensorFlow GraphDef。
4.定义路径。更改以下路径以反映Samples中包含的模型的位置:
model_file = '/data/mnist/mnist.uff'
5.创建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)
3.2.5 使用Python接口导入ONNX模型
限制:由于ONNX格式正在快速开发,因此您可能会遇到模型版本和解析器版本之间的版本不匹配。 TensorRT 5.0.0附带的ONNX Parser支持ONNX IR(中间表示)版本0.0.3,opset版本7。
通常,较新版本的ONNX Parser旨在向后兼容,因此,遇到早期版本的ONNX导出器生成的模型文件不应该导致问题。当更改不向后兼容时,可能会有一些例外情况。在这种情况下,将早期的ONNX模型文件转换为以后支持的版本。有关此主题的更多信息,请参阅ONNX Model Opset版本转换器。
用户模型也可能是由支持后续opset的导出工具生成的,而不是由TensorRT附带的ONNX解析器支持的。在这种情况下,请检查发布到GitHub onnx-tensorrt的最新版本的TensorRT是否支持所需的版本。有关更多信息,请参阅yolov3_onnx。
支持的版本由onnx_trt_backend.cpp中的BACKEND_OPSET_VERSION变量定义。从GitHub下载并构建最新版本的ONNX TensorRT Parser。有关构建的说明,请访问:ONNX的TensorRT后端。
以下步骤说明了如何使用OnnxParser和Python API直接导入ONNX模型。有关更多信息,请参阅introduction_parser_samples Python示例。
1.导入TensorRT
import tensorrt as trt
2.创建builder、network和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())
3.2.6 导入PyTorch或其他框架的模型
将TensorRT与PyTorch(或任何其他具有NumPy兼容权重的框架)一起使用涉及使用TensorRT API复制网络架构(请参阅使用Python API从头创建网络定义),然后从PyTorch模型中复制权重。有关更多信息,请参阅使用PyTorch和其他框架(示例代码为/usr/src/tensorrt/
samples/python/network_api_pytorch_mnist)。
3.3 使用Python接口创建引擎
builder的一个功能是搜索其CUDA内核目录以获得最快的可用实现,因此必须使用相同的GPU来构建优化引擎将运行的GPU。
builder具有许多属性,您可以设置这些属性以控制网络应运行的精度,以及自动调整参数,例如TensorRT在确定哪个最快时(多次迭代导致更长的运行时间)应该为每个内核计时多少次同时对噪声的敏感性较低)。您还可以查询builder以找出硬件本身支持的混合精度类型。
两个特别重要的属性是最大batch大小和最大workspace大小。
- 最大batch大小指定TensorRT将优化的批量大小。在运行时,可以选择较小的批量大小;
- 层算法通常需要临时工作空间。此参数限制网络中任何层可以使用的最大大小。如果提供的scratch不足,则TensorRT可能无法找到给定层的实现。
有关使用Python构建引擎的更多信息,请参阅introduction_orySamples示例。
1.使用builder对象创建引擎:
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会复制权重。
2.进行推理,请参考3.5节。
3.4 在Python中序列化引擎
序列化时,您正在将引擎转换为一种格式,以便以后存储和使用以进行推理。要用于推理,您只需反序列化引擎即可。序列化和反序列化是可选的。由于从网络定义创建引擎可能非常耗时,因此每次应用程序重新运行时都可以通过序列化一次并在推理时对其进行反序列化来避免重建引擎。因此,在构建引擎之后,用户通常希望将其序列化以供以后使用。
从这里开始,您可以序列化引擎,也可以直接使用引擎进行推理。在将模型用于推理之前,序列化和反序列化是一个可选步骤 - 如果需要,可以直接使用引擎对象进行推理。
序列化引擎不能跨平台或TensorRT版本移植。引擎特定于它们构建的精确GPU模型(包括平台和TensorRT版本)。
1.序列化模型为modelstream:
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())
3.5 使用Python接口执行推理
以下步骤说明了如何在Python中执行推理,现在您已拥有引擎。
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(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()
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(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上创建。