目录
快捷方式

torch_tensorrt

函数

torch_tensorrt.compile(module: Any, ir: str = 'default', inputs: Optional[Sequence[Input | torch.Tensor | InputTensorSpec]] = None, arg_inputs: Optional[Sequence[Sequence[Any]]] = None, kwarg_inputs: Optional[dict[Any, Any]] = None, enabled_precisions: Optional[Set[Union[dtype, dtype]]] = None, **kwargs: Any) → Union[Module, ScriptModule, GraphModule, Callable[[...], Any]][源代码]

使用 TensorRT 为 NVIDIA GPU 编译 PyTorch 模块

接收一个现有的 PyTorch 模块和一组用于配置编译器的设置，并使用在 ir 中指定的路径来降级和编译模块为 TensorRT，并返回一个 PyTorch 模块。

特别转换模块的 forward 方法

参数

module (Union(torch.nn.Module,torch.jit.ScriptModule) – 源模块

关键字参数

  • inputs (List[Union(Input, torch.Tensor)]) –

    必需 模块输入的输入形状、数据类型和内存布局的规范列表。此参数是必需的。输入大小可以指定为 torch 大小、元组或列表。数据类型可以使用 torch 数据类型或 torch_tensorrt 数据类型指定，您可以使用 torch 设备或 torch_tensorrt 设备类型枚举来选择设备类型。

    inputs=[
    torch_tensorrt.Input((1, 3, 224, 224)), # Static NCHW input shape for input #1
    torch_tensorrt.Input(
        min_shape=(1, 224, 224, 3),
        opt_shape=(1, 512, 512, 3),
        max_shape=(1, 1024, 1024, 3),
        dtype=torch.int32
        format=torch.channel_last
    ), # Dynamic input shape for input #2
    torch.randn((1, 3, 224, 244)) # Use an example tensor and let torch_tensorrt infer settings
]

  • arg_inputs (Tuple[Any, ...]) – 与 inputs 相同。为便于与 kwarg_inputs 一起理解而使用的别名。

  • kwarg_inputs (dict[Any, ...]) – 可选，传递给模块 forward 函数的关键字参数输入。

  • enabled_precision (Set(Union(torch.dpython:type, torch_tensorrt.dpython:type))) – TensorRT 在选择内核时可以使用的数据类型集合

  • ir (str) – 请求的编译策略。（选项：default - 让 Torch-TensorRT 决定，ts - 带脚本路径的 TorchScript）

  • **kwargs – 特定请求策略的附加设置（更多信息请参见子模块）

返回

已编译的模块，运行时将通过 TensorRT 执行

返回类型

torch.nn.Module

torch_tensorrt.convert_method_to_trt_engine(module: Any, method_name: str = 'forward', inputs: Optional[Sequence[Input | torch.Tensor | InputTensorSpec]] = None, arg_inputs: Optional[Sequence[Sequence[Any]]] = None, kwarg_inputs: Optional[dict[Any, Any]] = None, ir: str = 'default', enabled_precisions: Optional[Set[Union[dtype, dtype]]] = None, **kwargs: Any) → bytes[源代码]

将 TorchScript 模块方法转换为序列化的 TensorRT 引擎

给定一个转换设置字典，将模块的指定方法转换为序列化的 TensorRT 引擎

参数

module (Union(torch.nn.Module,torch.jit.ScriptModule) – 源模块

关键字参数

  • inputs (List[Union(Input, torch.Tensor)]) –

    必需 模块输入的输入形状、数据类型和内存布局的规范列表。此参数是必需的。输入大小可以指定为 torch 大小、元组或列表。数据类型可以使用 torch 数据类型或 torch_tensorrt 数据类型指定，您可以使用 torch 设备或 torch_tensorrt 设备类型枚举来选择设备类型。

    input=[
    torch_tensorrt.Input((1, 3, 224, 224)), # Static NCHW input shape for input #1
    torch_tensorrt.Input(
        min_shape=(1, 224, 224, 3),
        opt_shape=(1, 512, 512, 3),
        max_shape=(1, 1024, 1024, 3),
        dtype=torch.int32
        format=torch.channel_last
    ), # Dynamic input shape for input #2
    torch.randn((1, 3, 224, 244)) # Use an example tensor and let torch_tensorrt infer settings
]

  • arg_inputs (Tuple[Any, ...]) – 与 inputs 相同。为便于与 kwarg_inputs 一起理解而使用的别名。

  • kwarg_inputs (dict[Any, ...]) – 可选，传递给模块 forward 函数的关键字参数输入。

  • enabled_precision (Set(Union(torch.dpython:type, torch_tensorrt.dpython:type))) – TensorRT 在选择内核时可以使用的数据类型集合

  • ir (str) – 请求的编译策略。（选项：default - 让 Torch-TensorRT 决定，ts - 带脚本路径的 TorchScript）

  • **kwargs – 特定请求策略的附加设置（更多信息请参见子模块）

返回

序列化的 TensorRT 引擎，可以保存到文件或通过 TensorRT API 进行反序列化

返回类型

字节

torch_tensorrt.cross_compile_for_windows(*args: List[Any], **kwargs: Dict[str, Any]) → Any
torch_tensorrt.load_cross_compiled_exported_program(file_path: str = '') → Any[源代码]

在 Windows 中加载一个之前在 Linux 中交叉编译的 ExportedProgram 文件

参数

file_path (str) – 磁盘上的文件路径

引发

ValueError – 如果 API 未在 Windows 中调用，或者没有文件，或者文件不是有效的 ExportedProgram 文件

torch_tensorrt.save(module: Any, file_path: str = '', *, output_format: str = 'exported_program', inputs: Optional[Sequence[Tensor]] = None, arg_inputs: Optional[Sequence[Tensor]] = None, kwarg_inputs: Optional[dict[str, Any]] = None, retrace: bool = False, pickle_protocol: int = 2, **kwargs: Any) → None[源代码]

以指定的输出格式将模型保存到磁盘。

参数

  • module (Optional(torch.jit.ScriptModule | torch.export.ExportedProgram | torch.fx.GraphModule | CudaGraphsTorchTensorRTModule)) – 已编译的 Torch-TensorRT 模块

  • inputs (torch.Tensor) – Torch 输入张量

  • arg_inputs (Tuple[Any, ...]) – 与 inputs 相同。为便于与 kwarg_inputs 一起理解而使用的别名。

  • kwarg_inputs (dict[Any, ...]) – 可选，传递给模块 forward 函数的关键字参数输入。

  • output_format (str) – 保存模型的格式。选项包括 exported_program | torchscript | aot_inductor。

  • retrace (bool) – 当模块类型为 fx.GraphModule 时，此选项使用 torch.export.export(strict=False) 重新导出图以保存它。此标志目前为实验性功能。

  • pickle_protocol (python:int) – 用于保存模型的 pickle 协议。默认为 2。对于大型模型，请将其增加到 4 或更高。

torch_tensorrt.load(file_path: str = '') → Any[源代码]

加载 Torchscript 模型或 ExportedProgram。

从磁盘加载 TorchScript 或 ExportedProgram 文件。文件类型将使用 try, except 来检测。

参数

file_path (str) – 磁盘上的文件路径

引发

ValueError – 如果没有文件，或者文件既不是 TorchScript 文件也不是 ExportedProgram 文件

类

class torch_tensorrt.MutableTorchTensorRTModule(pytorch_model: Module, *, device: Optional[Union[Device, device, str]] = None, use_python_runtime: bool = False, immutable_weights: bool = False, strict: bool = True, allow_complex_guards_as_runtime_asserts: bool = False, weight_streaming_budget: Optional[int] = None, enabled_precisions: Optional[Set[Union[dtype, dtype]]] = None, **kwargs: Any)[源代码]

初始化一个 MutableTorchTensorRTModule，以便像操作常规 PyTorch 模块一样无缝地操作它。所有 TensorRT 编译和重装过程都会在您使用模块时自动处理。对其属性的任何更改或加载不同的 state_dict 都会触发重装或重新编译，这将在下一次前向传播期间进行管理。

MutableTorchTensorRTModule 接收一个 PyTorch 模块和一组用于编译器的配置设置。编译完成后，该模块会维护 TensorRT 图模块和原始 PyTorch 模块之间的连接。对 MutableTorchTensorRTModule 所做的任何修改都将反映在 TensorRT 图模块和原始 PyTorch 模块中。

__init__(pytorch_model: Module, *, device: Optional[Union[Device, device, str]] = None, use_python_runtime: bool = False, immutable_weights: bool = False, strict: bool = True, allow_complex_guards_as_runtime_asserts: bool = False, weight_streaming_budget: Optional[int] = None, enabled_precisions: Optional[Set[Union[dtype, dtype]]] = None, **kwargs: Any) → None[源代码]
参数

pytorch_model (torch.nn.module) – 需要加速的源模块

关键字参数

  • device (Union(Device, torch.device, dict)) –

    TensorRT 引擎运行的目标设备

    device=torch_tensorrt.Device("dla:1", allow_gpu_fallback=True)

  • disable_tf32 (bool) – 强制 FP32 层使用传统的 FP32 格式，而不是默认行为，即在相乘前将输入四舍五入到 10 位尾数，但使用 23 位尾数累加求和

  • assume_dynamic_shape_support (bool) – 将此设置为 true 可使转换器适用于动态和静态形状。默认值：False

  • sparse_weights (bool) – 为卷积层和全连接层启用稀疏性。

  • enabled_precision (Set(Union(torch.dpython:type, torch_tensorrt.dpython:type))) – TensorRT 在选择内核时可以使用的数据类型集合

  • immutable_weights (bool) – 构建不可重装的引擎。这对于某些不可重装的层很有用。

  • capability (EngineCapability) – 将内核选择限制为安全的 GPU 内核或安全的 DLA 内核

  • num_avg_timing_iters (python:int) – 用于选择内核的平均计时迭代次数

  • workspace_size (python:int) – 分配给 TensorRT 的最大工作空间大小

  • dla_sram_size (python:int) – DLA 用于在层内通信的快速软件管理 RAM。

  • dla_local_dram_size (python:int) – DLA 用于在操作间共享中间张量数据的主机 RAM

  • dla_global_dram_size (python:int) – DLA 用于存储权重和元数据以供执行的主机 RAM

  • truncate_double (bool) – 将以 double (float64) 提供的权重截断为 float32

  • calibrator (Union(torch_tensorrt._C.IInt8Calibrator, tensorrt.IInt8Calibrator)) – 校准器对象，它将为 INT8 校准的 PTQ 系统提供数据

  • require_full_compilation (bool) – 要求模块端到端编译，否则返回错误，而不是返回一个混合图，其中无法在 TensorRT 中运行的操作在 PyTorch 中运行

  • min_block_size (python:int) – 为了在 TensorRT 中运行一组操作，连续的可转换 TensorRT 操作的最小数量

  • torch_executed_ops (Collection[Target]) – 必须在 PyTorch 中运行的 aten 操作符集合。如果此集合不为空但 require_full_compilation 为 True，则会抛出错误

  • torch_executed_modules (List[str]) – 必须在 PyTorch 中运行的模块列表。如果此列表不为空但 require_full_compilation 为 True，则会抛出错误

  • pass_through_build_failures (bool) – 如果在编译过程中出现问题，则报错（仅适用于 torch.compile 工作流）

  • max_aux_stream (Optional[python:int]) – 引擎中的最大流数

  • version_compatible (bool) – 构建与未来版本 TensorRT 兼容的 TensorRT 引擎（限制为精简运行时操作符以提供引擎的向前版本兼容性）

  • optimization_level – (Optional[int]): 设置更高的优化级别允许 TensorRT 花费更长的引擎构建时间来搜索更多优化选项。与使用较低优化级别构建的引擎相比，生成的引擎可能具有更好的性能。默认优化级别为 3。有效值包括从 0 到当前最大优化级别 5 的整数。将其设置为大于最大级别会产生与最大级别相同的行为。

  • use_python_runtime – (bool): 返回使用纯 Python 运行时的图，减少了序列化的选项

  • use_fast_partitioner – (bool): 使用基于邻接的分区方案代替全局分区器。邻接分区速度更快，但可能不是最优的。如果寻求最佳性能，请使用全局分区器 (False)。

  • enable_experimental_decompositions (bool) – 使用全套操作符分解。这些分解可能未经测试，但有助于使图更容易转换为 TensorRT，从而可能增加在 TensorRT 中运行的图的数量。

  • dryrun (bool) – “演练”模式的开关，运行除转换为 TRT 和记录输出之外的所有内容

  • hardware_compatible (bool) – 构建与构建引擎的 GPU 之外的 GPU 架构兼容的 TensorRT 引擎（目前适用于 NVIDIA Ampere 及更新版本）

  • timing_cache_path (str) – 计时缓存的路径（如果存在）或编译后将保存到的位置

  • lazy_engine_init (bool) – 推迟设置引擎，直到所有引擎的编译完成。可以允许具有多个图中断的较大型号进行编译，但可能导致运行时 GPU 内存超额订阅。

  • enabled_precisions (Set(Union(torch.dpython:type, torch_tensorrt.dpython:type))) – TensorRT 在选择内核时可以使用的数据类型集合

  • **kwargs – 任何,

返回

MutableTorchTensorRTModule

compile() → None[源代码]

使用 PyTorch 模块（重新）编译 TRT 图模块。每当权重结构发生变化（形状、更多层……）时，都应调用此函数。MutableTorchTensorRTModule 会自动捕获权重值的更新并调用此函数进行重新编译。如果未能捕获到更改，请手动调用此函数以重新编译 TRT 图模块。

refit_gm() → None[源代码]

用任何更新来重装 TRT 图模块。每当权重值发生变化但权重结构保持不变时，都应调用此函数。MutableTorchTensorRTModule 会自动捕获权重值的更新并调用此函数来重装模块。如果未能捕获到更改，请手动调用此函数以更新 TRT 图模块。

set_expected_dynamic_shape_range(args_dynamic_shape: tuple[dict[Any

,

Any]], kwargs_dynamic_shape: dict[str

,

Any]) → None[源代码]

设置动态形状范围。形状提示应严格遵循传递给 forward 函数的 arg_inputs 和 kwarg_inputs，且不应省略任何条目（kwarg_inputs 中的 None 除外）。如果输入中存在嵌套的 dict/list，则该条目的动态形状也应为嵌套的 dict/list。如果某个输入不需要动态形状，则应为其提供一个空字典作为形状提示。请注意，您应排除值为 None 的关键字参数，因为它们将被过滤掉。

示例：def forward(a, b, c=0, d=0)

pass

seq_len = torch.export.Dim("seq_len", min=1, max=10) args_dynamic_shape = ({0: seq_len}, {}) # b 没有动态形状 kwargs_dynamic_shape = {'c': {0, seq_len}, 'd': {}} # d 没有动态形状 set_expected_dynamic_shape_range(args_dynamic_shape, kwargs_dynamic_shape) # 稍后当您调用函数 forward(*(a, b), **{c:…, d:…})

参考: https://pytorch.ac.cn/docs/stable/export.html#expressing-dynamism :param args_dynamic_shape: arg_inputs 的动态形状提示, :type args_dynamic_shape: tuple[dict[Any, Any]] :param kwargs_dynamic_shape: (dict[str, Any]): kwarg_inputs 的动态形状提示

set_weight_streaming_ctx(requested_budget: Optional[int] = None) → None[源代码]

设置权重流预算。如果未设置预算，则使用自动权重流预算。

class torch_tensorrt.Input(*args: Any, **kwargs: Any)[源代码]

根据预期的形状、数据类型和张量格式定义模块的输入。

变量

  • shape_mode (torch_tensorrt.Input._ShapeMode) – 输入是静态形状还是动态形状

  • shape (Tuple or Dict) –

    单个元组或定义输入形状的元组字典。静态形状的输入将有一个单独的元组。动态输入将有一个形式如下的字典

    {"min_shape": Tuple, "opt_shape": Tuple, "max_shape": Tuple}

  • dtype (torch_tensorrt.dpython:type) – 输入张量的预期数据类型（默认：torch_tensorrt.dtype.float32）

  • format (torch_tensorrt.TensorFormat) – 输入张量的预期格式（默认：torch_tensorrt.TensorFormat.NCHW）

__init__(*args: Any, **kwargs: Any) → None[源代码]

torch_tensorrt.Input 的 __init__ 方法

Input 接受以下几种构造模式之一

参数

shape (Tuple or List, optional) – 输入张量的静态形状

关键字参数

  • shape (Tuple or List, optional) – 输入张量的静态形状

  • min_shape (Tuple or List, optional) – 输入张量形状范围的最小值。注意：min_shape、opt_shape、max_shape 必须同时提供，不能有位置参数，shape 不能被定义，并且这会隐式地将 Input 的 shape_mode 设置为 DYNAMIC

  • opt_shape (Tuple or List, optional) – 输入张量形状范围的最佳大小。注意：min_shape、opt_shape、max_shape 必须同时提供，不能有位置参数，shape 不能被定义，并且这会隐式地将 Input 的 shape_mode 设置为 DYNAMIC

  • max_shape (Tuple or List, optional) – 输入张量形状范围的最大大小。注意：min_shape、opt_shape、max_shape 必须同时提供，不能有位置参数，shape 不能被定义，并且这会隐式地将 Input 的 shape_mode 设置为 DYNAMIC

  • dtype (torch.dpython:type or torch_tensorrt.dpython:type) – 输入张量的预期数据类型（默认：torch_tensorrt.dtype.float32）

  • format (torch.memory_format or torch_tensorrt.TensorFormat) – 输入张量的预期格式（默认：torch_tensorrt.TensorFormat.NCHW）

  • tensor_domain (Tuple(python:float, python:float), optional) – 张量允许值的域，以区间表示：[tensor_domain[0], tensor_domain[1])。注意：输入 "None"（或不指定）将把界限设置为 [0, 2)

  • torch_tensor (torch.Tensor) – 持有与此 Input 对应的 torch 张量。

  • name (str, optional) – 此输入在输入 nn.Module 的 forward 函数中的名称。用于在 dynamo 跟踪器中为相应输入指定动态形状。

示例

  • Input([1,3,32,32], dtype=torch.float32, format=torch.channel_last)

  • Input(shape=(1,3,32,32), dtype=torch_tensorrt.dtype.int32, format=torch_tensorrt.TensorFormat.NCHW)

  • Input(min_shape=(1,3,32,32), opt_shape=[2,3,32,32], max_shape=(3,3,32,32)) # 隐式 dtype=torch_tensorrt.dtype.float32, format=torch_tensorrt.TensorFormat.NCHW

example_tensor(optimization_profile_field: Optional[str] = None) → Tensor[源代码]

获取一个由 Input 对象指定的形状的示例张量

参数

optimization_profile_field (Optional(str)) – 在 Input 是动态形状的情况下，用于形状的字段名称

返回

一个 PyTorch 张量

classmethod from_tensor(t: Tensor, disable_memory_format_check: bool = False) → Input[源代码]

生成一个包含给定 PyTorch 张量信息的 Input 对象。

参数

  • tensor (torch.Tensor) – 一个 PyTorch 张量。

  • disable_memory_format_check (bool) – 是否验证输入张量的内存格式

返回

一个 Input 对象。

classmethod from_tensors(ts: Sequence[Tensor], disable_memory_format_check: bool = False) → List[Input][源代码]

生成一个 Input 列表，其中包含所有给定 PyTorch 张量的信息。

参数

  • tensors (Iterable[torch.Tensor]) – 一个 PyTorch 张量列表。

  • disable_memory_format_check (bool) – 是否验证输入张量的内存格式

返回

一个 Input 列表。

dtype: dtype = 1

torch_tensorrt.dtype.float32)

类型

输入张量的预期数据类型（默认

format: memory_format = 1

torch_tensorrt.memory_format.linear)

