目錄
快捷方式

torch_tensorrt

函式

torch_tensorrt.compile(module: Any, ir: str = 'default', inputs: Optional[Sequence[Input | torch.Tensor |InputTensorSpec]] = None, arg_inputs: Optional[Sequence[Sequence[Any]]] = None, kwarg_inputs: Optional[dict[Any, Any]] = None, enabled_precisions: Optional[Set[Union[dtype, dtype]]] = None, **kwargs: Any) → Union[Module, ScriptModule, GraphModule, Callable[[...], Any]][source]

使用 TensorRT 為 NVIDIA GPU 編譯 PyTorch 模組

接受一個現有的 PyTorch 模組和一組配置編譯器的設定，並使用 ir 中指定的路徑將模組降低並編譯為 TensorRT，然後返回一個 PyTorch 模組

專門轉換 Module 的 forward 方法

引數

module (Union(torch.nn.Module,torch.jit.ScriptModule) – 源模組

關鍵字引數

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

    必需 模組輸入的形狀、dtype 和記憶體佈局規範列表。此引數為必需項。輸入大小可以指定為 torch 大小、元組或列表。dtype 可以使用 torch 資料型別或 torch_tensorrt 資料型別來指定，並且可以使用 torch 裝置或 torch_tensorrt 裝置型別列舉來選擇裝置型別。

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

  • arg_inputs (Tuple[Any, ...]) – 與 inputs 相同。為了更好地與 kwarg_inputs 對應，可作為別名。

  • kwarg_inputs (dict[Any, ...]) – 可選，模組 forward 函式的關鍵字引數輸入。

  • enabled_precision (Set(Union(torch.dpython:type, torch_tensorrt.dpython:type))) – TensorRT 在選擇 Kernel 時可以使用的基本資料型別集合

  • ir (str) – 請求的編譯策略。（選項：default - 由 Torch-TensorRT 決定，ts - 使用指令碼化路徑的 TorchScript）

  • **kwargs – 特定請求策略的其他設定（更多資訊請參見子模組）

返回值

編譯後的模組，執行時將透過 TensorRT 執行

返回值型別

torch.nn.Module

torch_tensorrt.convert_method_to_trt_engine(module: Any, method_name: str = 'forward', inputs: Optional[Sequence[Input | torch.Tensor |InputTensorSpec]] = None, arg_inputs: Optional[Sequence[Sequence[Any]]] = None, kwarg_inputs: Optional[dict[Any, Any]] = None, ir: str = 'default', enabled_precisions: Optional[Set[Union[dtype, dtype]]] = None, **kwargs: Any) → bytes[source]

將 TorchScript 模組方法轉換為序列化的 TensorRT 引擎

給定一個轉換設定字典，將模組的指定方法轉換為序列化的 TensorRT 引擎

引數

module (Union(torch.nn.Module,torch.jit.ScriptModule) – 源模組

關鍵字引數

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

  • arg_inputs (Tuple[Any, ...]) – 與 inputs 相同。為了更好地與 kwarg_inputs 對應，可作為別名。

  • kwarg_inputs (dict[Any, ...]) – 可選，模組 forward 函式的關鍵字引數輸入。

  • enabled_precision (Set(Union(torch.dpython:type, torch_tensorrt.dpython:type))) – TensorRT 在選擇 Kernel 時可以使用的基本資料型別集合

  • ir (str) – 請求的編譯策略。（選項：default - 由 Torch-TensorRT 決定，ts - 使用指令碼化路徑的 TorchScript）

  • **kwargs – 特定請求策略的其他設定（更多資訊請參見子模組）

返回值

序列化的 TensorRT 引擎，可以儲存到檔案或透過 TensorRT API 反序列化

返回值型別

bytes

torch_tensorrt.cross_compile_for_windows(module: Module, file_path: str, inputs: Optional[Sequence[Input | torch.Tensor]] = None, arg_inputs: Optional[Sequence[Sequence[Any]]] = None, kwarg_inputs: Optional[dict[Any, Any]] = None, enabled_precisions: Optional[Set[Union[dtype, dtype]]] = None, **kwargs: Any) → None[source]

在 Linux 中使用 TensorRT 編譯 PyTorch 模組，用於 Windows 中的推理

接受一個現有的 PyTorch 模組和一組配置編譯器的設定，並將方法轉換為 AOT 圖，這些圖會將等效的 TensorRT 序列化引擎資訊寫入使用者提供的指定 file_path 中。然後使用者可以在 Windows 中從磁碟載入反序列化的模型。注意：在 Linux 環境中為 Windows 交叉編譯的模型只能在 Windows 中載入。

引數

module (torch.nn.Module): 源模組 file_path (str): 用於將序列化模組儲存到磁碟的檔案路徑

關鍵字引數

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

    必需 模組輸入的形狀、dtype 和記憶體佈局規範列表。此引數為必需項。輸入大小可以指定為 torch 大小、元組或列表。dtype 可以使用 torch 資料型別或 torch_tensorrt 資料型別來指定，並且可以使用 torch 裝置或 torch_tensorrt 裝置型別列舉來選擇裝置型別。

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

  • arg_inputs (Tuple[Any, ...]) – 與 inputs 相同。為了更好地與 kwarg_inputs 對應，可作為別名。

  • kwarg_inputs (dict[Any, ...]) – 可選，模組 forward 函式的關鍵字引數輸入。

  • enabled_precision (Set(Union(torch.dpython:type, torch_tensorrt.dpython:type))) – TensorRT 在選擇 Kernel 時可以使用的基本資料型別集合

  • **kwargs – 特定請求策略的其他設定（更多資訊請參見子模組）

torch_tensorrt.load_cross_compiled_exported_program(file_path: str = '') → Any[source]

在 Windows 中載入之前在 Linux 中交叉編譯的 ExportedProgram 檔案

引數

file_path (str) – 磁碟上的檔案路徑

丟擲

ValueError – 如果在 Windows 中呼叫此 API，或檔案不存在，或檔案不是有效的 ExportedProgram 檔案

torch_tensorrt.save(module: Any, file_path: str = '', *, output_format: str = 'exported_program', inputs: Optional[Sequence[Tensor]] = None, arg_inputs: Optional[Sequence[Tensor]] = None, kwarg_inputs: Optional[dict[str, Any]] = None, retrace: bool = False) → None[source]

以指定的輸出格式將模型儲存到磁碟。

引數

  • module (Optional(torch.jit.ScriptModule | torch.export.ExportedProgram | torch.fx.GraphModule | CudaGraphsTorchTensorRTModule)) – 已編譯的 Torch-TensorRT 模組

  • inputs (torch.Tensor) – Torch 輸入張量

  • arg_inputs (Tuple[Any, ...]) – 與 inputs 相同。為了更好地與 kwarg_inputs 對應，可作為別名。

  • kwarg_inputs (dict[Any, ...]) – 可選，模組 forward 函式的關鍵字引數輸入。

  • output_format (str) – 儲存模型的格式。選項包括 exported_program | torchscript。

  • retrace (bool) – 當模組型別為 fx.GraphModule 時，此選項使用 torch.export.export(strict=False) 重新匯出圖並儲存。此標誌目前處於實驗階段。

torch_tensorrt.load(file_path: str = '') → Any[source]

載入 Torchscript 模型或 ExportedProgram。

從磁碟載入 TorchScript 或 ExportedProgram 檔案。將使用 try, except 檢測檔案型別。

引數

file_path (str) – 磁碟上的檔案路徑

丟擲

ValueError – 如果檔案不存在，或者檔案既不是 TorchScript 檔案也不是 ExportedProgram 檔案

類

class torch_tensorrt.MutableTorchTensorRTModule(pytorch_model: Module, *, device: Optional[Union[Device, device, str]] = None, disable_tf32: bool = False, assume_dynamic_shape_support: bool = False, sparse_weights: bool = False, enabled_precisions: Set[Union[dtype, dtype]] = {dtype.f32}, engine_capability: EngineCapability = EngineCapability.STANDARD, immutable_weights: bool = False, debug: bool = False, 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_double: bool = False, require_full_compilation: bool = False, min_block_size: int = 5, torch_executed_ops: Optional[Collection[Union[Callable[[...], Any], str]]] = None, torch_executed_modules: Optional[List[str]] = None, pass_through_build_failures: bool = False, max_aux_streams: Optional[int] = None, version_compatible: bool = False, optimization_level: Optional[int] = None, use_python_runtime: bool = False, use_fast_partitioner: bool = True, enable_experimental_decompositions: bool = False, dryrun: bool = False, hardware_compatible: bool = False, timing_cache_path: str = '/tmp/torch_tensorrt_engine_cache/timing_cache.bin', **kwargs: Any)[source]

初始化一個 MutableTorchTensorRTModule，以便像操作普通 PyTorch 模組一樣無縫地對其進行操作。在使用該模組時，所有的 TensorRT 編譯和重擬合過程都會自動處理。對其屬性的任何更改或載入不同的 state_dict 將觸發重擬合或重新編譯，這些操作將在下一次前向傳播期間進行管理。

MutableTorchTensorRTModule 接受一個 PyTorch 模組以及一組用於編譯器的配置設定。編譯完成後，該模組會保持 TensorRT 圖模組與原始 PyTorch 模組之間的連線。對 MutableTorchTensorRTModule 所做的任何修改都將反映在 TensorRT 圖模組和原始 PyTorch 模組中。

__init__(pytorch_model: Module, *, device: Optional[Union[Device, device, str]] = None, disable_tf32: bool = False, assume_dynamic_shape_support: bool = False, sparse_weights: bool = False, enabled_precisions: Set[Union[dtype, dtype]] = {dtype.f32}, engine_capability: EngineCapability = EngineCapability.STANDARD, immutable_weights: bool = False, debug: bool = False, 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_double: bool = False, require_full_compilation: bool = False, min_block_size: int = 5, torch_executed_ops: Optional[Collection[Union[Callable[[...], Any], str]]] = None, torch_executed_modules: Optional[List[str]] = None, pass_through_build_failures: bool = False, max_aux_streams: Optional[int] = None, version_compatible: bool = False, optimization_level: Optional[int] = None, use_python_runtime: bool = False, use_fast_partitioner: bool = True, enable_experimental_decompositions: bool = False, dryrun: bool = False, hardware_compatible: bool = False, timing_cache_path: str = '/tmp/torch_tensorrt_engine_cache/timing_cache.bin', **kwargs: Any) → None[source]
引數

pytorch_model (torch.nn.module) – 需要加速的源模組

關鍵字引數

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

    TensorRT 引擎執行的目標裝置

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

  • disable_tf32 (bool) – 強制 FP32 層使用傳統的 FP32 格式，而非預設行為（即在乘法前將輸入四捨五入到 10 位尾數，但使用 23 位尾數累加求和）

  • assume_dynamic_shape_support (bool) – 將此設定為 true 可使轉換器同時支援動態和靜態形狀。預設值：False

  • sparse_weights (bool) – 為卷積層和全連線層啟用稀疏性。

  • enabled_precision (Set(Union(torch.dpython:type, torch_tensorrt.dpython:type))) – TensorRT 在選擇 Kernel 時可以使用的基本資料型別集合

  • immutable_weights (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_double (bool) – 將雙精度 (float64) 提供的權重截斷為 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 (Collection[Target]) – 必須在 PyTorch 中執行的 aten 算子集合。如果此集合非空但 require_full_compilation 為 True，則會丟擲錯誤。

  • torch_executed_modules (List[str]) – 必須在 PyTorch 中執行的模組列表。如果此列表非空但 require_full_compilation 為 True，則會丟擲錯誤。

  • pass_through_build_failures (bool) – 如果編譯期間出現問題則報錯（僅適用於 torch.compile 工作流）

  • max_aux_stream (Optional[python:int]) – 引擎中的最大流數量

  • version_compatible (bool) – 構建與未來 TensorRT 版本相容的引擎（限制使用精簡執行時運算元以提供引擎的版本前向相容性）

  • optimization_level – (Optional[int]): 設定更高的最佳化級別允許 TensorRT 花費更長的引擎構建時間來尋找更多的最佳化選項。生成的引擎可能比使用較低最佳化級別構建的引擎效能更好。預設最佳化級別為 3。有效值包括從 0 到最大最佳化級別（當前為 5）的整數。設定為大於最大級別的值與最大級別行為相同。

  • use_python_runtime – (bool): 使用純 Python 執行時返回圖，減少序列化選項

  • use_fast_partitioner – (bool): 使用基於鄰接的劃分方案而非全域性劃分器。鄰接劃分速度更快，但可能不是最優的。如果追求最佳效能，請使用全域性劃分器 (False)。

  • enable_experimental_decompositions (bool) – 使用完整的運算元分解集合。這些分解可能未經充分測試，但有助於使圖更容易轉換為 TensorRT，從而可能增加在 TensorRT 中執行的圖的數量。

  • dryrun (bool) – “試執行”模式開關，除了轉換為 TRT 和記錄輸出外，執行所有內容

  • hardware_compatible (bool) – 構建與引擎構建所在的 GPU 架構之外的其他 GPU 架構相容的 TensorRT 引擎（目前適用於 NVIDIA Ampere 及更新架構）

  • timing_cache_path (str) – 時間快取的路徑（如果存在），或者編譯後儲存的路徑

  • lazy_engine_init (bool) – 推遲引擎設定直到所有引擎編譯完成。可以允許編譯具有多個圖中斷點的更大模型，但在執行時可能導致 GPU 記憶體超額訂閱。

  • **kwargs – 任意型別,

返回值

MutableTorchTensorRTModule

compile() → None[source]

(重新)編譯 TRT 圖模組，使用 PyTorch 模組。當權重結構發生變化（形狀、更多層等）時應呼叫此函式。MutableTorchTensorRTModule 會自動捕獲權重值更新並呼叫此函式進行重新編譯。如果它未能捕獲更改，請手動呼叫此函式重新編譯 TRT 圖模組。

refit_gm() → None[source]

使用任何更新重擬合 TRT 圖模組。當權重值發生變化但權重結構保持不變時應呼叫此函式。MutableTorchTensorRTModule 會自動捕獲權重值更新並呼叫此函式重擬合模組。如果它未能捕獲更改，請手動呼叫此函式更新 TRT 圖模組。

set_expected_dynamic_shape_range(args_dynamic_shape: tuple[dict[Any, Any]], kwargs_dynamic_shape: dict[str, Any]) → None[source]

設定動態形狀範圍。形狀提示應嚴格遵循傳遞給前向函式的 arg_inputs 和 kwarg_inputs，不應省略任何條目（kwarg_inputs 中的 None 除外）。如果輸入中有巢狀的字典/列表，則該條目的動態形狀也應是巢狀的字典/列表。如果輸入不需要動態形狀，則應為該輸入提供一個空字典作為形狀提示。請注意，您應排除值為 None 的關鍵字引數，因為它們將被過濾掉。

示例： def forward(a, b, c=0, d=0)

pass

seq_len = torch.export.Dim(“seq_len”, min=1, max=10) args_dynamic_shape = ({0: seq_len}, {}) # b 沒有動態形狀 kwargs_dynamic_shape = {‘c’: {0: seq_len}, ‘d’: {}} # d 沒有動態形狀 set_expected_dynamic_shape_range(args_dynamic_shape, kwargs_dynamic_shape) # 稍後呼叫函式 forward(*(a, b), **{c:…, d:…}) 時

參考： https://pytorch.com.tw/docs/stable/export.html#expressing-dynamism :param args_dynamic_shape: args_inputs 的動態形狀提示, :type args_dynamic_shape: tuple[dict[Any, Any]] :param kwargs_dynamic_shape: (dict[str, Any]): kwarg_inputs 的動態形狀提示

class torch_tensorrt.Input(*args: Any, **kwargs: Any)[source]

根據預期的形狀、資料型別和張量格式定義模組的輸入。

變數

  • shape_mode (torch_tensorrt.Input._ShapeMode) – 輸入是靜態形狀還是動態形狀

  • shape (Tuple or Dict) –

    定義輸入形狀的單個 Tuple 或 Dict。靜態形狀的輸入將有一個 Tuple。動態輸入將有一個如下形式的 Dict：

    {"min_shape": Tuple, "opt_shape": Tuple, "max_shape": Tuple}

  • dtype (torch_tensorrt.dpython:type) – 輸入張量的預期資料型別（預設值：torch_tensorrt.dtype.float32）

  • format (torch_tensorrt.TensorFormat) – 輸入張量的預期格式（預設值：torch_tensorrt.TensorFormat.NCHW）

__init__(*args: Any, **kwargs: Any) → None[source]

torch_tensorrt.Input 的 __init__ 方法

Input 接受以下幾種構造模式之一

引數

shape (Tuple or List, optional) – 輸入張量的靜態形狀

關鍵字引數

  • shape (Tuple or List, optional) – 輸入張量的靜態形狀

  • min_shape (Tuple or List, optional) – 輸入張量形狀範圍的最小值 注：必須同時提供 min_shape, opt_shape, max_shape，不能有位置引數，不能定義 shape，這意味著 Input 的 shape_mode 被設定為 DYNAMIC

  • opt_shape (Tuple or List, optional) – 輸入張量形狀範圍的最佳尺寸。注意：min_shape、opt_shape、max_shape 三者都必須提供，不能有位置引數，不能定義 shape，並且這隱式地將 Input 的 shape_mode 設定為 DYNAMIC

  • max_shape (Tuple or List, optional) – 輸入張量形狀範圍的最大尺寸。注意：min_shape、opt_shape、max_shape 三者都必須提供，不能有位置引數，不能定義 shape，並且這隱式地將 Input 的 shape_mode 設定為 DYNAMIC

  • dtype (torch.dpython:type or torch_tensorrt.dpython:type) – 輸入張量預期的資料型別（預設值：torch_tensorrt.dtype.float32）

  • format (torch.memory_format or torch_tensorrt.TensorFormat) – 輸入張量的預期格式（預設值：torch_tensorrt.TensorFormat.NCHW）

  • tensor_domain (Tuple(python:float, python:float), optional) – 張量允許的值的範圍，以區間表示法表示：[tensor_domain[0], tensor_domain[1])。注意：輸入“None”（或不指定）將把範圍設定為 [0, 2)

  • torch_tensor (torch.Tensor) – 持有與此 Input 對應的 torch 張量。

  • name (str, optional) – 此輸入在輸入 nn.Module 的 forward 函式中的名稱。用於在 dynamo tracer 中為相應輸入指定動態形狀。

示例

  • Input([1,3,32,32], dtype=torch.float32, format=torch.channel_last)

  • Input(shape=(1,3,32,32), dtype=torch_tensorrt.dtype.int32, format=torch_tensorrt.TensorFormat.NCHW)

  • Input(min_shape=(1,3,32,32), opt_shape=[2,3,32,32], max_shape=(3,3,32,32)) #隱式地 dtype=torch_tensorrt.dtype.float32, format=torch_tensorrt.TensorFormat.NCHW

example_tensor(optimization_profile_field: Optional[str] = None) → Tensor[source]

獲取 Input 物件指定的形狀的示例張量

引數

optimization_profile_field (Optional(str)) – 如果 Input 是動態形狀的，用於指定形狀的欄位名稱

返回值

一個 PyTorch 張量

classmethod from_tensor(t: Tensor, disable_memory_format_check: bool = False) → Input[source]

生成一個 Input，其中包含給定 PyTorch 張量的資訊。

引數

  • tensor (torch.Tensor) – 一個 PyTorch 張量。

  • disable_memory_format_check (bool) – 是否驗證輸入張量的記憶體格式

返回值

一個 Input 物件。

classmethod from_tensors(ts: Sequence[Tensor], disable_memory_format_check: bool = False) → List[Input][source]

生成一個 Input 列表，其中包含所有給定 PyTorch 張量的資訊。

引數

  • tensors (Iterable[torch.Tensor]) – 一個 PyTorch 張量列表。

  • disable_memory_format_check (bool) – 是否驗證輸入張量的記憶體格式

返回值

一個 Input 列表。

dtype: dtype = 1

torch_tensorrt.dtype.float32)

型別

輸入張量的預期資料型別（預設值

format: memory_format = 1

torch_tensorrt.memory_format.linear)

型別

輸入張量的預期格式（預設值

class torch_tensorrt.Device(*args: Any, **kwargs: Any)[source]

定義可用於指定引擎目標裝置的裝置

變數

  • device_type (DeviceType) – 目標裝置型別（GPU 或 DLA）。基於是否指定了 dla_core 隱式設定。

  • gpu_id (python:int) – 目標 GPU 的裝置 ID

  • dla_core (python:int) – 目標 DLA 核心的核心 ID

  • allow_gpu_fallback (bool) – 如果 DLA 不支援某個操作，是否允許回退到 GPU

__init__(*args: Any, **kwargs: Any)[source]

torch_tensorrt.Device 的 __init__ 方法

Device 接受以下幾種構造模式之一

引數

spec (str) – 帶有裝置規格的字串，例如“dla:0”表示 DLA，核心 ID 0

關鍵字引數

  • gpu_id (python:int) – 目標 GPU 的 ID（如果指定了 dla_core，將由管理 DLA 的 GPU 覆蓋）。如果指定，不應提供位置引數

  • dla_core (python:int) – 目標 DLA 核心的 ID。如果指定，不應提供位置引數。

  • allow_gpu_fallback (bool) – 如果操作在 DLA 上不受支援，則允許 TensorRT 將其排程到 GPU 上（如果裝置型別不是 DLA，則忽略）

示例

  • Device(“gpu:1”)

  • Device(“cuda:1”)

  • Device(“dla:0”, allow_gpu_fallback=True)

  • Device(gpu_id=0, dla_core=0, allow_gpu_fallback=True)

  • Device(dla_core=0, allow_gpu_fallback=True)

  • Device(gpu_id=1)

device_type: DeviceType = 1

目標裝置型別（GPU 或 DLA）。基於是否指定了 dla_core 隱式設定。

dla_core: int = -1

目標 DLA 核心的核心 ID

gpu_id: int = -1

目標 GPU 的裝置 ID

列舉

class torch_tensorrt.dtype(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)[source]

描述 Torch-TensorRT 資料型別的列舉，與 torch、tensorrt 和 numpy 的 dtypes 相容

to(t: Union[Type[dtype], Type[DataType], Type[dtype], Type[dtype]], use_default: bool = False) → Union[dtype, DataType, dtype, dtype][source]

將 dtype 轉換為 [torch, numpy, tensorrt] 中的等效型別

將 self 轉換為 numpy、torch 和 tensorrt 中的等效 dtypes。如果 self 在目標庫中不受支援，則將引發異常。因此，不建議直接使用此方法。

或者使用 torch_tensorrt.dtype.try_to()

引數

  • t (Union(Type(torch.dpython:type), Type(tensorrt.DataType), Type(numpy.dpython:type), Type(dpython:type))) – 要轉換到的另一個庫的資料型別列舉

  • use_default (bool) – 在某些情況下，一個通用的型別（例如 torch.float）就足夠了，因此 instead of throwing an exception, return default value。

返回值

庫列舉 t 中等效的 torch_tensorrt.dtype dtype

返回值型別

Union(torch.dtype, tensorrt.DataType, numpy.dtype, dtype)

丟擲

TypeError – 不支援的資料型別或未知目標

示例

# Succeeds
float_dtype = torch_tensorrt.dtype.f32.to(torch.dtype) # Returns torch.float

# Failure
float_dtype = torch_tensorrt.dtype.bf16.to(numpy.dtype) # Throws exception
classmethod try_from(t: Union[dtype, DataType, dtype, dtype], use_default: bool = False) → Optional[dtype][source]

從另一個庫的 dtype 系統建立 Torch-TensorRT dtype。

接受來自 numpy、torch 和 tensorrt 之一的 dtype 列舉，並建立一個 torch_tensorrt.dtype。如果源 dtype 系統不受支援或該型別在 Torch-TensorRT 中不受支援，則返回 None。

引數

  • t (Union(torch.dpython:type, tensorrt.DataType, numpy.dpython:type, dpython:type)) – 來自另一個庫的資料型別列舉

  • use_default (bool) – 在某些情況下，一個通用的型別（例如 torch_tensorrt.dtype.f32）就足夠了，因此 вместо выбрасывания исключения，возвращается значение по умолчанию。

返回值

與 t 等效的 torch_tensorrt.dtype 或 None

返回值型別

Optional(dtype)

示例

# Succeeds
float_dtype = torch_tensorrt.dtype.try_from(torch.float) # Returns torch_tensorrt.dtype.f32

# Unsupported type
float_dtype = torch_tensorrt.dtype.try_from(torch.complex128) # Returns None
try_to(t: Union[Type[dtype], Type[DataType], Type[dtype], Type[dtype]], use_default: bool) → Optional[Union[dtype, DataType, dtype, dtype]][source]

將 dtype 轉換為 [torch, numpy, tensorrt] 中的等效型別

將 self 轉換為 numpy、torch 和 tensorrt 中的等效 dtypes。如果 self 在目標庫中不受支援，則返回 None。

引數

  • t (Union(Type(torch.dpython:type), Type(tensorrt.DataType), Type(numpy.dpython:type), Type(dpython:type))) – 要轉換到的另一個庫的資料型別列舉

  • use_default (bool) – 在某些情況下，一個通用的型別（例如 torch.float）就足夠了，因此 instead of throwing an exception, return default value。

返回值

庫列舉 t 中等效的 torch_tensorrt.dtype dtype

返回值型別

Optional(Union(torch.dtype, tensorrt.DataType, numpy.dtype, dtype))

示例

# Succeeds
float_dtype = torch_tensorrt.dtype.f32.to(torch.dtype) # Returns torch.float

# Failure
float_dtype = torch_tensorrt.dtype.bf16.to(numpy.dtype) # Returns None
b

布林值，等同於 dtype.bool

bf16

16 位“Brain”浮點數，等同於 dtype.bfloat16

f16

16 位浮點數，等同於 dtype.half、dtype.fp16 和 dtype.float16

f32

32 位浮點數，等同於 dtype.float、dtype.fp32 和 dtype.float32

f64

64 位浮點數，等同於 dtype.double、dtype.fp64 和 dtype.float64

f8

8 位浮點數，等同於 dtype.fp8 和 dtype.float8

i32

有符號 32 位整數，等同於 dtype.int32 和 dtype.int

i64

有符號 64 位整數，等同於 dtype.int64 和 dtype.long

i8

有符號 8 位整數，等同於 dtype.int8，當啟用為核心精度時通常需要模型支援量化

u8

無符號 8 位整數，等同於 dtype.uint8

unknown

標記值

class torch_tensorrt.DeviceType(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)[source]

TensorRT 將面向的裝置型別

to(t: Union[Type[DeviceType], Type[DeviceType]], use_default: bool = False) → Union[DeviceType, DeviceType][source]

將 DeviceType 轉換為 tensorrt 中的等效型別

將 self 轉換為 torch 或 tensorrt 中的等效裝置型別。如果 self 在目標庫中不受支援，則將引發異常。因此，不建議直接使用此方法。

或者使用 torch_tensorrt.DeviceType.try_to()

引數

t (Union(Type(tensorrt.DeviceType), Type(DeviceType))) – 要轉換到的另一個庫的裝置型別列舉

返回值

列舉 t 中等效於 torch_tensorrt.DeviceType 的裝置型別

返回值型別

Union(tensorrt.DeviceType, DeviceType)

丟擲

TypeError – 未知目標型別或不支援的裝置型別

示例

# Succeeds
trt_dla = torch_tensorrt.DeviceType.DLA.to(tensorrt.DeviceType) # Returns tensorrt.DeviceType.DLA
classmethod try_from(d: Union[DeviceType, DeviceType]) → Optional[DeviceType][source]

從 TensorRT 裝置型別列舉建立 Torch-TensorRT 裝置型別列舉。

接受來自 tensorrt 的裝置型別列舉，並建立一個 torch_tensorrt.DeviceType。如果源不受支援或該裝置型別在 Torch-TensorRT 中不受支援，則將引發異常。因此，不建議直接使用此方法。

或者使用 torch_tensorrt.DeviceType.try_from()

引數

d (Union(tensorrt.DeviceType, DeviceType)) – 來自另一個庫的裝置型別列舉

返回值

與 d 等效的 torch_tensorrt.DeviceType

返回值型別

DeviceType

示例

torchtrt_dla = torch_tensorrt.DeviceType._from(tensorrt.DeviceType.DLA)
try_to(t: Union[Type[DeviceType], Type[DeviceType]], use_default: bool = False) → Optional[Union[DeviceType, DeviceType]][source]

將 DeviceType 轉換為 tensorrt 中的等效型別

將 self 轉換為 torch 或 tensorrt 中的等效記憶體格式。如果 self 在目標庫中不受支援，則返回 None。

引數

t (Union(Type(tensorrt.DeviceType), Type(DeviceType))) – 要轉換到的另一個庫的裝置型別列舉

返回值

列舉 t 中等效於 torch_tensorrt.DeviceType 的裝置型別

返回值型別

Optional(Union(tensorrt.DeviceType, DeviceType))

示例

# Succeeds
trt_dla = torch_tensorrt.DeviceType.DLA.to(tensorrt.DeviceType) # Returns tensorrt.DeviceType.DLA
DLA

目標是 DLA 核心

GPU

目標是 GPU

UNKNOWN

標記值

class torch_tensorrt.EngineCapability(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)[source]

EngineCapability 決定了構建時網路的限制以及它面向的執行時。

to(t: Union[Type[EngineCapability], Type[EngineCapability]]) → Union[EngineCapability, EngineCapability][source]

將 EngineCapability 轉換為 tensorrt 中的等效型別

將 self 轉換為 torch 或 tensorrt 中的等效引擎能力。如果 self 在目標庫中不受支援，則將引發異常。因此，不建議直接使用此方法。

或者使用 torch_tensorrt.EngineCapability.try_to()

引數

t (Union(Type(tensorrt.EngineCapability), Type(EngineCapability))) – 要轉換到的另一個庫的引擎能力列舉

返回值

列舉 t 中等效於 torch_tensorrt.EngineCapability 的引擎能力

返回值型別

Union(tensorrt.EngineCapability, EngineCapability)

丟擲

TypeError – 未知目標型別或不支援的引擎能力

示例

# Succeeds
torchtrt_dla_ec = torch_tensorrt.EngineCapability.DLA_STANDALONE.to(tensorrt.EngineCapability) # Returns tensorrt.EngineCapability.DLA
classmethod try_from() → Optional[EngineCapability][source]

從 TensorRT 引擎能力列舉建立 Torch-TensorRT 引擎能力列舉。

獲取 TensorRT 的裝置型別列舉，並建立一個 torch_tensorrt.EngineCapability。如果源不支援，或者 Torch-TensorRT 不支援該引擎能力級別，則會引發異常。因此，不建議直接使用此方法。

或者使用 torch_tensorrt.EngineCapability.try_from()

引數

c (Union(tensorrt.EngineCapability, EngineCapability)) – 來自其他庫的引擎能力列舉

返回值

與 c 等效的 torch_tensorrt.EngineCapability

返回值型別

EngineCapability

示例

torchtrt_safety_ec = torch_tensorrt.EngineCapability._from(tensorrt.EngineCapability.SAEFTY)
try_to(t: Union[Type[EngineCapability], Type[EngineCapability]]) → Optional[Union[EngineCapability, EngineCapability]][source]

將 EngineCapability 轉換為 tensorrt 中的等效型別

將 self 轉換為等效的 torch 或 tensorrt 引擎能力之一。如果目標庫不支援 self，則返回 None。

引數

t (Union(Type(tensorrt.EngineCapability), Type(EngineCapability))) – 要轉換到的另一個庫的引擎能力列舉

返回值

列舉 t 中等效於 torch_tensorrt.EngineCapability 的引擎能力

返回值型別

Optional(Union(tensorrt.EngineCapability, EngineCapability))

示例

# Succeeds
trt_dla_ec = torch_tensorrt.EngineCapability.DLA.to(tensorrt.EngineCapability) # Returns tensorrt.EngineCapability.DLA_STANDALONE
DLA_STANDALONE

EngineCapability.DLA_STANDALONE 提供 DLA 相容的網路操作受限子集，生成的序列化引擎可以使用獨立 DLA 執行時 API 執行。

SAFETY

EngineCapability.SAFETY 提供安全認證的網路操作受限子集，生成的序列化引擎可以使用 TensorRT 的 tensorrt.safe 名稱空間中的安全執行時 API 執行。

STANDARD

EngineCapability.STANDARD 對功能沒有任何限制，生成的序列化引擎可以使用 TensorRT 的標準執行時 API 執行。

class torch_tensorrt.memory_format(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)[source]
to(t: Union[Type[memory_format], Type[TensorFormat], Type[memory_format]]) → Union[memory_format, TensorFormat, memory_format][source]

將 memory_format 轉換為 torch 或 tensorrt 中等效的型別

將 self 轉換為等效的 torch 或 tensorrt 記憶體格式之一。如果目標庫不支援 self，則會引發異常。因此，不建議直接使用此方法。

或者使用 torch_tensorrt.memory_format.try_to()

引數

t (Union(Type(torch.memory_format), Type(tensorrt.TensorFormat), Type(memory_format))) – 要轉換到的來自另一個庫的記憶體格式型別列舉

返回值

列舉 t 中等效的 torch_tensorrt.memory_format 記憶體格式

返回值型別

Union(torch.memory_format, tensorrt.TensorFormat, memory_format)

丟擲

TypeError – 未知目標型別或不支援的記憶體格式

示例

# Succeeds
tf = torch_tensorrt.memory_format.linear.to(torch.dtype) # Returns torch.contiguous
classmethod try_from(f: Union[memory_format, TensorFormat, memory_format]) → Optional[memory_format][source]

從另一個庫的記憶體格式列舉建立 Torch-TensorRT 記憶體格式列舉。

獲取來自 torch 和 tensorrt 之一的記憶體格式列舉，並建立一個 torch_tensorrt.memory_format。如果源不支援，或者 Torch-TensorRT 不支援該記憶體格式，則返回 None。

引數

f (Union(torch.memory_format, tensorrt.TensorFormat, memory_format)) – 來自另一個庫的記憶體格式列舉

返回值

與 f 等效的 torch_tensorrt.memory_format

返回值型別

Optional(memory_format)

示例

torchtrt_linear = torch_tensorrt.memory_format.try_from(torch.contiguous)
try_to(t: Union[Type[memory_format], Type[TensorFormat], Type[memory_format]]) → Optional[Union[memory_format, TensorFormat, memory_format]][source]

將 memory_format 轉換為 torch 或 tensorrt 中等效的型別

將 self 轉換為等效的 torch 或 tensorrt 記憶體格式之一。如果目標庫不支援 self，則返回 None

引數

t (Union(Type(torch.memory_format), Type(tensorrt.TensorFormat), Type(memory_format))) – 要轉換到的來自另一個庫的記憶體格式型別列舉

返回值

列舉 t 中等效的 torch_tensorrt.memory_format 記憶體格式

返回值型別

Optional(Union(torch.memory_format, tensorrt.TensorFormat, memory_format))

示例

# Succeeds
tf = torch_tensorrt.memory_format.linear.to(torch.dtype) # Returns torch.contiguous
cdhw32

具有 3 個空間維度的 32 通道向量化行主格式。

此格式繫結到 FP16 和 INT8。僅適用於維度 >= 4 的張量。

對於維度為 {N, C, D, H, W} 的張量，記憶體佈局等同於維度為 [N][(C+31)/32][D][H][W][32] 的 C 陣列，其中張量座標 (n, d, c, h, w) 對映到陣列下標 [n][c/32][d][h][w][c%32]。

chw16

16 通道向量化行主格式。

此格式繫結到 FP16。僅適用於維度 >= 3 的張量。

對於維度為 {N, C, H, W} 的張量，記憶體佈局等同於維度為 [N][(C+15)/16][H][W][16] 的 C 陣列，其中張量座標 (n, c, h, w) 對映到陣列下標 [n][c/16][h][w][c%16]。

chw2

2 通道向量化行主格式。

在 TensorRT 中，此格式繫結到 FP16。僅適用於維度 >= 3 的張量。

對於維度為 {N, C, H, W} 的張量，記憶體佈局等同於維度為 [N][(C+1)/2][H][W][2] 的 C 陣列，其中張量座標 (n, h, c, w) 對映到陣列下標 [n][c/2][h][w][c%2]。

chw32

32 通道向量化行主格式。

此格式僅適用於維度 >= 3 的張量。

對於維度為 {N, C, H, W} 的張量，記憶體佈局等同於維度為 [N][(C+31)/32][H][W][32] 的 C 陣列，其中張量座標 (n, c, h, w) 對映到陣列下標 [n][c/32][h][w][c%32]。

chw4

4 通道向量化行主格式。此格式繫結到 INT8。僅適用於維度 >= 3 的張量。

對於維度為 {N, C, H, W} 的張量，記憶體佈局等同於維度為 [N][(C+3)/4][H][W][4] 的 C 陣列，其中張量座標 (n, c, h, w) 對映到陣列下標 [n][c/4][h][w][c%4]。

dhwc

非向量化通道最後格式。此格式繫結到 FP32。僅適用於維度 >= 4 的張量。

等效於 memory_format.channels_last_3d

dhwc8

8 通道格式，其中 C 填充到 8 的倍數。

此格式繫結到 FP16，僅適用於維度 >= 4 的張量。

對於維度為 {N, C, D, H, W} 的張量，記憶體佈局等同於維度為 [N][D][H][W][(C+7)/8*8] 的陣列，其中張量座標 (n, c, d, h, w) 對映到陣列下標 [n][d][h][w][c]。

dla_hwc4

DLA 影像格式。通道最後格式。C 只能是 1、3、4。如果 C == 3，則會四捨五入到 4。沿 H 軸步進的步長向上取整到 32 位元組。

此格式繫結到 FP16/Int8，僅適用於維度 >= 3 的張量。

對於維度為 {N, C, H, W} 的張量，其中當 C 分別為 1、3、4 時，C’ 為 1、4、4，記憶體佈局等同於維度為 [N][H][roundUp(W, 32/C’/elementSize)][C’] 的 C 陣列，其中 elementSize 對於 FP16 為 2，對於 Int8 為 1，C’ 是取整後的 C。張量座標 (n, c, h, w) 對映到陣列下標 [n][h][w][c]。

dla_linear

DLA 平面格式。行主格式。沿 H 軸步進的步長向上取整到 64 位元組。

此格式繫結到 FP16/Int8，僅適用於維度 >= 3 的張量。

對於維度為 {N, C, H, W} 的張量，記憶體佈局等同於維度為 [N][C][H][roundUp(W, 64/elementSize)] 的 C 陣列，其中 elementSize 對於 FP16 為 2，對於 Int8 為 1，張量座標 (n, c, h, w) 對映到陣列下標 [n][c][h][w]。

hwc

非向量化通道最後格式。此格式繫結到 FP32，僅適用於維度 >= 3 的張量。

等效於 memory_format.channels_last

hwc16

16 通道格式，其中 C 填充到 16 的倍數。此格式繫結到 FP16。僅適用於維度 >= 3 的張量。

對於維度為 {N, C, H, W} 的張量，記憶體佈局等同於維度為 [N][H][W][(C+15)/16*16] 的陣列，其中張量座標 (n, c, h, w) 對映到陣列下標 [n][h][w][c]。

hwc8

8 通道格式，其中 C 填充到 8 的倍數。

此格式繫結到 FP16。僅適用於維度 >= 3 的張量。

對於維度為 {N, C, H, W} 的張量，記憶體佈局等同於維度為 [N][H][W][(C+7)/8*8] 的陣列，其中張量座標 (n, c, h, w) 對映到陣列下標 [n][h][w][c]。

linear

行主線性格式。

對於維度為 {N, C, H, W} 的張量，W 軸始終具有單位步長，而其他每個軸的步長至少是下一個維度乘以下一個步長的乘積。步長與維度為 [N][C][H][W] 的 C 陣列相同。

等效於 memory_format.contiguous

子模組

文件

訪問 PyTorch 的綜合開發者文件

檢視文件

教程

獲取面向初學者和高階開發者的深度教程

檢視教程

資源

查詢開發資源並獲取問題解答

檢視資源

為了分析流量並最佳化您的體驗，我們在本網站上提供 Cookie。透過點選或導航，您同意我們使用 Cookie。作為本網站的當前維護者，Facebook 的 Cookie 政策適用。瞭解更多，包括有關可用控制元件的資訊：Cookie 政策。