torch.backends#
创建日期：2020年9月16日 | 最后更新日期：2025年6月12日
torch.backends 控制 PyTorch 支持的各种后端行为。
这些后端包括
torch.backends.cputorch.backends.cudatorch.backends.cudnntorch.backends.cusparselttorch.backends.mhatorch.backends.mpstorch.backends.mkltorch.backends.mkldnntorch.backends.nnpacktorch.backends.openmptorch.backends.opt_einsumtorch.backends.xeon
torch.backends.cpu#
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 Tensor 核心。参见 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.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”，则在可能的情况下使用 cuBLAS。
如果设置为“cublaslt”，则在可能的情况下使用 cuBLASLt。
如果设置为“ck”，则在可能的情况下使用 CK。
如果设置为“default”（默认），则将使用启发式方法在其他选项之间进行选择。
如果没有输入，此函数返回当前首选的库。
用户可以使用环境变量 TORCH_BLAS_PREFER_CUBLASLT=1 来全局设置首选库为 cuBLASLt。此标志仅设置首选库的初始值，首选库仍可能在脚本中稍后通过此函数调用进行覆盖。
注意：当首选某个库时，如果首选库未实现所调用的操作，则仍可能使用其他库。如果 PyTorch 的库选择对于您的应用程序输入不正确，此标志可能会实现更好的性能。
- 返回类型
_BlasBackend
- torch.backends.cuda.preferred_rocm_fa_library(backend=None)[source]#
[仅限 ROCm] 在 ROCm 环境中覆盖 PyTorch 用于 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)[source]#
覆盖 PyTorch 用于在 cuSOLVER 和 MAGMA 之间为 CUDA 线性代数操作进行选择的启发式方法。
警告
此标志是实验性的，可能会更改。
当 PyTorch 运行 CUDA 线性代数操作时，它通常使用 cuSOLVER 或 MAGMA 库，如果两者都可用，它会通过启发式方法决定使用哪个。此标志（一个
str）允许覆盖这些启发式方法。如果设置为“cusolver”，则在可能的情况下使用 cuSOLVER。
如果设置为“magma”，则在可能的情况下使用 MAGMA。
如果设置为“default”（默认），则在 cuSOLVER 和 MAGMA 都可用时，将使用启发式方法在它们之间进行选择。
如果没有输入，此函数返回当前首选的库。
用户可以使用环境变量 TORCH_LINALG_PREFER_CUSOLVER=1 来全局设置首选库为 cuSOLVER。此标志仅设置首选库的初始值，首选库仍可能在脚本中稍后通过此函数调用进行覆盖。
注意：当首选某个库时，如果首选库未实现所调用的操作，则仍可能使用其他库。如果 PyTorch 的启发式库选择对于您的应用程序输入不正确，此标志可能会实现更好的性能。
当前支持的线性代数运算符
torch.linalg.eighvals()
- 返回类型
_LinalgBackend
- class torch.backends.cuda.SDPAParams#
- torch.backends.cuda.flash_sdp_enabled()[source]#
警告
此标志处于测试阶段，可能会更改。
返回是否启用了 Flash Scaled Dot Product Attention。
- torch.backends.cuda.enable_mem_efficient_sdp(enabled)[source]#
警告
此标志处于测试阶段，可能会更改。
启用或禁用内存高效的 Scaled Dot Product Attention。
- torch.backends.cuda.mem_efficient_sdp_enabled()[source]#
警告
此标志处于测试阶段，可能会更改。
返回是否启用了内存高效的 Scaled Dot Product Attention。
- torch.backends.cuda.enable_flash_sdp(enabled)[source]#
警告
此标志处于测试阶段，可能会更改。
启用或禁用 Flash Scaled Dot Product Attention。
- torch.backends.cuda.math_sdp_enabled()[source]#
警告
此标志处于测试阶段，可能会更改。
返回是否启用了数学 Scaled Dot Product Attention。
- torch.backends.cuda.enable_math_sdp(enabled)[source]#
警告
此标志处于测试阶段，可能会更改。
启用或禁用数学 Scaled Dot Product Attention。
- torch.backends.cuda.fp16_bf16_reduction_math_sdp_allowed()[source]#
警告
此标志处于测试阶段，可能会更改。
返回是否在数学 Scaled Dot Product Attention 中启用了 fp16/bf16 归约。
- torch.backends.cuda.allow_fp16_bf16_reduction_math_sdp(enabled)[source]#
警告
此标志处于测试阶段，可能会更改。
启用或禁用数学 Scaled Dot Product Attention 中的 fp16/bf16 归约。
- torch.backends.cuda.cudnn_sdp_enabled()[source]#
警告
此标志处于测试阶段，可能会更改。
返回是否启用了 cuDNN Scaled Dot Product Attention。
- torch.backends.cuda.enable_cudnn_sdp(enabled)[source]#
警告
此标志处于测试阶段，可能会更改。
启用或禁用 cuDNN Scaled Dot Product Attention。
- torch.backends.cuda.is_flash_attention_available()[source]#
检查 PyTorch 是否使用 FlashAttention 构建以用于 scaled_dot_product_attention。
- 返回
如果 FlashAttention 已构建并可用，则为 True；否则为 False。
- 返回类型
注意
此函数依赖于 PyTorch 的 CUDA 启用构建。在非 CUDA 环境中将返回 False。
- torch.backends.cuda.can_use_flash_attention(params, debug=False)[source]#
检查 FlashAttention 是否可用于 scaled_dot_product_attention。
- 参数
params (_SDPAParams) – 包含查询、键、值张量、可选注意力掩码、dropout 率以及指示注意力是否因果的标志的 SDPAParams 实例。
debug (bool) – 是否 logging.warn 调试信息，说明为何无法运行 FlashAttention。默认为 False。
- 返回
如果 FlashAttention 可以与给定参数一起使用，则为 True；否则为 False。
- 返回类型
注意
此函数依赖于 PyTorch 的 CUDA 启用构建。在非 CUDA 环境中将返回 False。
- torch.backends.cuda.can_use_efficient_attention(params, debug=False)[source]#
检查 efficient_attention 是否可用于 scaled_dot_product_attention。
- 参数
params (_SDPAParams) – 包含查询、键、值张量、可选注意力掩码、dropout 率以及指示注意力是否因果的标志的 SDPAParams 实例。
debug (bool) – 是否 logging.warn 信息，说明为何无法运行 efficient_attention。默认为 False。
- 返回
如果 efficient_attention 可以与给定参数一起使用，则为 True；否则为 False。
- 返回类型
注意
此函数依赖于 PyTorch 的 CUDA 启用构建。在非 CUDA 环境中将返回 False。
- torch.backends.cuda.can_use_cudnn_attention(params, debug=False)[source]#
检查 cudnn_attention 是否可用于 scaled_dot_product_attention。
- 参数
params (_SDPAParams) – 包含查询、键、值张量、可选注意力掩码、dropout 率以及指示注意力是否因果的标志的 SDPAParams 实例。
debug (bool) – 是否 logging.warn 信息，说明为何无法运行 cuDNN attention。默认为 False。
- 返回
如果 cuDNN 可以与给定参数一起使用，则为 True；否则为 False。
- 返回类型
注意
此函数依赖于 PyTorch 的 CUDA 启用构建。在非 CUDA 环境中将返回 False。
torch.backends.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.cusparselt#
torch.backends.mha#
torch.backends.mps#
torch.backends.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#
- 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.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)。