类型

输入张量的预期格式（默认

class torch_tensorrt.Device(*args: Any, **kwargs: Any)[源代码]

定义可用于为引擎指定目标设备的设备

变量

  • device_type (DeviceType) – 目标设备类型（GPU 或 DLA）。根据是否指定 dla_core 隐式设置。

  • gpu_id (python:int) – 目标 GPU 的设备 ID

  • dla_core (python:int) – 目标 DLA 核心的核心 ID

  • allow_gpu_fallback (bool) – 在 DLA 不支持某个操作时，是否允许回退到 GPU

__init__(*args: Any, **kwargs: Any)[源代码]

torch_tensorrt.Device 的 __init__ 方法

Device 接受以下几种构造模式之一

参数

spec (str) – 带有设备规格的字符串，例如 “dla:0” 表示 DLA，核心 ID 为 0

关键字参数

  • gpu_id (python:int) – 目标 GPU 的 ID（如果指定了 dla_core，则会被管理 DLA 的 GPU 覆盖）。如果指定，则不应提供位置参数

  • dla_core (python:int) – 目标 DLA 核心的 ID。如果指定，则不应提供位置参数。

  • allow_gpu_fallback (bool) – 如果操作在 DLA 上不受支持，则允许 TensorRT 在 GPU 上调度操作（如果设备类型不是 DLA，则忽略）

