Python API 文档#

注意

TensorRT Python API 使 Python 基础开发环境中的开发人员以及希望试验 TensorRT 的人员能够轻松解析模型（例如，来自 ONNX）并生成和运行 PLAN 文件。

本节说明了 Python API 的基本用法，假设您从 ONNX 模型开始。onnx_resnet50.py 示例更详细地说明了此用例。

可以通过 tensorrt 模块访问 Python API

import tensorrt as trt

构建阶段#

要创建构建器，您必须首先创建一个记录器。Python 绑定包含一个简单的记录器实现，该实现将所有先于特定严重性的消息记录到 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 = trt.Builder(logger)

构建引擎旨在作为离线过程，因此可能需要相当长的时间。优化构建器性能 部分提供了有关使构建器运行更快的技巧。

创建网络定义#

创建构建器后，优化模型的第一步是创建网络定义。网络定义选项是使用标志的组合（OR 运算）指定的。

您可以使用 NetworkDefinitionCreationFlag.STRONGLY_TYPED 标志指定应将网络视为强类型。有关更多信息，请参阅 强类型网络 部分。

最后，创建一个网络

network = builder.create_network(flag)

从头开始创建网络定义（高级）#

除了使用解析器之外，您还可以通过网络定义 API 直接向 TensorRT 定义网络。此方案假定每层权重已在主机内存中准备就绪，以便在网络创建期间传递给 TensorRT。

与本节对应的代码可以在 network_api_pytorch_mnist 中找到。

此示例创建了一个简单的网络，其中包含输入、卷积、池化、矩阵乘法、Shuffle、激活和 Softmax 层。此示例使用一个辅助类来保存有关模型的一些元数据

class ModelData(object):
    INPUT_NAME = "data"
    INPUT_SHAPE = (1, 1, 28, 28)
    OUTPUT_NAME = "prob"
    OUTPUT_SIZE = 10
    DTYPE = trt.float32

在此示例中，权重从 PyTorch MNIST 模型导入。

weights = mnist_model.get_weights()

创建记录器、构建器和网络类。

TRT_LOGGER = trt.Logger(trt.Logger.ERROR)
builder = trt.Builder(TRT_LOGGER)
network = builder.create_network(0)

接下来，为网络创建输入张量，指定张量的名称、数据类型和形状。

input_tensor = network.add_input(name=ModelData.INPUT_NAME, dtype=ModelData.DTYPE, shape=ModelData.INPUT_SHAPE)

添加卷积层，指定输入、输出映射的数量、内核形状、权重、偏差和步幅

conv1_w = weights["conv1.weight"].cpu().numpy()
    conv1_b = weights["conv1.bias"].cpu().numpy()
    conv1 = network.add_convolution_nd(
    input=input_tensor, num_output_maps=20, kernel_shape=(5, 5), kernel=conv1_w, bias=conv1_b
    )
    conv1.stride_nd = (1, 1)

添加池化层，指定输入（先前卷积层的输出）、池化类型、窗口大小和步幅

pool1 = network.add_pooling_nd(input=conv1.get_output(0), type=trt.PoolingType.MAX, window_size=(2, 2))
    pool1.stride_nd = trt.Dims2(2, 2)

添加下一对卷积层和池化层

conv2_w = weights["conv2.weight"].cpu().numpy()
conv2_b = weights["conv2.bias"].cpu().numpy()
conv2 = network.add_convolution_nd(pool1.get_output(0), 50, (5, 5), conv2_w, conv2_b)
conv2.stride_nd = (1, 1)

pool2 = network.add_pooling_nd(conv2.get_output(0), trt.PoolingType.MAX, (2, 2))
pool2.stride_nd = trt.Dims2(2, 2)

添加 Shuffle 层以重塑输入，为矩阵乘法做准备

batch = input.shape[0]
mm_inputs = np.prod(input.shape[1:])
input_reshape = net.add_shuffle(input)
input_reshape.reshape_dims = trt.Dims2(batch, mm_inputs)

现在，添加一个矩阵乘法层。模型导出器提供了转置权重，因此指定了 kTRANSPOSE 选项。

auto prob = network->addSoftMax(*relu1->getOutput(0));

添加偏差，它将在批次维度上广播

bias_const = net.add_constant(trt.Dims2(1, nbOutputs), weights["fc1.bias"].numpy())
bias_add = net.add_elementwise(mm.get_output(0), bias_const.get_output(0), trt.ElementWiseOperation.SUM)

添加 ReLU 激活层

relu1 = network.add_activation(input=fc1.get_output(0), type=trt.ActivationType.RELU)

添加最终全连接层，并将此层的输出标记为整个网络的输出

fc2_w = weights['fc2.weight'].numpy()
fc2_b = weights['fc2.bias'].numpy()
fc2 = add_matmul_as_fc(network, relu1.get_output(0), ModelData.OUTPUT_SIZE, fc2_w, fc2_b)

