运行器

模块： polygraphy.backend.trt

class TrtRunner(engine, name: str | None = None, optimization_profile: int | None = None, allocation_strategy: str | None = None, weight_streaming_budget: int | None = None, weight_streaming_percent: float | None = None)[source]

基类： BaseRunner

使用 TensorRT 运行推理。

请注意，运行器并非为生产环境部署而设计，通常仅应用于原型设计、测试和调试。

参数:

  • engine (Union[Union[trt.ICudaEngine, trt.IExecutionContext], Callable() -> Union[trt.ICudaEngine, trt.IExecutionContext]]) – TensorRT 引擎或执行上下文，或返回其中一个的可调用对象。如果提供了引擎，运行器将自动创建上下文。

  • name (str) – 用于此运行器的人类可读名称前缀。运行器计数和时间戳将附加到此前缀。

  • optimization_profile (int) – 每次激活此运行器时要设置的优化配置文件的索引。如果未提供，则不会显式设置配置文件，并将默认为第 0 个配置文件。您也可以在使用 set_profile() 方法在运行器激活后更改配置文件。

  • allocation_strategy (str) –

    为执行上下文分配设备内存（内部激活和暂存内存）的方式。此参数的值可以是

    • “static”：默认值。执行上下文将预先分配一块内存，该内存足以满足所有配置文件中任何可能的输入大小。

    • “profile”：根据配置文件最大形状为当前配置文件分配足够的设备内存。

    • “runtime”：为当前输入形状分配足够的设备内存。

  • weight_streaming_budget (int) –

    TensorRT 在运行时可以用于权重的 GPU 内存量。它可以采用以下值

    None 或 -2：在运行时禁用权重流式传输。-1：TensorRT 将自动决定流式传输预算。>= 0：TensorRT 允许用于权重的最大 GPU 内存量（以字节为单位）。

  • weight_streaming_percent (float) –

    TRT 将保留在 GPU 上的权重的百分比。它可以采用以下值

    None 或 100%：在运行时禁用权重流式传输。[0 到 100]：TRT 将流式传输的权重的百分比。0 将流式传输最大数量的权重。

set_profile(index: int)[source]

为此运行器设置活动的优化配置文件。运行器必须已激活（请参阅 __enter__() 或 activate()）。

这仅在您的引擎构建时使用了多个优化配置文件时适用。

在 TensorRT 8.0 及更高版本中，配置文件将使用此运行器的 CUDA 流（runner.stream）异步设置。

默认情况下，运行器使用第一个配置文件（配置文件 0）。

参数:

index (int) – 要使用的优化配置文件的索引。

infer_impl(feed_dict, copy_outputs_to_host=None, return_raw_buffers=None)[source]

使用 TensorRT 运行推理的实现。请勿直接调用此方法 - 请改用 infer()，它会将无法识别的参数转发给此方法。

参数:

  • feed_dict (OrderedDict[str, Union[numpy.ndarray, DeviceView, torch.Tensor]]) – 输入张量名称到相应输入 NumPy 数组、Polygraphy DeviceView 或 PyTorch 张量的映射。如果在 feed_dict 中提供了 PyTorch 张量，则此函数还将返回 PyTorch 张量形式的输出。如果提供的输入已驻留在 GPU 内存中，则不会进行额外的复制。

  • copy_outputs_to_host (bool) – 是否将推理输出复制回主机内存。如果为 False，则返回 PyTorch GPU 张量或 Polygraphy DeviceView，而不是 PyTorch CPU 张量或 NumPy 数组。默认为 True。

返回值:

输出张量名称到相应输出 NumPy 数组、Polygraphy DeviceView 或 PyTorch 张量的映射。

返回类型:

OrderedDict[str, Union[numpy.ndarray, DeviceView, torch.Tensor]]

__enter__()

激活运行器以进行推理。例如，这可能涉及分配 CPU 或 GPU 内存。

__exit__(exc_type, exc_value, traceback)

停用运行器。例如，这可能涉及释放 CPU 或 GPU 内存。

activate()

激活运行器以进行推理。例如，这可能涉及分配 CPU 或 GPU 内存。

通常，您应该使用上下文管理器，而不是手动激活和停用。例如

with RunnerType(...) as runner:
    runner.infer(...)
deactivate()

停用运行器。例如，这可能涉及释放 CPU 或 GPU 内存。

通常，您应该使用上下文管理器，而不是手动激活和停用。例如

with RunnerType(...) as runner:
    runner.infer(...)
get_input_metadata(use_numpy_dtypes=None)

返回有关模型输入的信息。此处的形状可能包括动态维度，以 None 表示。必须仅在 activate() 之后和 deactivate() 之前调用。

参数:

use_numpy_dtypes (bool) – [已弃用] 是否返回 NumPy 数据类型而不是 Polygraphy DataType。提供此功能是为了保持向后兼容性。将来，此参数将被删除，并且始终返回 Polygraphy DataType。可以通过调用 numpy() 方法将它们转换为 NumPy 数据类型。默认为 True。

返回值:

输入名称、形状和数据类型。

返回类型:

TensorMetadata

infer(feed_dict, check_inputs=True, *args, **kwargs)

使用提供的 feed_dict 运行推理。

必须仅在 activate() 之后和 deactivate() 之前调用。

注意：某些运行器可能在 infer() 中接受其他参数。有关这些参数的详细信息，请参阅其 infer_impl() 方法的文档。

参数:

  • feed_dict (OrderedDict[str, numpy.ndarray]) – 输入张量名称到相应输入 NumPy 数组的映射。

  • check_inputs (bool) – 是否检查提供的 feed_dict 是否包含具有预期数据类型和形状的预期输入。禁用此项可能会提高性能。默认为 True。

inference_time

运行推理所需的时间，以秒为单位。

类型:

float

返回值:

输出张量名称到其相应 NumPy 数组的映射。

重要提示：运行器可能会重复使用这些输出缓冲区。因此，如果您需要保存来自多次推理的输出，则应使用 copy.deepcopy(outputs) 制作副本。

返回类型:

OrderedDict[str, numpy.ndarray]

last_inference_time()

返回上次调用 infer() 期间所需的总推理时间（以秒为单位）。

必须仅在 activate() 之后和 deactivate() 之前调用。

返回值:

时间（以秒为单位），如果运行器未测量运行时，则为 None。

返回类型:

float

is_active

此运行器是否已激活，可以通过上下文管理器激活，也可以通过调用 activate() 激活。

类型:

bool