示例

  • Device(“gpu:1”)

  • Device(“cuda:1”)

  • Device(“dla:0”, allow_gpu_fallback=True)

  • Device(gpu_id=0, dla_core=0, allow_gpu_fallback=True)

  • Device(dla_core=0, allow_gpu_fallback=True)

  • Device(gpu_id=1)

device_type: DeviceType = 1

目标设备类型（GPU 或 DLA）。根据是否指定 dla_core 隐式设置。

dla_core: int = -1

目标 DLA 核心的核心 ID

gpu_id: int = -1

目标 GPU 的设备 ID

枚举

class torch_tensorrt.dtype(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)[源代码]

用于向 Torch-TensorRT 描述数据类型的枚举，与 torch、tensorrt 和 numpy 的 dtypes 兼容

to(t: Union[Type[dtype], Type[DataType], Type[dtype], Type[dtype]], use_default: bool = False) → Union[dtype, DataType, dtype, dtype][源代码]

将 dtype 转换为 [torch, numpy, tensorrt] 中的等效类型

将 self 转换为 numpy、torch 和 tensorrt 等效 dtypes 之一。如果 self 在目标库中不受支持，则会引发异常。因此，不建议直接使用此方法。

或者使用 torch_tensorrt.dtype.try_to()

参数

  • t (Union(Type(torch.dpython:type), Type(tensorrt.DataType), Type(numpy.dpython:type), Type(dpython:type))) – 来自其他库的数据类型枚举，用于转换

  • use_default (bool) – 在某些情况下，一个通用的类型（例如 torch.float）就足够了，因此不抛出异常，而是返回默认值。

