目录
快捷方式

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[源代码]

使用 TensorRT 为 NVIDIA GPU 编译 TorchScript 模块

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

特别转换 TorchScript 模块的 forward 方法

参数

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

关键字参数

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

    必需 模块输入的输入形状、数据类型和内存布局的规范列表。此参数是必需的。输入大小可以指定为 torch.Size、元组或列表。数据类型可以使用 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.Size、元组或列表。数据类型可以使用 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 用于在操作间共享中间张量数据的主机 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)) – 校准器对象，它将为 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[源代码]

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

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

参数

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

关键字参数

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

    必需 模块输入的输入形状、数据类型和内存布局的规范列表。此参数是必需的。输入大小可以指定为 torch.Size、元组或列表。数据类型可以使用 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.Size、元组或列表。数据类型可以使用 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 用于在操作间共享中间张量数据的主机 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)) – 校准器对象，它将为 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[源代码]

检查一个方法是否完全被 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[源代码]

将一个预构建的序列化 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 0x7f39bceea830>[源代码]

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

关键字参数

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

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