fc2.get_output(0).name = ModelData.OUTPUT_NAME
network.mark_output(tensor=fc2.get_output(0))

现在已完全构建表示 MNIST 模型的网络。有关如何构建引擎并使用此网络运行推理的说明，请参阅构建引擎 和 执行推理 部分。

有关图层的更多信息，请参阅 TensorRT 运算符文档。

使用 ONNX 解析器导入模型#

现在，必须从 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

构建引擎#

下一步是创建一个构建配置，指定 TensorRT 应如何优化模型

config = builder.create_builder_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)

注意

序列化引擎在平台之间不可移植。它们特定于构建它们的精确 GPU 模型（除了平台之外）。

反序列化 Plan#

当您有一个先前序列化的优化模型并且想要执行推理时，您必须首先创建 Runtime 接口的实例。与构建器一样，运行时需要记录器的实例

runtime = trt.Runtime(logger)

可以使用以下方法反序列化引擎：[1]

内存中反序列化

此方法简单明了，适用于较小的模型或内存不是约束的情况。将 plan 文件读取到内存缓冲区中。

with open("model.plan", "rb") as f:
    model_data = f.read()
engine = runtime.deserialize_cuda_engine(model_data)

您可以从序列化的引擎对象反序列化引擎

serialized_engine = builder.build_serialized_network(network, config)
engine = runtime.deserialize_cuda_engine(serialized_engine)

IStreamReaderV2 反序列化

这是最先进的方法，支持读取到主机和设备指针，并实现潜在的性能改进。使用此方法，无需将整个 plan 文件读取到缓冲区进行反序列化，因为 IStreamReaderV2 允许根据需要以块的形式读取文件。

class StreamReaderV2(trt.IStreamReaderV2):
    def __init__(self, bytes):
        trt.IStreamReaderV2.__init__(self)
        self.bytes = bytes
        self.len = len(bytes)
        self.index = 0

    def read(self, size, cudaStreamPtr):
        assert self.index + size <= self.len
        data = self.bytes[self.index:self.index + size]
        self.index += size
        return data

    def seek(self, offset, where):
        if where == SeekPosition.SET:
            self.index = offset
        elif where == SeekPosition.CUR:
            self.index += offset
        elif where == SeekPosition.END:
            self.index = self.len - offset
        else:
            raise ValueError(f"Invalid seek position: {where}")

reader_v2 = MyStreamReaderV2("model.plan")
engine = runtime.deserialize_cuda_engine(reader_v2)

trt.IStreamReaderV2 方法对于大型模型或使用高级功能（如 GPUDirect 或权重流式传输）尤其有利。它可以显着减少引擎加载时间和内存使用量。

选择反序列化方法时，请考虑您的具体要求

  • 对于小型模型或简单的用例，内存中反序列化通常就足够了。

  • 对于大型模型或内存效率至关重要的情况，请考虑使用 trt.IStreamReaderV2。

  • 如果您需要自定义文件处理或流式传输功能，trt.IStreamReaderV2 提供了必要的灵活性。

执行推理#

引擎保存优化的模型，但推理需要中间激活的附加状态。这通过 IExecutionContext 接口完成

context = engine.create_execution_context()

一个引擎可以有多个执行上下文，允许一组权重用于多个重叠的推理任务。（当前的一个例外是使用动态形状时，除非指定了预览功能 PROFILE_SHARING_0806，否则每个优化配置文件只能有一个执行上下文。）

要执行推理，您必须为输入和输出指定缓冲区

context.set_tensor_address(name, ptr)

多个 Python 包允许您在 GPU 上分配内存，包括但不限于官方 CUDA Python 绑定、PyTorch、cuPy 和 Numba。

填充输入缓冲区后，您可以调用 TensorRT 的 execute_async_v3 方法以使用 CUDA 流启动推理。网络将异步或同步执行，具体取决于网络的结构和功能。可能导致同步行为的非详尽列表包括数据相关的形状、DLA 使用、循环和同步插件。

首先，创建 CUDA 流。如果您已经有一个，请使用指向它的指针，例如，对于 PyTorch CUDA 流 torch.cuda.Stream()，您可以使用 cuda_stream 属性访问指针；对于 Polygraphy CUDA 流，使用 ptr 属性；或者您可以通过调用 cudaStreamCreate() 直接使用 CUDA Python 绑定创建流。

接下来，启动推理

context.execute_async_v3(buffers, stream_ptr)

通常在内核之前和之后排队异步传输 (cudaMemcpyAsync()) 以将数据从 GPU 移动，如果数据尚未在那里。

要确定推理（和异步传输）何时完成，请使用标准 CUDA 同步机制，例如事件或等待流。例如，对于 PyTorch CUDA 流或 Polygraphy CUDA 流，发出 stream。要 synchronize() 与使用 CUDA Python 绑定创建的流同步，请发出 cudaStreamSynchronize(stream)。

脚注

在本页中