返回

来自库枚举 t 的 dtype 等效 torch_tensorrt.dtype

返回类型

Union(torch.dtype, tensorrt.DataType, numpy.dtype, dtype)

引发

TypeError – 不支持的数据类型或未知目标

示例

# Succeeds
float_dtype = torch_tensorrt.dtype.f32.to(torch.dtype) # Returns torch.float

# Failure
float_dtype = torch_tensorrt.dtype.bf16.to(numpy.dtype) # Throws exception
classmethod try_from(t: Union[dtype, DataType, dtype, dtype], use_default: bool = False) → Optional[dtype][源代码]

从其他库的 dtype 系统创建一个 Torch-TensorRT dtype。

从 numpy、torch 和 tensorrt 之一中获取一个 dtype 枚举，并创建一个 torch_tensorrt.dtype。如果源 dtype 系统不受支持或该类型在 Torch-TensorRT 中不受支持，则返回 None。

参数

  • t (Union(torch.dpython:type, tensorrt.DataType, numpy.dpython:type, dpython:type)) – 来自其他库的数据类型枚举

  • use_default (bool) – 在某些情况下，一个通用的类型（例如 torch_tensorrt.dtype.f32）就足够了，因此不抛出异常，而是返回默认值。

返回

等效于 t 的 torch_tensorrt.dtype 或 None

