目录
快捷方式

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 编译用于 NVIDIA GPU 的 TorchScript 模块

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

具体转换 TorchScript 模块的 forward 方法

参数

module (torch.jit.ScriptModule) – 源模块，由 PyTorch torch.nn.Module 追踪或脚本化得到

关键字参数

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

    必需的模块输入形状、数据类型和内存布局规范列表。此参数是必需的。输入大小可以指定为 torch sizes、元组或列表。数据类型可以使用 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 sizes、元组或列表。数据类型可以使用 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) – 启用 refitting

  • 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_long_and_double (bool) – 将以 int64 或 double (float64) 提供的权重截断为 int32 和 float32

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

  • 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 sizes、元组或列表。数据类型可以使用 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 sizes、元组或列表。数据类型可以使用 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) – 启用 refitting

  • 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_long_and_double (bool) – 将以 int64 或 double (float64) 提供的权重截断为 int32 和 float32

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

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

返回

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

返回类型

bytes

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

返回类型

bool

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

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

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

例如 - [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]) – TensorRT 绑定的名称列表，以便按顺序传递给包含它的 PyTorch 模块

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

  • 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: Optional[Union[device, Device]] = None, 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 0x7f74a704bb30>[source]

用于为 PyTorch TensorRT 后端创建格式化规范字典的实用工具

关键字参数

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

    必需的模块输入形状、数据类型和内存布局规范列表。此参数是必需的。输入大小可以指定为 torch sizes、元组或列表。数据类型可以使用 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) – 启用 refitting

  • 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)) – Calibrator 对象，将为 PTQ 系统提供数据用于 INT8 校准

  • allow_shape_tensors –

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

    返回

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

文档

访问 PyTorch 的完整开发者文档

查看文档

教程

获取面向初学者和高级开发者的深度教程

查看教程

资源

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

查看资源

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