torch.backends

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()

返回 CPU 功能，以字符串值表示。

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

返回类型

str

torch.backends.cuda

torch.backends.cuda.is_built()

返回 PyTorch 是否构建时支持 CUDA。

请注意，这不一定意味着 CUDA 可用；只是如果这个 PyTorch 二进制文件在具有工作 CUDA 驱动程序和设备的机器上运行，我们将能够使用它。

torch.backends.cuda.matmul.allow_tf32

一个 bool，用于控制在 Ampere 或更新的 GPU 上的矩阵乘法中是否可以使用 TensorFloat-32 张量核心。请参阅 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)

覆盖 PyTorch 用于 BLAS 操作的库。在 cuBLAS、cuBLASLt 和 CK [仅限 ROCm] 之间选择。

警告

此标志是实验性的，可能会发生变化。

当 PyTorch 运行 CUDA BLAS 操作时，即使 cuBLAS 和 cuBLASLt 都可用，默认情况下也会使用 cuBLAS。对于为 ROCm 构建的 PyTorch，hipBLAS、hipBLASLt 和 CK 可能会提供不同的性能。此标志（一个 str）允许覆盖要使用的 BLAS 库。

• 如果设置了 “cublas”，则尽可能使用 cuBLAS。

• 如果设置了 “cublaslt”，则尽可能使用 cuBLASLt。

• 如果设置了 “ck”，则尽可能使用 CK。

• 当没有输入时，此函数返回当前首选库。

• 用户可以使用环境变量 TORCH_BLAS_PREFER_CUBLASLT=1 将首选库全局设置为 cuBLASLt。此标志仅设置首选库的初始值，首选库仍可能在稍后的脚本中被此函数调用覆盖。

注意：当首选某个库时，如果首选库未实现调用的操作，则仍可能使用其他库。如果 PyTorch 的库选择对于您的应用程序的输入不正确，则此标志可能会获得更好的性能。

返回类型

_BlasBackend

torch.backends.cuda.preferred_linalg_library(backend=None)

覆盖 PyTorch 用于在 cuSOLVER 和 MAGMA 之间为 CUDA 线性代数运算进行选择的启发式方法。

警告

此标志是实验性的，可能会发生变化。

当 PyTorch 运行 CUDA 线性代数运算时，它通常使用 cuSOLVER 或 MAGMA 库，如果两者都可用，它会使用启发式方法决定使用哪个库。此标志（一个 str）允许覆盖这些启发式方法。

• 如果设置了 “cusolver”，则尽可能使用 cuSOLVER。

• 如果设置了 “magma”，则尽可能使用 MAGMA。

• 如果设置了 “default”（默认值），则如果 cuSOLVER 和 MAGMA 都可用，将使用启发式方法在两者之间进行选择。

• 当没有输入时，此函数返回当前首选库。

• 用户可以使用环境变量 TORCH_LINALG_PREFER_CUSOLVER=1 将首选库全局设置为 cuSOLVER。此标志仅设置首选库的初始值，首选库仍可能在稍后的脚本中被此函数调用覆盖。

注意：当首选某个库时，如果首选库未实现调用的操作，则仍可能使用其他库。如果 PyTorch 的启发式库选择对于您的应用程序的输入不正确，则此标志可能会获得更好的性能。

当前支持的 linalg 运算符

• torch.linalg.inv()

• torch.linalg.inv_ex()

• torch.linalg.cholesky()

• torch.linalg.cholesky_ex()

• torch.cholesky_solve()

• torch.cholesky_inverse()

• torch.linalg.lu_factor()

• torch.linalg.lu()

• torch.linalg.lu_solve()

• torch.linalg.qr()

• torch.linalg.eigh()

• torch.linalg.eighvals()

• torch.linalg.svd()

• torch.linalg.svdvals()

返回类型

_LinalgBackend

class torch.backends.cuda.SDPAParams

torch.backends.cuda.flash_sdp_enabled()

警告

此标志是 beta 版，可能会发生更改。

返回是否启用了 flash scaled dot product attention。

torch.backends.cuda.enable_mem_efficient_sdp(enabled)

警告

此标志是 beta 版，可能会发生更改。

启用或禁用内存高效的 scaled dot product attention。

torch.backends.cuda.mem_efficient_sdp_enabled()

警告

此标志是 beta 版，可能会发生更改。

返回是否启用了内存高效的 scaled dot product attention。

torch.backends.cuda.enable_flash_sdp(enabled)

警告

此标志是 beta 版，可能会发生更改。

启用或禁用 flash scaled dot product attention。

torch.backends.cuda.math_sdp_enabled()

警告

此标志是 beta 版，可能会发生更改。

返回是否启用了 math scaled dot product attention。

torch.backends.cuda.enable_math_sdp(enabled)

警告

此标志是 beta 版，可能会发生更改。

启用或禁用 math scaled dot product attention。

torch.backends.cuda.fp16_bf16_reduction_math_sdp_allowed()

警告

此标志是 beta 版，可能会发生更改。

返回是否在 math scaled dot product attention 中启用了 fp16/bf16 归约。

torch.backends.cuda.allow_fp16_bf16_reduction_math_sdp(enabled)

警告

此标志是 beta 版，可能会发生更改。

在 math scaled dot product attention 中启用或禁用 fp16/bf16 归约。

torch.backends.cuda.cudnn_sdp_enabled()

警告

此标志是 beta 版，可能会发生更改。

返回是否启用了 cuDNN scaled dot product attention。

torch.backends.cuda.enable_cudnn_sdp(enabled)

警告

此标志是 beta 版，可能会发生更改。