返回类型

Optional(dtype)

示例

# Succeeds
float_dtype = torch_tensorrt.dtype.try_from(torch.float) # Returns torch_tensorrt.dtype.f32

# Unsupported type
float_dtype = torch_tensorrt.dtype.try_from(torch.complex128) # Returns None
try_to(t: Union[Type[dtype], Type[DataType], Type[dtype], Type[dtype]], use_default: bool) → Optional[Union[dtype, DataType, dtype, dtype]][源代码]

将 dtype 转换为 [torch, numpy, tensorrt] 中的等效类型

将 self 转换为 numpy、torch 和 tensorrt 等效 dtypes 之一。如果 self 在目标库中不受支持，则返回 None。

参数

  • t (Union(Type(torch.dpython:type), Type(tensorrt.DataType), Type(numpy.dpython:type), Type(dpython:type))) – 来自其他库的数据类型枚举，用于转换

  • use_default (bool) – 在某些情况下，一个通用的类型（例如 torch.float）就足够了，因此不抛出异常，而是返回默认值。

返回

来自库枚举 t 的 dtype 等效 torch_tensorrt.dtype

返回类型

Optional(Union(torch.dtype, tensorrt.DataType, numpy.dtype, dtype))

示例

# Succeeds
float_dtype = torch_tensorrt.dtype.f32.to(torch.dtype) # Returns torch.float

