評價此頁
★ ★ ★ ★ ★

torch.backends#

創建於: 2020年9月16日 | 最後更新於: 2025年8月26日

torch.backends 控制 PyTorch 支援的各種後端行為。

這些後端包括

  • torch.backends.cpu

  • torch.backends.cuda

  • torch.backends.cudnn

  • torch.backends.cusparselt

  • torch.backends.mha

  • torch.backends.mps

  • torch.backends.mkl

  • torch.backends.mkldnn

  • torch.backends.nnpack

  • torch.backends.openmp

  • torch.backends.opt_einsum

  • torch.backends.xeon

torch.backends.cpu#

torch.backends.cpu.get_cpu_capability()[source]#

返回 CPU 功能作為字串值。

可能的值： - “DEFAULT” - “VSX” - “Z VECTOR” - “NO AVX” - “AVX2” - “AVX512” - “SVE256”

返回型別

str

torch.backends.cuda#

torch.backends.cuda.is_built()[source]#

返回 PyTorch 是否已構建支援 CUDA。

請注意，這並不一定意味著 CUDA 可用；它僅僅表示如果此 PyTorch 二進位制檔案在具有工作 CUDA 驅動程式和裝置的機器上執行，我們將能夠使用它。

torch.backends.cuda.matmul.allow_tf32#

一個 bool，用於控制在 Ampere 或更新的 GPU 上進行矩陣乘法時是否可以使用 TensorFloat-32 張量核心。allow_tf32 將被棄用。請參閱 Ampere（及更高版本）裝置上的 TensorFloat-32 (TF32)。

torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction#

一個 bool，用於控制在使用 fp16 GEMM 時是否允許使用低精度歸約（例如，使用 fp16 累加型別）。

torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction#

一個 bool，用於控制在使用 bf16 GEMM 時是否允許使用低精度歸約。

torch.backends.cuda.cufft_plan_cache#

cufft_plan_cache 包含每個 CUDA 裝置的 cuFFT 計劃快取。透過 torch.backends.cuda.cufft_plan_cache[i] 查詢特定裝置 i 的快取。

torch.backends.cuda.cufft_plan_cache.size#

一個只讀 int，顯示當前 cuFFT 計劃快取中的計劃數量。

torch.backends.cuda.cufft_plan_cache.max_size#

一個 int，用於控制 cuFFT 計劃快取的容量。

torch.backends.cuda.cufft_plan_cache.clear()#

清除 cuFFT 計劃快取。

torch.backends.cuda.preferred_blas_library(backend=None)[source]#

覆蓋 PyTorch 用於 BLAS 操作的庫。在 cuBLAS、cuBLASLt 和 CK [僅限 ROCm] 之間進行選擇。

警告

此標誌是實驗性的，可能會發生變化。

當 PyTorch 執行 CUDA BLAS 操作時，即使 cuBLAS 和 cuBLASLt 都可用，它也預設使用 cuBLAS。對於為 ROCm 構建的 PyTorch，hipBLAS、hipBLASLt 和 CK 可能會提供不同的效能。此標誌（一個 str）允許覆蓋要使用的 BLAS 庫。

  • 如果設定為 “cublas”，則 wherever possible 將使用 cuBLAS。

  • 如果設定為 “cublaslt”，則 wherever possible 將使用 cuBLASLt。

  • 如果設定為 “ck”，則 wherever possible 將使用 CK。

  • 如果設定為 “default”（預設值），則將使用啟發式方法在其他選項之間進行選擇。

  • 不提供輸入時，此函式返回當前首選的庫。

  • 使用者可以使用環境變數 TORCH_BLAS_PREFER_CUBLASLT=1 全域性設定首選庫為 cuBLASLt。此標誌僅設定首選庫的初始值，並且首選庫仍可能在此函式呼叫後被指令碼中的後續呼叫覆蓋。

注意：當某個庫被首選時，如果首選庫不實現所呼叫的操作，仍可能使用其他庫。如果 PyTorch 的庫選擇不適合您的應用程式輸入，此標誌可能會獲得更好的效能。

返回型別

_BlasBackend

torch.backends.cuda.preferred_rocm_fa_library(backend=None)[source]#

[僅限 ROCm] 覆蓋 PyTorch 在 ROCm 環境中用於 Flash Attention 的後端。在 AOTriton 和 CK 之間進行選擇

警告

此標誌是實驗性的，可能會發生變化。

當啟用了 Flash Attention 並需要時，PyTorch 預設使用 AOTriton 作為後端。此標誌（一個 str）允許使用者覆蓋此後端以使用 composable_kernel

  • 如果設定為 “default”，則 wherever possible 將使用預設後端。目前是 AOTriton。

  • 如果設定為 “aotriton”，則 wherever possible 將使用 AOTriton。

  • 如果設定為 “ck”，則 wherever possible 將使用 CK。

  • 不提供輸入時，此函式返回當前首選的庫。

  • 使用者可以使用環境變數 TORCH_ROCM_FA_PREFER_CK=1 將首選庫設定為 CK 全域性。