启用或禁用 cuDNN scaled dot product attention。

torch.backends.cuda.is_flash_attention_available()

检查 PyTorch 构建时是否包含用于 scaled_dot_product_attention 的 FlashAttention。

返回

如果 FlashAttention 已构建且可用，则为 True；否则为 False。

返回类型

bool

注意

此函数依赖于启用 CUDA 的 PyTorch 构建版本。在非 CUDA 环境中，它将返回 False。

torch.backends.cuda.can_use_flash_attention(params, debug=False)

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

参数

• params (_SDPAParams) – SDPAParams 的实例，包含查询、键、值张量，以及可选的注意力掩码、dropout 率和指示注意力是否为因果关系的标志。

• debug (bool) – 是否记录警告调试信息，说明为什么无法运行 FlashAttention。默认为 False。

返回

如果可以使用给定的参数使用 FlashAttention，则为 True；否则为 False。

返回类型

bool

注意

此函数依赖于启用 CUDA 的 PyTorch 构建版本。在非 CUDA 环境中，它将返回 False。

torch.backends.cuda.can_use_efficient_attention(params, debug=False)

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

参数

• params (_SDPAParams) – SDPAParams 的实例，包含查询、键、值张量，以及可选的注意力掩码、dropout 率和指示注意力是否为因果关系的标志。

• debug (bool) – 是否记录警告信息，说明为什么无法运行 efficient_attention。默认为 False。

返回

如果可以使用给定的参数使用 efficient_attention，则为 True；否则为 False。

返回类型

bool

注意

此函数依赖于启用 CUDA 的 PyTorch 构建版本。在非 CUDA 环境中，它将返回 False。

torch.backends.cuda.can_use_cudnn_attention(params, debug=False)

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

参数

• params (_SDPAParams) – SDPAParams 的实例，包含查询、键、值张量，以及可选的注意力掩码、dropout 率和指示注意力是否为因果关系的标志。

• debug (bool) – 是否记录警告信息，说明为什么无法运行 cuDNN attention。默认为 False。

返回

如果可以使用给定的参数使用 cuDNN，则为 True；否则为 False。

返回类型

bool

注意

此函数依赖于启用 CUDA 的 PyTorch 构建版本。在非 CUDA 环境中，它将返回 False。

torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=True, enable_cudnn=True)

警告

此标志是 beta 版，可能会发生更改。

此上下文管理器可用于临时启用或禁用 scaled dot product attention 的三个后端中的任何一个。退出上下文管理器后，标志的先前状态将恢复。

torch.backends.cudnn

torch.backends.cudnn.version()

返回 cuDNN 的版本。

torch.backends.cudnn.is_available()

返回一个布尔值，指示当前 CUDNN 是否可用。

torch.backends.cudnn.enabled

一个 bool，用于控制是否启用 cuDNN。

torch.backends.cudnn.allow_tf32

一个 bool，用于控制在 Ampere 或更新的 GPU 上的 cuDNN 卷积中是否可以使用 TensorFloat-32 张量核心。请参阅 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()

返回 cuSPARSELt 的版本

返回类型

Optional[int]

torch.backends.cusparselt.is_available()

返回一个布尔值，指示 cuSPARSELt 当前是否可用。

返回类型

bool

torch.backends.mha

torch.backends.mha.get_fastpath_enabled()

返回是否为 TransformerEncoder 和 MultiHeadAttention 启用了快速路径，如果 jit 正在进行脚本编译，则返回 True。

..note

即使 get_fastpath_enabled 返回 True，快速路径也可能不会运行，除非满足输入的所有条件。

返回类型

bool

torch.backends.mha.set_fastpath_enabled(value)

设置是否启用快速路径

torch.backends.mps

torch.backends.mps.is_available()

返回一个布尔值，指示 MPS 当前是否可用。

返回类型

bool

torch.backends.mps.is_built()

返回 PyTorch 是否构建时启用了 MPS 支持。

请注意，这并不一定意味着 MPS 可用；只是如果这个 PyTorch 二进制文件在具有正常工作的 MPS 驱动程序和设备的机器上运行，我们将能够使用它。

返回类型

bool

torch.backends.mkl

torch.backends.mkl.is_available()

返回 PyTorch 是否构建时启用了 MKL 支持。

class torch.backends.mkl.verbose(enable)

按需 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()

返回 PyTorch 是否构建时启用了 MKL-DNN 支持。

class torch.backends.mkldnn.verbose(level)

按需 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()

返回 PyTorch 是否构建时启用了 NNPACK 支持。

torch.backends.nnpack.flags(enabled=False)

用于全局设置是否启用 nnpack 的上下文管理器

torch.backends.nnpack.set_flags(_enabled)

全局设置是否启用 nnpack

torch.backends.openmp

torch.backends.openmp.is_available()

返回 PyTorch 是否构建时启用了 OpenMP 支持。

torch.backends.opt_einsum

torch.backends.opt_einsum.is_available()

返回一个布尔值，指示 opt_einsum 当前是否可用。

您必须安装 opt-einsum 才能使 torch 自动优化 einsum。要使 opt-einsum 可用，您可以将其与 torch 一起安装：pip install torch[opt-einsum] 或单独安装： pip install opt-einsum。如果安装了软件包，torch 将自动导入它并相应地使用它。使用此函数检查 opt-einsum 是否已安装并由 torch 正确导入。

返回类型

bool

torch.backends.opt_einsum.get_opt_einsum()

如果 opt_einsum 当前可用，则返回 opt_einsum 包，否则返回 None。

返回类型

Any

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