# Failure
float_dtype = torch_tensorrt.dtype.bf16.to(numpy.dtype) # Returns None
b

布尔值，等效于 dtype.bool

bf16

16位“大脑”浮点数，等效于 dtype.bfloat16

f16

16位浮点数，等效于 dtype.half、dtype.fp16 和 dtype.float16

f32

32位浮点数，等效于 dtype.float、dtype.fp32 和 dtype.float32

f4

4位浮点数，等同于 dtype.fp4 和 dtype.float4

f64

64位浮点数，等效于 dtype.double、dtype.fp64 和 dtype.float64

f8

8位浮点数，等同于 dtype.fp8 和 dtype.float8

i32

有符号32位整数，等效于 dtype.int32 和 dtype.int

i64

有符号64位整数，等效于 dtype.int64 和 dtype.long

i8

有符号8位整数，等效于 dtype.int8，当作为内核精度启用时，通常要求模型支持量化

u8

无符号8位整数，等效于 dtype.uint8

unknown

哨兵值

class torch_tensorrt.DeviceType(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)[源代码]

TensorRT 将要定位的设备类型

to(t: Union[Type[DeviceType], Type[DeviceType]], use_default: bool = False) → Union[DeviceType, DeviceType][源代码]

将 DeviceType 转换为 tensorrt 中的等效类型

将 self 转换为 torch 或 tensorrt 等效设备类型之一。如果 self 在目标库中不受支持，则会引发异常。因此，不建议直接使用此方法。

或者使用 torch_tensorrt.DeviceType.try_to()

参数

t (Union(Type(tensorrt.DeviceType), Type(DeviceType))) – 来自其他库的设备类型枚举，用于转换

返回

枚举 t 中等效的 torch_tensorrt.DeviceType 设备类型

返回类型

Union(tensorrt.DeviceType, DeviceType)

引发

TypeError – 未知目标类型或不支持的设备类型

示例

# Succeeds
trt_dla = torch_tensorrt.DeviceType.DLA.to(tensorrt.DeviceType) # Returns tensorrt.DeviceType.DLA
classmethod try_from(d: Union[DeviceType, DeviceType]) → Optional[DeviceType][源代码]

从 TensorRT 设备类型枚举创建一个 Torch-TensorRT 设备类型枚举。

从 tensorrt 获取一个设备类型枚举，并创建一个 torch_tensorrt.DeviceType。如果源不受支持或 Torch-TensorRT 不支持该设备类型，则会引发异常。因此，不建议直接使用此方法。

或者，使用 torch_tensorrt.DeviceType.try_from()

参数

d (Union(tensorrt.DeviceType, DeviceType)) – 来自其他库的设备类型枚举

返回

与 d 等效的 torch_tensorrt.DeviceType

返回类型

DeviceType

示例

torchtrt_dla = torch_tensorrt.DeviceType._from(tensorrt.DeviceType.DLA)
try_to(t: Union[Type[DeviceType], Type[DeviceType]], use_default: bool = False) → Optional[Union[DeviceType, DeviceType]][源代码]

将 DeviceType 转换为 tensorrt 中的等效类型

将 self 转换为 torch 或 tensorrt 等效的内存格式之一。如果目标库不支持 self，则将返回 None。

参数

t (Union(Type(tensorrt.DeviceType), Type(DeviceType))) – 来自其他库的设备类型枚举，用于转换

返回

枚举 t 中等效的 torch_tensorrt.DeviceType 设备类型

返回类型

Optional(Union(tensorrt.DeviceType, DeviceType))

示例

# Succeeds
trt_dla = torch_tensorrt.DeviceType.DLA.to(tensorrt.DeviceType) # Returns tensorrt.DeviceType.DLA
DLA

目标是 DLA 核心

GPU

目标是 GPU

UNKNOWN

哨兵值

class torch_tensorrt.EngineCapability(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)[源代码]

EngineCapability 决定了网络在构建时的限制以及它所针对的运行时。

to(t: Union[Type[EngineCapability], Type[EngineCapability]]) → Union[EngineCapability, EngineCapability][源代码]

将 EngineCapability 转换为 tensorrt 中的等效类型

将 self 转换为 torch 或 tensorrt 等效的引擎能力之一。如果目标库不支持 self，则会引发异常。因此，不建议直接使用此方法。

或者，使用 torch_tensorrt.EngineCapability.try_to()

参数

t (Union(Type(tensorrt.EngineCapability), Type(EngineCapability))) – 要转换到的其他库的引擎能力枚举

返回

在枚举 t 中与 torch_tensorrt.EngineCapability 等效的引擎能力

返回类型

Union(tensorrt.EngineCapability, EngineCapability)

引发

TypeError – 未知目标类型或不支持的引擎能力

示例

# Succeeds
torchtrt_dla_ec = torch_tensorrt.EngineCapability.DLA_STANDALONE.to(tensorrt.EngineCapability) # Returns tensorrt.EngineCapability.DLA
classmethod try_from() → Optional[EngineCapability][源代码]