注意：當某個庫被首選時，如果首選庫不實現所呼叫的操作，仍可能使用其他庫。如果 PyTorch 的庫選擇不適合您的應用程式輸入，此標誌可能會獲得更好的效能。

返回型別

_ROCmFABackend

torch.backends.cuda.preferred_linalg_library(backend=None)[source]#

覆蓋 PyTorch 在 CUDA 線性代數運算中選擇 cuSOLVER 和 MAGMA 的啟發式方法。

警告

此標誌是實驗性的，可能會發生變化。

當 PyTorch 執行 CUDA 線性代數運算時，它經常使用 cuSOLVER 或 MAGMA 庫，如果兩者都可用，它會透過啟發式方法決定使用哪一個。此標誌（一個 str）允許覆蓋這些啟發式方法。

  • 如果設定為 “cusolver”，則 wherever possible 將使用 cuSOLVER。

  • 如果設定為 “magma”，則 wherever possible 將使用 MAGMA。

  • 如果設定為 “default”（預設值），則當兩者都可用時，將使用啟發式方法在 cuSOLVER 和 MAGMA 之間進行選擇。

  • 不提供輸入時，此函式返回當前首選的庫。

  • 使用者可以使用環境變數 TORCH_LINALG_PREFER_CUSOLVER=1 全域性設定首選庫為 cuSOLVER。此標誌僅設定首選庫的初始值，並且首選庫仍可能在此函式呼叫後被指令碼中的後續呼叫覆蓋。

注意：當某個庫被首選時，如果首選庫不實現所呼叫的操作，仍可能使用其他庫。此標誌可能會在 PyTorch 的啟發式庫選擇不適合您的應用程式輸入時獲得更好的效能。

當前支援的線性代數運算子

返回型別

_LinalgBackend

class torch.backends.cuda.SDPAParams#
torch.backends.cuda.flash_sdp_enabled()[source]#

警告

此標誌為 beta 版，可能會發生變化。

返回 flash scaled dot product attention 是否啟用。

torch.backends.cuda.enable_mem_efficient_sdp(enabled)[source]#

警告

此標誌為 beta 版，可能會發生變化。

啟用或停用記憶體高效的 scaled dot product attention。

torch.backends.cuda.mem_efficient_sdp_enabled()[source]#

警告

此標誌為 beta 版，可能會發生變化。

返回 memory efficient scaled dot product attention 是否啟用。

torch.backends.cuda.enable_flash_sdp(enabled)[source]#

警告

此標誌為 beta 版，可能會發生變化。

啟用或停用 flash scaled dot product attention。

torch.backends.cuda.math_sdp_enabled()[source]#

警告

此標誌為 beta 版，可能會發生變化。

返回 math scaled dot product attention 是否啟用。

torch.backends.cuda.enable_math_sdp(enabled)[source]#

警告

此標誌為 beta 版，可能會發生變化。

啟用或停用 math scaled dot product attention。

torch.backends.cuda.fp16_bf16_reduction_math_sdp_allowed()[source]#

警告

此標誌為 beta 版，可能會發生變化。

返回 math scaled dot product attention 中的 fp16/bf16 歸約是否啟用。

torch.backends.cuda.allow_fp16_bf16_reduction_math_sdp(enabled)[source]#

警告

此標誌為 beta 版，可能會發生變化。

啟用或停用 math scaled dot product attention 中的 fp16/bf16 歸約。

torch.backends.cuda.cudnn_sdp_enabled()[source]#

警告

此標誌為 beta 版，可能會發生變化。

返回 cuDNN scaled dot product attention 是否啟用。

torch.backends.cuda.enable_cudnn_sdp(enabled)[source]#

警告

此標誌為 beta 版，可能會發生變化。

啟用或停用 cuDNN scaled dot product attention。

torch.backends.cuda.is_flash_attention_available()[source]#

檢查 PyTorch 是否已構建用於 scaled_dot_product_attention 的 FlashAttention。

返回

如果 FlashAttention 已構建並可用，則返回 True；否則返回 False。

返回型別

布林值

注意

此函式依賴於啟用 CUDA 的 PyTorch 構建。在非 CUDA 環境中將返回 False。

torch.backends.cuda.can_use_flash_attention(params, debug=False)[source]#

檢查是否可以在 scaled_dot_product_attention 中使用 FlashAttention。

引數

  • params (_SDPAParams) – 一個 SDPAParams 例項，包含 query、key、value 張量、可選的 attention mask、dropout 率以及一個指示 attention 是否為 causal 的標誌。

  • debug (bool) – 是否記錄 logging.warn 關於為什麼無法執行 FlashAttention 的除錯資訊。預設為 False。

