评价此页
★ ★ ★ ★ ★

torch.backends#

创建于: 2020年9月16日 | 最后更新于: 2025年12月9日

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 张量核心。allow_tf32 将被弃用。请参阅 Ampere (及更高版本) 设备上的 TensorFloat-32 (TF32)。

torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction#

一个 bool，用于控制在使用 fp16 GEMM 时是否允许使用低精度归约（例如，使用 fp16 累加类型）。赋值一个元组 (allow_reduced_precision, allow_splitk) 还可以切换在使用 cuBLASLt 分派时是否可以使用 split-K 启发式。 allow_splitk 默认为 True。

torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction#

一个 bool，用于控制在使用 bf16 GEMM 时是否允许使用低精度归约。赋值一个元组 (allow_reduced_precision, allow_splitk) 还可以切换在使用 cuBLASLt 分派时是否可以使用 split-K 启发式。 allow_splitk 默认为 True。

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。

  • 如果设置为 “default”（默认值），则使用启发式方法在其他选项之间进行选择。

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

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

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

返回类型:

_BlasBackend

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

[仅限 ROCm] 覆盖 PyTorch 在 ROCm 环境中用于 Flash Attention 的后端。在 AOTriton 和 CK 之间选择

警告

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

当启用并期望 Flash Attention 时，PyTorch 默认使用 AOTriton 作为后端。此标志（一个 str）允许用户覆盖此后端以使用 composable_kernel

  • 如果设置为 “default”，则尽可能使用默认后端。目前是 AOTriton。

  • 如果设置为 “aotriton”，则尽可能使用 AOTriton。

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

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

  • 用户可以使用环境变量 TORCH_ROCM_FA_PREFER_CK=1 将首选库全局设置为 CK。

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

返回类型:

_ROCmFABackend

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

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

警告

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

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

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

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

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

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

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

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

当前支持的线性代数运算符

返回类型:

_LinalgBackend

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

警告

此标志是 Beta 版，可能会发生变化。

返回 flash 缩放点积注意力是否已启用。

torch.backends.cuda.enable_mem_efficient_sdp(enabled)[源码]#

警告

此标志是 Beta 版，可能会发生变化。

启用或禁用内存高效缩放点积注意力。

torch.backends.cuda.mem_efficient_sdp_enabled()[源码]#

警告

此标志是 Beta 版，可能会发生变化。

返回内存高效缩放点积注意力是否已启用。

torch.backends.cuda.enable_flash_sdp(enabled)[源码]#

警告

此标志是 Beta 版，可能会发生变化。

启用或禁用 flash 缩放点积注意力。

torch.backends.cuda.math_sdp_enabled()[源码]#

警告

此标志是 Beta 版，可能会发生变化。

返回 math 缩放点积注意力是否已启用。

torch.backends.cuda.enable_math_sdp(enabled)[源码]#

警告

此标志是 Beta 版，可能会发生变化。

启用或禁用 math 缩放点积注意力。

torch.backends.cuda.fp16_bf16_reduction_math_sdp_allowed()[源码]#

警告

此标志是 Beta 版，可能会发生变化。

返回 math 缩放点积注意力中的 fp16/bf16 归约是否已启用。

torch.backends.cuda.allow_fp16_bf16_reduction_math_sdp(enabled)[源码]#

警告

此标志是 Beta 版，可能会发生变化。

启用或禁用 math 缩放点积注意力中的 fp16/bf16 归约。

torch.backends.cuda.cudnn_sdp_enabled()[源码]#

警告

此标志是 Beta 版，可能会发生变化。

返回 cuDNN 缩放点积注意力是否已启用。

torch.backends.cuda.enable_cudnn_sdp(enabled)[源码]#

警告

此标志是 Beta 版，可能会发生变化。

启用或禁用 cuDNN 缩放点积注意力。

torch.backends.cuda.is_flash_attention_available()[源码]#

检查 PyTorch 是否使用 FlashAttention 构建用于 scaled_dot_product_attention。

返回:

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

返回类型:

布尔值

注意

此函数依赖于支持 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。

返回类型:

布尔值

注意

此函数依赖于支持 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。

返回类型:

布尔值

注意

此函数依赖于支持 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。

返回类型:

布尔值

注意

此函数依赖于支持 CUDA 的 PyTorch 构建。在非 CUDA 环境中，它将返回 False。

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

警告

此标志是 Beta 版，可能会发生变化。

此上下文管理器可用于临时启用或禁用缩放点积注意力的任何三个后端。退出上下文管理器后，将恢复标志的先前状态。

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 张量核心。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#

当 torch.backends.cudnn.benchmark 为 True 时，指定要尝试的最大 cuDNN 卷积算法数量的 int。将 benchmark_limit 设置为零可尝试所有可用的算法。请注意，此设置仅影响通过 cuDNN v8 API 分派的卷积。

torch.backends.cusparselt#

torch.backends.cusparselt.version()[源码]#

返回 cuSPARSELt 的版本

返回类型:

int | None

torch.backends.cusparselt.is_available()[源码]#

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

返回类型:

布尔值

torch.backends.mha#

torch.backends.mha.get_fastpath_enabled()[源码]#

返回 TransformerEncoder 和 MultiHeadAttention 的快速路径是否已启用，或者如果 jit 正在脚本化，则返回 True。

注意

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

返回类型:

布尔值

torch.backends.mha.set_fastpath_enabled(value)[源码]#

设置是否启用快速路径

torch.backends.miopen#

torch.backends.miopen.immediate#

一个 bool，如果为 True，则会导致 MIOpen 使用即时模式（https://rocm.docs.amd.com/projects/MIOpen/en/latest/how-to/find-and-immediate.html）。

torch.backends.mps#

torch.backends.mps.is_available()[源码]#

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

返回类型:

布尔值

torch.backends.mps.is_built()[源码]#

返回 PyTorch 是否使用 MPS 支持构建。

请注意，这并不一定意味着 MPS 可用；仅仅表示如果此 PyTorch 二进制文件在具有有效 MPS 驱动程序和设备的机器上运行，我们将能够使用它。

返回类型:

布尔值

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()[源码]#
class torch.backends.mkldnn.verbose(level)[源码]#

按需 oneDNN (原 MKL-DNN) 详细输出功能。

为了更方便地调试性能问题，oneDNN 在执行 kernel 时可以转储包含 kernel 大小、输入数据大小和执行时长等信息的详细消息。详细消息功能可以通过名为 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 kernel 创建

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 库