从 TensorRT 引擎能力枚举创建一个 Torch-TensorRT 引擎能力枚举。

从 tensorrt 获取一个设备类型枚举，并创建一个 torch_tensorrt.EngineCapability。如果源不受支持或 Torch-TensorRT 不支持该引擎能力级别，则会引发异常。因此，不建议直接使用此方法。

或者，使用 torch_tensorrt.EngineCapability.try_from()

参数

c (Union(tensorrt.EngineCapability, EngineCapability)) – 来自其他库的引擎能力枚举

返回

与 c 等效的 torch_tensorrt.EngineCapability

返回类型

EngineCapability

示例

torchtrt_safety_ec = torch_tensorrt.EngineCapability._from(tensorrt.EngineCapability.SAEFTY)
try_to(t: Union[Type[EngineCapability], Type[EngineCapability]]) → Optional[Union[EngineCapability, EngineCapability]][源代码]

将 EngineCapability 转换为 tensorrt 中的等效类型

将 self 转换为 torch 或 tensorrt 等效的引擎能力之一。如果目标库不支持 self，则将返回 None。

参数

t (Union(Type(tensorrt.EngineCapability), Type(EngineCapability))) – 要转换到的其他库的引擎能力枚举

返回

在枚举 t 中与 torch_tensorrt.EngineCapability 等效的引擎能力

返回类型

Optional(Union(tensorrt.EngineCapability, EngineCapability))

示例

# Succeeds
trt_dla_ec = torch_tensorrt.EngineCapability.DLA.to(tensorrt.EngineCapability) # Returns tensorrt.EngineCapability.DLA_STANDALONE
DLA_STANDALONE

EngineCapability.DLA_STANDALONE 提供了一个受限的网络操作子集，这些操作与 DLA 兼容，并且生成的序列化引擎可以使用独立的 DLA 运行时 API 执行。

SAFETY

EngineCapability.SAFETY 提供了一个经过安全认证的受限网络操作子集，生成的序列化引擎可以在 tensorrt.safe 命名空间中使用 TensorRT 的安全运行时 API 执行。

STANDARD

EngineCapability.STANDARD 不对功能施加任何限制，生成的序列化引擎可以使用 TensorRT 的标准运行时 API 执行。

class torch_tensorrt.memory_format(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)[源代码]
to(t: Union[Type[memory_format], Type[TensorFormat], Type[memory_format]]) → Union[memory_format, TensorFormat, memory_format][源代码]

将 memory_format 转换为 torch 或 tensorrt 中的等效类型

将 self 转换为 torch 或 tensorrt 等效的内存格式之一。如果目标库不支持 self，则会引发异常。因此，不建议直接使用此方法。

或者，使用 torch_tensorrt.memory_format.try_to()

参数

t (Union(Type(torch.memory_format), Type(tensorrt.TensorFormat), Type(memory_format))) – 要转换到的其他库的内存格式类型枚举

返回

在枚举 t 中与 torch_tensorrt.memory_format 等效的内存格式

返回类型

Union(torch.memory_format, tensorrt.TensorFormat, memory_format)

引发

TypeError – 未知目标类型或不支持的内存格式

示例

# Succeeds
tf = torch_tensorrt.memory_format.linear.to(torch.dtype) # Returns torch.contiguous
classmethod try_from(f: Union[memory_format, TensorFormat, memory_format]) → Optional[memory_format][源代码]

从其他库的内存格式枚举创建一个 Torch-TensorRT 内存格式枚举。

从 torch 和 tensorrt 中的一个获取内存格式枚举，并创建一个 torch_tensorrt.memory_format。如果源不受支持或 Torch-TensorRT 不支持该内存格式，则将返回 None。

参数

f (Union(torch.memory_format, tensorrt.TensorFormat, memory_format)) – 来自其他库的内存格式枚举

返回

与 f 等效的 torch_tensorrt.memory_format

返回类型

Optional(memory_format)

示例

torchtrt_linear = torch_tensorrt.memory_format.try_from(torch.contiguous)
try_to(t: Union[Type[memory_format], Type[TensorFormat], Type[memory_format]]) → Optional[Union[memory_format, TensorFormat, memory_format]][源代码]

将 memory_format 转换为 torch 或 tensorrt 中的等效类型

将 self 转换为 torch 或 tensorrt 等效的内存格式之一。如果目标库不支持 self，则将返回 None。

参数

t (Union(Type(torch.memory_format), Type(tensorrt.TensorFormat), Type(memory_format))) – 要转换到的其他库的内存格式类型枚举

返回

在枚举 t 中与 torch_tensorrt.memory_format 等效的内存格式

返回类型

Optional(Union(torch.memory_format, tensorrt.TensorFormat, memory_format))

示例

# Succeeds
tf = torch_tensorrt.memory_format.linear.to(torch.dtype) # Returns torch.contiguous
cdhw32

具有 3 个空间维度的三十二宽通道矢量化行主格式。

此格式绑定到 FP16 和 INT8。它仅适用于维度 >= 4 的情况。