返回

如果可以使用 FlashAttention 並滿足給定引數，則返回 True；否則返回 False。

返回型別

布林值

注意

此函式依賴於啟用 CUDA 的 PyTorch 構建。在非 CUDA 環境中將返回 False。

torch.backends.cuda.can_use_efficient_attention(params, debug=False)[source]#

檢查是否可以在 scaled_dot_product_attention 中使用 efficient_attention。

引數

  • params (_SDPAParams) – 一個 SDPAParams 例項，包含 query、key、value 張量、可選的 attention mask、dropout 率以及一個指示 attention 是否為 causal 的標誌。

  • debug (bool) – 是否記錄 logging.warn 關於為什麼無法執行 efficient_attention 的資訊。預設為 False。

返回

如果可以使用 efficient_attention 並滿足給定引數，則返回 True；否則返回 False。

返回型別

布林值

注意

此函式依賴於啟用 CUDA 的 PyTorch 構建。在非 CUDA 環境中將返回 False。

torch.backends.cuda.can_use_cudnn_attention(params, debug=False)[source]#

檢查是否可以在 scaled_dot_product_attention 中使用 cudnn_attention。

引數

  • params (_SDPAParams) – 一個 SDPAParams 例項，包含 query、key、value 張量、可選的 attention mask、dropout 率以及一個指示 attention 是否為 causal 的標誌。

  • debug (bool) – 是否記錄 logging.warn 關於為什麼無法執行 cuDNN attention 的資訊。預設為 False。

返回

如果可以使用 cuDNN 並滿足給定引數，則返回 True；否則返回 False。

返回型別

布林值

注意

此函式依賴於啟用 CUDA 的 PyTorch 構建。在非 CUDA 環境中將返回 False。

torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=True, enable_cudnn=True)[source]#

警告

此標誌為 beta 版，可能會發生變化。

此上下文管理器可用於臨時啟用或停用 scaled dot product attention 的任何三個後端。退出上下文管理器後，將恢復標誌的先前狀態。

torch.backends.cudnn#

torch.backends.cudnn.version()[source]#

返回 cuDNN 的版本。

torch.backends.cudnn.is_available()[source]#

返回一個布林值，指示 CUDNN 當前是否可用。

torch.backends.cudnn.enabled#

一個 bool，用於控制是否啟用 cuDNN。

torch.backends.cudnn.allow_tf32#

一個 bool，用於控制在 Ampere 或更新的 GPU 上是否可以在 cuDNN 卷積中使用 TensorFloat-32 張量核心。allow_tf32 將被棄用。請參閱 Ampere（及更高版本）裝置上的 TensorFloat-32 (TF32)。

torch.backends.cudnn.deterministic#

一個 bool，如果為 True，則會導致 cuDNN 只使用確定性卷積演算法。另請參閱 torch.are_deterministic_algorithms_enabled() 和 torch.use_deterministic_algorithms()。

torch.backends.cudnn.benchmark#

一個 bool，如果為 True，則會導致 cuDNN 對多個卷積演算法進行基準測試並選擇最快的。

torch.backends.cudnn.benchmark_limit#

一個 int，指定當 torch.backends.cudnn.benchmark 為 True 時嘗試的 cuDNN 卷積演算法的最大數量。將 benchmark_limit 設定為零以嘗試所有可用演算法。請注意，此設定僅影響透過 cuDNN v8 API 分派的卷積。

torch.backends.cusparselt#

torch.backends.cusparselt.version()[source]#

返回 cuSPARSELt 的版本

返回型別

Optional[int]

torch.backends.cusparselt.is_available()[source]#

返回一個布林值，指示 cuSPARSELt 當前是否可用。

返回型別

布林值

torch.backends.mha#

torch.backends.mha.get_fastpath_enabled()[source]#

返回 TransformerEncoder 和 MultiHeadAttention 的快速路徑是否啟用，或者如果 jit 正在指令碼化，則返回 True。

注意

除非滿足輸入的所有條件，否則即使 get_fastpath_enabled 返回 True，快速路徑也可能不會執行。

返回型別

布林值

torch.backends.mha.set_fastpath_enabled(value)[source]#

設定是否啟用快速路徑

torch.backends.miopen#

torch.backends.miopen.immediate#

一個 bool，如果為 True，則會導致 MIOpen 使用 Immediate Mode（https://rocm.docs.amd.com/projects/MIOpen/en/latest/how-to/find-and-immediate.html）。

torch.backends.mps#

torch.backends.mps.is_available()[source]#

返回一個布林值，指示 MPS 當前是否可用。

返回型別

布林值

torch.backends.mps.is_built()[source]#

