torch.backends

torch.backends 控制 PyTorch 支持的各种后端的行为。

这些后端包括：

• torch.backends.cpu

• torch.backends.cuda

• torch.backends.cudnn

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

Return type

字符串

torch.backends.cuda

torch.backends.cuda.is_built()

返回 PyTorch 是否具有 CUDA 支持。

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

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精度的归约（例如，具有fp16累加类型的归约）与fp16 GEMM结合。

torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction

一个bool，用于控制是否允许使用bf16 GEMMs进行精度降低的计算。

torch.backends.cuda.cufft_plan_cache

cufft_plan_cache 包含每个 CUDA 设备的 cuFFT 计划缓存。 查询特定设备 i 的缓存通过 torch.backends.cuda.cufft_plan_cache[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 之间进行选择。

警告

此标志为实验性质，可能会发生变化。

当 PyTorch 运行 CUDA BLAS 操作时，默认使用 cuBLAS，即使同时可用 cuBLAS 和 cuBLASLt 也是如此。 对于为 ROCm 构建的 PyTorch，hipBLAS 和 hipBLASLt 可能提供不同的性能。 此标志（str）允许覆盖使用的 BLAS 库。

• 如果设置为“cublas”，则在可能的情况下会使用cuBLAS。

• 如果设置为“cublaslt”，则在可能的情况下使用cuBLASLt。

• 在没有输入的情况下，此函数返回当前首选的库。

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

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

Return type

_BlasBackend

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 的启发式库选择对应用程序输入不正确的情况下，可能会获得更好的性能。

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

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

Return type

_LinalgBackend

class torch.backends.cuda.SDPAParams

torch.backends.cuda.flash_sdp_enabled()

警告

此标志处于测试阶段，可能会发生变化。

返回是否启用了 flash 点积注意力计算。

torch.backends.cuda.enable_mem_efficient_sdp(enabled)

警告

此标志处于测试阶段，可能会发生变化。

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

torch.backends.cuda.mem_efficient_sdp_enabled()

警告

此标志处于测试阶段，可能会发生变化。

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

torch.backends.cuda.enable_flash_sdp(enabled)

警告

此标志处于测试阶段，可能会发生变化。

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

torch.backends.cuda.math_sdp_enabled()

警告

此标志处于测试阶段，可能会发生变化。

返回数学缩放点积注意力是否启用。

torch.backends.cuda.enable_math_sdp(enabled)

警告

此标志处于测试阶段，可能会发生变化。

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

torch.backends.cuda.cudnn_sdp_enabled()

警告

此标志处于测试阶段，可能会发生变化。

返回 cuDNN 规模点积注意力是否启用。

torch.backends.cuda.enable_cudnn_sdp(enabled)

警告

此标志处于测试阶段，可能会发生变化。

启用或禁用 cuDNN 规模点积注意力。

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

检查 FlashAttention 是否可以用于缩放点积注意力。

Parameters

• 参数 (_SDPAParams) – SDPAParams 的一个实例，包含查询、键、值张量，可选的注意力掩码，丢弃率，以及一个指示注意力是否为因果的标志。

• 调试 (布尔值) – 是否记录为何FlashAttention无法运行的调试信息。 默认为 False。

Returns

如果给定参数可以使用FlashAttention，则为真；否则为假。

Return type

布尔

注意

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

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

检查 efficient_attention 是否可以用于 scaled_dot_product_attention。

Parameters

• 参数 (_SDPAParams) – SDPAParams 的一个实例，包含查询、键、值张量，可选的注意力掩码，丢弃率，以及一个指示注意力是否为因果的标志。

• 调试 (布尔值) – 是否使用日志记录.warn 记录有关为什么无法运行 efficient_attention 的信息。 默认为 False。

Returns

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

Return type

布尔

注意

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

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

警告

此标志处于测试阶段，可能会发生变化。

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

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

如果为 True，则会导致 cuDNN 只使用确定性的卷积算法的 bool。 另请参见 torch.are_deterministic_algorithms_enabled() 和 torch.use_deterministic_algorithms()。

torch.backends.cudnn.benchmark

如果为 bool，则会导致 cuDNN 对多个卷积算法进行基准测试并选择最快的一种。

torch.backends.cudnn.benchmark_limit

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

torch.backends.mha

torch.backends.mha.get_fastpath_enabled()

返回是否启用了TransformerEncoder和MultiHeadAttention的快速路径，或者如果jit处于脚本模式，则返回True。

..note:

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

Return type

布尔

torch.backends.mha.set_fastpath_enabled(value)

设置是否启用快速路径

torch.backends.mps

torch.backends.mps.is_available()

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

Return type

布尔

torch.backends.mps.is_built()

返回 PyTorch 是否带有 MPS 支持构建。

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

Return type

布尔

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)

Parameters

级别 – 详细程度级别 - 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)

Parameters

级别 – 详细程度级别 - 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 当前是否可用。

Return type

布尔

torch.backends.opt_einsum.get_opt_einsum()

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

Return type

任何

torch.backends.opt_einsum.enabled

A :class: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

A :class:str that specifies which strategies to try when torch.backends.opt_einsum.enabled is True. By default, torch.einsum will try the “auto” strategy, but the “greedy” and “optimal” strategies are also supported. Note that the “optimal” strategy is factorial on the number of inputs as it tries all possible paths. See more details in opt_einsum’s docs (https://optimized-einsum.readthedocs.io/en/stable/path_finding.html).

torch.backends.xeon