目录
快捷方式

torch_tensorrt.ts

函数

torch_tensorrt.ts.compile(module: ScriptModule, inputs: Optional[Sequence[Input | torch.Tensor]] = None, input_signature: Optional[Tuple[Union[Input, Tensor, Sequence[Any]]]] = None, device: Device = Device(type=DeviceType.GPU, gpu_id=0), disable_tf32: bool = False, sparse_weights: bool = False, enabled_precisions: Optional[Set[Union[dtype, dtype]]] = None, refit: bool = False, debug: bool = False, capability: EngineCapability = EngineCapability.STANDARD, 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, calibrator: object = None, truncate_long_and_double: bool = False, require_full_compilation: bool = False, min_block_size: int = 3, torch_executed_ops: Optional[List[str]] = None, torch_executed_modules: Optional[List[str]] = None, allow_shape_tensors: bool = False) → ScriptModule[source]

使用 TensorRT 将 TorchScript 模块编译到 NVIDIA GPU 上

接收一个现有的 TorchScript 模块和一组设置来配置编译器，并将方法转换为调用等效 TensorRT 引擎的 JIT 图

具体转换 TorchScript 模块的 forward 方法

参数

module (torch.jit.ScriptModule) – 源模块，是跟踪或脚本化 PyTorch torch.nn.Module 的结果

