目錄
快捷方式

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 政策。