返回 PyTorch 是否已構建支援 MPS。

請注意，這並不一定意味著 MPS 可用；它僅僅表示如果此 PyTorch 二進位制檔案在具有工作 MPS 驅動程式和裝置的機器上執行，我們將能夠使用它。

返回型別

布林值

torch.backends.mkl#

torch.backends.mkl.is_available()[source]#

返回 PyTorch 是否已構建支援 MKL。

class torch.backends.mkl.verbose(enable)[source]#

按需 oneMKL 詳細輸出功能。

為了更容易地除錯效能問題，oneMKL 可以在執行核心時轉儲包含執行資訊的詳細訊息，例如持續時間。詳細輸出功能可以透過名為 MKL_VERBOSE 的環境變數呼叫。但是，這種方法會在所有步驟中轉儲訊息。這些是大量的詳細訊息。此外，為了調查效能問題，通常只需一個迭代的詳細訊息就足夠了。這種按需詳細輸出功能使得可以控制詳細訊息轉儲的範圍。在下面的示例中，詳細訊息僅在第二次推理時轉儲。

import torch

model(data)
with torch.backends.mkl.verbose(torch.backends.mkl.VERBOSE_ON):
    model(data)
引數

level – 詳細級別 - VERBOSE_OFF：停用詳細輸出 - VERBOSE_ON：啟用詳細輸出

torch.backends.mkldnn#

torch.backends.mkldnn.is_available()[source]#
class torch.backends.mkldnn.verbose(level)[source]#

按需 oneDNN（前 MKL-DNN）詳細輸出功能。

為了更容易地除錯效能問題，oneDNN 可以在執行核心時轉儲包含核心大小、輸入資料大小和執行時間的詳細訊息。詳細輸出功能可以透過名為 DNNL_VERBOSE 的環境變數呼叫。但是，這種方法會在所有步驟中轉儲訊息。這些是大量的詳細訊息。此外，為了調查效能問題，通常只需一個迭代的詳細訊息就足夠了。這種按需詳細輸出功能使得可以控制詳細訊息轉儲的範圍。在下面的示例中，詳細訊息僅在第二次推理時轉儲。

import torch

model(data)
with torch.backends.mkldnn.verbose(torch.backends.mkldnn.VERBOSE_ON):
    model(data)
引數

level – 詳細級別 - VERBOSE_OFF：停用詳細輸出 - VERBOSE_ON：啟用詳細輸出 - VERBOSE_ON_CREATION：啟用詳細輸出，包括 oneDNN 核心建立

torch.backends.nnpack#

torch.backends.nnpack.is_available()[source]#

返回 PyTorch 是否已構建支援 NNPACK。

torch.backends.nnpack.flags(enabled=False)[source]#

用於全域性設定 nnpack 是否啟用的上下文管理器

torch.backends.nnpack.set_flags(_enabled)[source]#

設定 nnpack 是否全域性啟用

torch.backends.openmp#

torch.backends.openmp.is_available()[source]#

返回 PyTorch 是否已構建支援 OpenMP。

torch.backends.opt_einsum#

torch.backends.opt_einsum.is_available()[source]#

返回一個布林值，指示 opt_einsum 當前是否可用。

您必須安裝 opt-einsum 才能使 torch 自動最佳化 einsum。要使 opt-einsum 可用，您可以將其與 torch 一起安裝：pip install torch[opt-einsum] 或單獨安裝：pip install opt-einsum。如果安裝了該軟體包，torch 將自動匯入並按需使用它。使用此函式檢查 opt-einsum 是否已由 torch 安裝並正確匯入。

返回型別

布林值

torch.backends.opt_einsum.get_opt_einsum()[source]#

如果 opt_einsum 當前可用，則返回 opt_einsum 包；否則返回 None。

返回型別

任何

torch.backends.opt_einsum.enabled#

一個 bool，用於控制是否啟用 opt_einsum（預設值為 True）。如果啟用，torch.einsum 將使用 opt_einsum（https://optimized-einsum.readthedocs.io/en/stable/path_finding.html）以計算最佳收縮路徑以獲得更快的效能。

如果 opt_einsum 不可用，torch.einsum 將回退到從左到右的預設收縮路徑。

torch.backends.opt_einsum.strategy#

一個 str，用於指定當 torch.backends.opt_einsum.enabled 為 True 時要嘗試的策略。預設情況下，torch.einsum 將嘗試“auto”策略，但“greedy”和“optimal”策略也受支援。請注意，“optimal”策略對輸入數量是階乘的，因為它會嘗試所有可能的路徑。有關更多詳細資訊，請參閱 opt_einsum 的文件（https://optimized-einsum.readthedocs.io/en/stable/path_finding.html）。

torch.backends.xeon#

在此頁面上
在 GitHub 上編輯
顯示原始碼
PyTorch 庫