关键字参数

  • 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
]

  • Union (input_signature) –

    模块输入规范的格式化集合。输入大小可以指定为 torch 大小、元组或列表。数据类型可以使用 torch 数据类型或 torch_tensorrt 数据类型指定，并且可以使用 torch 设备或 torch_tensorrt 设备类型枚举来选择设备类型。**此 API 应被视为 beta 级稳定，将来可能会更改**

    input_signature=([
    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 for input #3

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

    TensorRT 引擎运行的目标设备

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

  • disable_tf32 (bool) – 强制 FP32 层使用传统 FP32 格式，而不是默认行为（在乘法之前将输入舍入到 10 位尾数，但使用 23 位尾数累加总和）

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

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

  • refit (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 用于在操作之间共享中间张量数据的 Host RAM

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

  • truncate_long_and_double (bool) – 将以 int64 或 double (float64) 提供的权重截断为 int32 和 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 (List[str]) – 必须在 PyTorch 中运行的 aten 运算符列表。如果此列表不为空但 require_full_compilation 为 True，则会抛出错误

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

  • allow_shape_tensors –（实验性）允许 aten::size 使用 TensorRT 中的 IShapeLayer 输出形状张量

返回值

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

返回类型

torch.jit.ScriptModule

torch_tensorrt.ts.convert_method_to_trt_engine(module: ScriptModule, method_name: str = 'forward', inputs: Optional[Sequence[Input | torch.Tensor]] = None, device: Device = Device(type=DeviceType.GPU, gpu_id=0), disable_tf32: bool = False, sparse_weights: bool = False, enabled_precisions: Optional[Set[Union[dtype, dtype]]] = None, refit: bool = False, debug: bool = False, capability: EngineCapability = EngineCapability.STANDARD, 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_long_and_double: int = False, calibrator: object = None, allow_shape_tensors: bool = False) → bytes[source]

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

根据转换设置字典，将模块的指定方法转换为序列化后的 TensorRT 引擎。

参数

module (torch.jit.ScriptModule) – 源模块，是跟踪或脚本化 PyTorch torch.nn.Module 的结果

关键字参数

  • 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
]

  • method_name (str) – 要转换的方法名称

  • Union (input_signature) –

    模块输入规范的格式化集合。输入大小可以指定为 torch 大小、元组或列表。数据类型可以使用 torch 数据类型或 torch_tensorrt 数据类型指定，并且可以使用 torch 设备或 torch_tensorrt 设备类型枚举来选择设备类型。**此 API 应被视为 beta 级稳定，将来可能会更改**

    input_signature=([
    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 for input #3

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

    TensorRT 引擎运行的目标设备

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

  • disable_tf32 (bool) – 强制 FP32 层使用传统 FP32 格式，而不是默认行为（在乘法之前将输入舍入到 10 位尾数，但使用 23 位尾数累加总和）

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

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

  • refit (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 用于在操作之间共享中间张量数据的 Host RAM

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

  • truncate_long_and_double (bool) – 将以 int64 或 double (float64) 提供的权重截断为 int32 和 float32

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

  • allow_shape_tensors –（实验性）允许 aten::size 使用 TensorRT 中的 IShapeLayer 输出形状张量

返回值

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

返回类型

字节

torch_tensorrt.ts.check_method_op_support(module: ScriptModule, method_name: str = 'forward') → bool[source]

检查方法是否完全受 torch_tensorrt 支持

检查 TorchScript 模块的方法是否可以由 torch_tensorrt 编译，如果不能，则打印出不支持的操作符列表并返回 false，否则返回 true。

参数

  • module (torch.jit.ScriptModule) – 源模块，是跟踪或脚本化 PyTorch torch.nn.Module 的结果

  • method_name (str) – 要检查的方法名称

返回值

如果支持方法，则为 True

返回类型

布尔值

torch_tensorrt.ts.embed_engine_in_new_module(serialized_engine: bytes, input_binding_names: Optional[List[str]] = None, output_binding_names: Optional[List[str]] = None, device: Device = Device(type=DeviceType.GPU, gpu_id=0)) → ScriptModule[source]

获取预构建的序列化 TensorRT 引擎并将其嵌入到 TorchScript 模块中。

获取预构建的序列化 TensorRT 引擎（以字节表示）并将其嵌入到 TorchScript 模块中。注册 forward 方法以使用以下函数签名执行 TensorRT 引擎：

forward(Tensor[]) -> Tensor[]

可以使用 [in/out]put_binding_names 显式指定 TensorRT 绑定，或者使用以下格式的名称：

  • [符号].[输入/输出数组中的索引]

例如：- [x.0, x.1, x.2] -> [y.0]

可以将包含嵌入引擎的模块保存为 torch.jit.save，并根据 torch_tensorrt 的可移植性规则移动/加载。

参数

serialized_engine (bytearray) – 来自 torch_tensorrt 或 TensorRT API 的序列化 TensorRT 引擎。

关键字参数

  • input_binding_names (List[str]) – 要传递到包含 PyTorch 模块的 TensorRT 绑定的名称列表，按顺序排列。

  • output_binding_names (List[str]) – 包含 PyTorch 模块应返回的 TensorRT 绑定的名称列表，按顺序排列。

  • device (Union(Device, torch.device, dict)) – 运行引擎的目标设备。必须与提供的引擎兼容。默认值：当前活动设备。

返回值

包含嵌入引擎的新 TorchScript 模块。

返回类型

torch.jit.ScriptModule

torch_tensorrt.ts.TensorRTCompileSpec(inputs: Optional[List[torch.Tensor | Input]] = None, input_signature: Optional[Any] = None, device: torch.device | Device = Device(type=DeviceType.GPU, gpu_id=0), disable_tf32: bool = False, sparse_weights: bool = False, enabled_precisions: Optional[Set[Union[dtype, dtype]]] = None, refit: bool = False, debug: bool = False, capability: EngineCapability = EngineCapability.STANDARD, 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_long_and_double: bool = False, calibrator: object = None, allow_shape_tensors: bool = False) → <torch.ScriptClass object at 0x7ff4f666d170>[source]

用于创建格式化规范字典，以便使用 PyTorch TensorRT 后端

关键字参数

  • 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
]

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

    TensorRT 引擎运行的目标设备

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

  • disable_tf32 (bool) – 强制 FP32 层使用传统 FP32 格式，而不是默认行为（在乘法之前将输入舍入到 10 位尾数，但使用 23 位尾数累加总和）

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

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

  • refit (bool) – 启用重新拟合

  • debug (bool) – 启用可调试引擎

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

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

  • workspace_size (python:int) – 提供给 TensorRT 的工作区最大大小

  • truncate_long_and_double (bool) – 将以 int64 或 double (float64) 提供的权重截断为 int32 和 float32

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

  • allow_shape_tensors –

    (实验性) 允许 aten::size 使用 TensorRT 中的 IShapeLayer 输出形状张量

    返回值

    torch.classes.tensorrt.CompileSpec：要提供给 torch._C._jit_to_tensorrt 的方法和格式化规范对象的列表

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取针对初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并获得问题的解答

查看资源

为了分析流量并优化您的体验，我们在本网站上提供 Cookie。点击或导航，即表示您同意我们使用 Cookie。作为本网站的当前维护方，Facebook 的 Cookie 政策适用。了解更多信息，包括可用的控件：Cookie 政策。