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]][source]¶
使用 TensorRT 为 NVIDIA GPU 编译 PyTorch 模块
接受现有的 PyTorch 模块和一组设置来配置编译器，并使用
ir中指定的路径来降低模块的级别并将其编译为 TensorRT，从而返回 PyTorch 模块专门转换模块的 forward 方法
- 参数
module (Union(torch.nn.Module,torch.jit.ScriptModule) – 源模块
- 关键字参数
inputs (List[Union(Input, torch.Tensor)]) –
必需 模块输入的输入形状、dtype 和内存布局的规范列表。此参数是必需的。输入大小可以指定为 torch 大小、元组或列表。dtype 可以使用 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 函数的 kwarg 输入。
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[source]¶
将 TorchScript 模块方法转换为序列化的 TensorRT 引擎
给定转换设置字典，将模块的指定方法转换为序列化的 TensorRT 引擎
- 参数
module (Union(torch.nn.Module,torch.jit.ScriptModule) – 源模块
- 关键字参数
inputs (List[Union(Input, torch.Tensor)]) –
必需 模块输入的输入形状、dtype 和内存布局的规范列表。此参数是必需的。输入大小可以指定为 torch 大小、元组或列表。dtype 可以使用 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 函数的 kwarg 输入。
enabled_precision (Set(Union(torch.dpython:type, torch_tensorrt.dpython:type))) – TensorRT 在选择内核时可以使用的数据类型集
ir (str) – 请求的编译策略。（选项：default - 让 Torch-TensorRT 决定，ts - 使用脚本路径的 TorchScript）
**kwargs – 特定请求策略的附加设置（有关更多信息，请参阅子模块）
- 返回值
序列化的 TensorRT 引擎，可以保存到文件或通过 TensorRT API 反序列化
- 返回类型
bytes
- 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) None[source]¶
以指定的输出格式将模型保存到磁盘。
- 参数
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 函数的 kwarg 输入。
output_format (str) – 保存模型的格式。选项包括 exported_program | torchscript。
retrace (bool) – 当模块类型为 fx.GraphModule 时，此选项使用 torch.export.export(strict=False) 重新导出图以保存它。此标志目前为实验性。
类¶
- class torch_tensorrt.MutableTorchTensorRTModule(pytorch_model: Module, *, device: Optional[Union[Device, device, str]] = None, disable_tf32: bool = False, assume_dynamic_shape_support: bool = False, sparse_weights: bool = False, enabled_precisions: Set[Union[dtype, dtype]]] = {dtype.f32}, engine_capability: EngineCapability = EngineCapability.STANDARD, immutable_weights: bool = True, debug: bool = False, num_avg_timing_iters: int = 1, workspace_size: int = 0, dla_sram_size: int = 1048576, dla_local_dram_size: int = 1073741824, dla_global_dram_size: int = 536870912, truncate_double: bool = False, require_full_compilation: bool = False, min_block_size: int = 5, torch_executed_ops: Optional[Collection[Union[Callable[[...], Any], str]]] = None, torch_executed_modules: Optional[List[str]] = None, pass_through_build_failures: bool = False, max_aux_streams: Optional[int] = None, version_compatible: bool = False, optimization_level: Optional[int] = None, use_python_runtime: bool = False, use_fast_partitioner: bool = True, enable_experimental_decompositions: bool = False, dryrun: bool = False, hardware_compatible: bool = False, timing_cache_path: str = '/tmp/torch_tensorrt_engine_cache/timing_cache.bin', **kwargs: Any)[source]¶
初始化 MutableTorchTensorRTModule，以便像常规 PyTorch 模块一样无缝地操作它。所有 TensorRT 编译和重新拟合过程都会在您使用模块时自动处理。对其属性的任何更改或加载不同的 state_dict 都将触发重新拟合或重新编译，这将在下一次正向传递期间进行管理。
MutableTorchTensorRTModule 接受一个 PyTorch 模块以及编译器的配置设置。编译完成后，该模块会保持 TensorRT 图模块和原始 PyTorch 模块之间的连接。对 MutableTorchTensorRTModule 所做的任何修改都将反映在 TensorRT 图模块和原始 PyTorch 模块中。
- __init__(pytorch_model: Module, *, device: Optional[Union[Device, device, str]] = None, disable_tf32: bool = False, assume_dynamic_shape_support: bool = False, sparse_weights: bool = False, enabled_precisions: Set[Union[dtype, dtype]] = {dtype.f32}, engine_capability: EngineCapability = EngineCapability.STANDARD, immutable_weights: bool = True, debug: bool = False, num_avg_timing_iters: int = 1, workspace_size: int = 0, dla_sram_size: int = 1048576, dla_local_dram_size: int = 1073741824, dla_global_dram_size: int = 536870912, truncate_double: bool = False, require_full_compilation: bool = False, min_block_size: int = 5, torch_executed_ops: Optional[Collection[Union[Callable[[...], Any], str]]] = None, torch_executed_modules: Optional[List[str]] = None, pass_through_build_failures: bool = False, max_aux_streams: Optional[/span>int] = None, version_compatible: bool = False, optimization_level: Optional[int] = None, use_python_runtime: bool = False, use_fast_partitioner: bool = True, enable_experimental_decompositions: bool = False, dryrun: bool = False, hardware_compatible: bool = False, timing_cache_path: str = '/tmp/torch_tensorrt_engine_cache/timing_cache.bin', **kwargs: Any) None[source]¶
- 参数
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) – 构建不可重拟合的引擎。这对于某些不可重拟合的层很有用。
debug (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) – “Dryrun” 模式的切换开关，运行除转换为 TRT 和记录输出之外的所有内容
hardware_compatible (bool) – 构建与构建引擎的 GPU 架构以外的 GPU 架构兼容的 TensorRT 引擎（目前适用于 NVIDIA Ampere 及更高版本）
timing_cache_path (str) – 如果存在，则为计时缓存的路径；如果不存在，则为编译后将保存计时缓存的路径
lazy_engine_init (bool) – 延迟引擎设置，直到所有引擎编译完成。可以允许编译具有多个图断点的大型模型，但可能导致运行时 GPU 内存的过度订阅。
**kwargs – Any,
- 返回值
MutableTorchTensorRTModule
- class torch_tensorrt.Input(*args: Any, **kwargs: Any)[source]¶
根据预期形状、数据类型和张量格式定义模块的输入。
- 变量
shape_mode (torch_tensorrt.Input._ShapeMode) – 输入是静态形状还是动态形状
shape (Tuple 或 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[source]¶
torch_tensorrt.Input 的 __init__ 方法
Input 接受以下几种构造模式之一
- 参数
shape (Tuple 或 List, 可选) – 输入张量的静态形状
- 关键字参数
shape (Tuple 或 List, 可选) – 输入张量的静态形状
min_shape (Tuple 或 List, 可选) – 输入张量形状范围的最小尺寸 注意：必须提供 min_shape、opt_shape、max_shape 这三个参数，不得有位置参数，不得定义 shape，并且隐式地将 Input 的 shape_mode 设置为 DYNAMIC
opt_shape (Tuple 或 List, 可选) – 输入张量形状范围的最佳尺寸 注意：必须提供 min_shape、opt_shape、max_shape 这三个参数，不得有位置参数，不得定义 shape，并且隐式地将 Input 的 shape_mode 设置为 DYNAMIC
max_shape (Tuple 或 List, 可选) – 输入张量形状范围的最大尺寸 注意：必须提供 min_shape、opt_shape、max_shape 这三个参数，不得有位置参数，不得定义 shape，并且隐式地将 Input 的 shape_mode 设置为 DYNAMIC
dtype (torch.dpython:type 或 torch_tensorrt.dpython:type) – 输入张量的预期数据类型（默认值：torch_tensorrt.dtype.float32）
format (torch.memory_format 或 torch_tensorrt.TensorFormat) – 输入张量的预期格式（默认值：torch_tensorrt.TensorFormat.NCHW）
tensor_domain (Tuple(python:float, python:float), 可选) – 张量允许值的域，以区间表示法表示：[tensor_domain[0], tensor_domain[1])。注意：输入 “None”（或不指定）会将边界设置为 [0, 2)
torch_tensor (torch.Tensor) – 持有与此 Input 对应的 torch 张量。
name (str, 可选) – 此输入在输入 nn.Module 的 forward 函数中的名称。用于在 dynamo tracer 中为相应的输入指定动态形状。
示例
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[source]¶
获取由 Input 对象指定的形状的示例张量
- 参数
optimization_profile_field (Optional(str)) – 在 Input 为动态形状的情况下，用于形状的字段名称
- 返回值
一个 PyTorch 张量
- classmethod from_tensor(t: Tensor, disable_memory_format_check: bool = False) Input[source]¶
生成一个 Input，其中包含给定 PyTorch 张量的信息。
- 参数
tensor (torch.Tensor) – 一个 PyTorch 张量。
disable_memory_format_check (bool) – 是否验证输入张量的内存格式
- 返回值
一个 Input 对象。
- classmethod from_tensors(ts: Sequence[Tensor], disable_memory_format_check: bool = False) List[Input][source]¶
生成一个 Input 列表，其中包含所有给定 PyTorch 张量的信息。
- 参数
tensors (Iterable[torch.Tensor]) – PyTorch 张量列表。
disable_memory_format_check (bool) – 是否验证输入张量的内存格式
- 返回值
一个 Input 列表。
- format: memory_format = 1¶
torch_tensorrt.memory_format.linear)
- 类型
输入张量的预期格式（默认值
- class torch_tensorrt.Device(*args: Any, **kwargs: Any)[source]¶
定义可用于为引擎指定目标设备的设备
- 变量
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)[source]¶
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 不支持在 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)[source]¶
枚举类型，用于描述 Torch-TensorRT 的数据类型，兼容 torch、tensorrt 和 numpy 的数据类型
- to(t: Union[Type[dtype], Type[DataType], Type[dtype], Type[dtype]], use_default: bool = False) Union[dtype, DataType, dtype, dtype][source]¶
将 dtype 转换为 [torch, numpy, tensorrt] 中等效的类型
将
self转换为 numpy、torch 和 tensorrt 等效的数据类型之一。如果目标库不支持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的等效torch_tensorrt.dtype- 返回类型
Union(torch.dtype, tensorrt.DataType, numpy.dtype, dtype)
- Raises
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])[source]¶
从另一个库的数据类型系统创建 Torch-TensorRT 数据类型。
接受来自 numpy、torch 和 tensorrt 之一的 dtype 枚举，并创建一个
torch_tensorrt.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])])[source]¶
将 dtype 转换为 [torch, numpy, tensorrt] 中等效的类型
将
self转换为 numpy、torch 和 tensorrt 等效的数据类型之一。如果目标库不支持self，则返回None。- 参数
t (Union(Type(torch.dpython:type), Type(tensorrt.DataType), Type(numpy.dpython:type), Type(dpython:type))) – 要转换到的另一个库的数据类型枚举
use_default (bool) – 在某些情况下，可以使用通用的类型（例如
torch.float）就足够了，因此，与其抛出异常，不如返回默认值。
- 返回值
来自库枚举
t的等效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 位 “Brain” 浮点数，等效于
dtype.bfloat16
- f16¶
16 位浮点数，等效于
dtype.half、dtype.fp16和dtype.float16
- f32¶
32 位浮点数，等效于
dtype.float、dtype.fp32和dtype.float32
- 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)[source]¶
TensorRT 将要面向的设备类型
- to(t: Union[Type([DeviceType]), Type([DeviceType])], use_default: bool = False) Union([DeviceType]), DeviceType])[source]¶
将
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)
- Raises
TypeError – 未知的目标类型或不支持的设备类型
示例
# Succeeds trt_dla = torch_tensorrt.DeviceType.DLA.to(tensorrt.DeviceType) # Returns tensorrt.DeviceType.DLA
- classmethod try_from(d: Union([DeviceType]), DeviceType])) Optional([DeviceType])[source]¶
从 TensorRT 设备类型枚举创建 Torch-TensorRT 设备类型枚举。
接受来自 tensorrt 的设备类型枚举，并创建一个
torch_tensorrt.DeviceType。如果不支持源或 Torch-TensorRT 不支持该设备类型，则会引发异常。因此，不建议直接使用此方法。或者使用
torch_tensorrt.DeviceType.try_from()- 参数
d (Union(tensorrt.DeviceType, DeviceType)) – 来自另一个库的设备类型枚举
- 返回值
与
d等效的torch_tensorrt.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])])[source]¶
将
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)[source]¶
EngineCapability 决定了网络在构建时的约束以及它面向的运行时。
- to(t: Union([Type([EngineCapability]), Type([EngineCapability])]) Union([EngineCapability]), EngineCapability])[source]¶
将
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)
- Raises
TypeError – 未知的目标类型或不支持的引擎能力
示例
# Succeeds torchtrt_dla_ec = torch_tensorrt.EngineCapability.DLA_STANDALONE.to(tensorrt.EngineCapability) # Returns tensorrt.EngineCapability.DLA
- classmethod try_from() Optional([EngineCapability])[source]¶
从 TensorRT 引擎能力枚举创建 Torch-TensorRT 引擎能力枚举。
接受来自 tensorrt 的设备类型枚举，并创建一个
torch_tensorrt.EngineCapability。如果不支持源或 Torch-TensorRT 不支持该引擎能力级别，则会引发异常。因此，不建议直接使用此方法。或者使用
torch_tensorrt.EngineCapability.try_from()- 参数
c (Union(tensorrt.EngineCapability, EngineCapability)) – 来自另一个库的引擎能力枚举
- 返回值
与
c等效的torch_tensorrt.EngineCapability- 返回类型
示例
torchtrt_safety_ec = torch_tensorrt.EngineCapability._from(tensorrt.EngineCapability.SAEFTY)
- try_to(t: Union([Type([EngineCapability]), Type([EngineCapability])]) Optional([Union([EngineCapability]), EngineCapability])])[source]¶
将
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 在 tensorrt.safe 命名空间中的安全运行时 API 执行生成的序列化引擎。
- STANDARD¶
EngineCapability.STANDARD 不对功能提供任何限制，并且可以使用 TensorRT 的标准运行时 API 执行生成的序列化引擎。
- class torch_tensorrt.memory_format(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)[source]¶
- to(t: Union([Type([memory_format]), Type([TensorFormat]), Type([memory_format])]) Union([memory_format]), TensorFormat, memory_format])[source]¶
将
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)
- Raises
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][source]¶
从另一个库内存格式枚举创建 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]][source]¶
将
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
