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 模組和一組設定來配置編譯器，並將方法轉換為 JIT 圖形，這些圖形會呼叫等效的 TensorRT 引擎

特別轉換 TorchScript 模組的 forward 方法

參數

module (torch.jit.ScriptModule) – 來源模組，追蹤或腳本化 PyTorch torch.nn.Module 的結果

關鍵字引數

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
]

Union (input_signature) –

模組輸入規格的格式化集合。輸入大小可以指定為 torch 大小、元組或列表。dtype 可以使用 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)) – 校準器物件，將為 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)]) –

必要模組輸入的輸入形狀、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
]

method_name (str) – 要轉換的方法名稱

Union (input_signature) –

模組輸入規格的格式化集合。輸入大小可以指定為 torch 大小、元組或列表。dtype 可以使用 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)) – 校準器物件，將為 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 引擎 (以 bytes 形式)，並將其嵌入 TorchScript 模組中。註冊 forward 方法以使用函數簽名執行 TensorRT 引擎

forward(Tensor[]) -> Tensor[]

TensorRT 綁定可以使用 [in/out]put_binding_names 明確指定，或具有以下格式的名稱

[symbol].[輸入/輸出陣列中的索引]

例如 - [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 0x7f97927623f0>[source]¶

用於為使用 PyTorch TensorRT 後端創建格式化規格字典的工具

關鍵字引數

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
]

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)) – 校準器物件，將為 PTQ 系統提供資料以進行 INT8 校準
allow_shape_tensors –
(實驗性) 允許 aten::size 使用 TensorRT 中的 IShapeLayer 輸出形狀張量

返回
torch.classes.tensorrt.CompileSpec：要提供給 torch._C._jit_to_tensorrt 的方法列表和格式化規格物件

torch_tensorrt.ts¶

函式¶

文件

教學

資源