对于维度为 {N, C, D, H, W} 的张量，内存布局等效于维度为 [N][(C+31)/32][D][H][W][32] 的 C 数组，其中张量坐标 (n, d, c, h, w) 映射到数组下标 [n][c/32][d][h][w][c%32]。

chw16

十六宽通道矢量化行主格式。

此格式绑定到 FP16。它仅适用于维度 >= 3 的情况。

对于维度为 {N, C, H, W} 的张量，内存布局等效于维度为 [N][(C+15)/16][H][W][16] 的 C 数组，其中张量坐标 (n, c, h, w) 映射到数组下标 [n][c/16][h][w][c%16]。

chw2

二宽通道矢量化行主格式。

在 TensorRT 中，此格式绑定到 FP16。它仅适用于维度 >= 3 的情况。

对于维度为 {N, C, H, W} 的张量，内存布局等效于维度为 [N][(C+1)/2][H][W][2] 的 C 数组，其中张量坐标 (n, c, h, w) 映射到数组下标 [n][c/2][h][w][c%2]。

chw32

三十二宽通道矢量化行主格式。

此格式仅适用于维度 >= 3 的情况。

对于维度为 {N, C, H, W} 的张量，内存布局等效于维度为 [N][(C+31)/32][H][W][32] 的 C 数组，其中张量坐标 (n, c, h, w) 映射到数组下标 [n][c/32][h][w][c%32]。

chw4

四宽通道矢量化行主格式。此格式绑定到 INT8。它仅适用于维度 >= 3 的情况。

对于维度为 {N, C, H, W} 的张量，内存布局等效于维度为 [N][(C+3)/4][H][W][4] 的 C 数组，其中张量坐标 (n, c, h, w) 映射到数组下标 [n][c/4][h][w][c%4]。

dhwc

非矢量化通道最后格式。此格式绑定到 FP32。它仅适用于维度 >= 4 的情况。

等效于 memory_format.channels_last_3d

dhwc8

八通道格式，其中 C 被填充为 8 的倍数。

此格式绑定到 FP16，并且仅适用于维度 >= 4 的情况。

对于维度为 {N, C, D, H, W} 的张量，内存布局等效于维度为 [N][D][H][W][(C+7)/8*8] 的数组，其中张量坐标 (n, c, d, h, w) 映射到数组下标 [n][d][h][w][c]。

dla_hwc4

DLA 图像格式。通道最后格式。C 只能是 1、3、4。如果 C == 3，它将被四舍五入到 4。沿 H 轴步进的步幅向上舍入到 32 字节。

此格式绑定到 FP16/Int8，并且仅适用于维度 >= 3 的情况。

对于维度为 {N, C, H, W} 的张量，当 C 分别为 1、3、4 时，C’ 为 1、4、4，内存布局等效于维度为 [N][H][roundUp(W, 32/C’/elementSize)][C’] 的 C 数组，其中 elementSize 对于 FP16 为 2，对于 Int8 为 1，C’ 是舍入后的 C。张量坐标 (n, c, h, w) 映射到数组下标 [n][h][w][c]。

dla_linear

DLA 平面格式。行主格式。沿 H 轴步进的步幅向上舍入到 64 字节。

此格式绑定到 FP16/Int8，并且仅适用于维度 >= 3 的情况。

对于维度为 {N, C, H, W} 的张量，内存布局等效于维度为 [N][C][H][roundUp(W, 64/elementSize)] 的 C 数组，其中 elementSize 对于 FP16 为 2，对于 Int8 为 1，张量坐标 (n, c, h, w) 映射到数组下标 [n][c][h][w]。

hwc

非矢量化通道最后格式。此格式绑定到 FP32，并且仅适用于维度 >= 3 的情况。

等效于 memory_format.channels_last

hwc16

十六通道格式，其中 C 被填充为 16 的倍数。此格式绑定到 FP16。它仅适用于维度 >= 3 的情况。

对于维度为 {N, C, H, W} 的张量，内存布局等效于维度为 [N][H][W][(C+15)/16*16] 的数组，其中张量坐标 (n, c, h, w) 映射到数组下标 [n][h][w][c]。

hwc8

八通道格式，其中 C 被填充为 8 的倍数。

此格式绑定到 FP16。它仅适用于维度 >= 3 的情况。

对于维度为 {N, C, H, W} 的张量，内存布局等效于维度为 [N][H][W][(C+7)/8*8] 的数组，其中张量坐标 (n, c, h, w) 映射到数组下标 [n][h][w][c]。

linear

行主线性格式。

对于维度为 {N, C, H, W} 的张量，W 轴始终具有单位步幅，而其他每个轴的步幅至少是下一个维度乘以下一个步幅的乘积。步幅与维度为 [N][C][H][W] 的 C 数组相同。

等效于 memory_format.contiguous

子模块

文档

访问全面的 PyTorch 开发者文档

查看文档

教程

为初学者和高级开发者提供深入的教程

查看教程

资源

查找开发资源并让您的问题得到解答

查看资源

为了分析流量和优化您的体验，我们在本网站上使用 cookie。通过点击或导航，您同意我们使用 cookie。作为本网站的当前维护者，Facebook 的 Cookie 政策适用。了解更多信息，包括可用的控制选项：Cookie 政策。