torch.backends

这些后端包括：

torch.backends.cpu

torch.backends.cpu.get_cpu_capability()[源代码]

以字符串形式返回CPU能力。

可能的值包括："DEFAULT"、"VSX"、"Z VECTOR"、"NO AVX"、"AVX2"、"AVX512"

返回类型

str

torch.backends.cuda

torch.backends.cuda.is_built()[源代码]

返回 PyTorch 是否是带 CUDA 支持构建的。

需要注意的是，这并不一定表示CUDA可用；只是说，如果在配备正常工作的CUDA驱动程序和设备的机器上运行此PyTorch二进制文件，我们就能使用它。

torch.backends.cuda.matmul.allow_tf32

一个bool，用于控制在安培或更新的GPU上进行矩阵乘法时是否可以使用TensorFloat-32张量核心。参见安培（及以后）设备上的TensorFloat-32 (TF32)。

torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction

一个bool值，用于控制是否允许使用 fp16 精度的矩阵乘法（例如，具有 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

一个用于控制 cuFFT 计划缓存容量的 int。

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 选择的库不适合你的应用输入时，这个标志可能会提高性能。

返回类型

_BlasBackend

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

覆盖 PyTorch 在 CUDA 线性代数操作中选择 cuSOLVER 和 MAGMA 的启发式算法。

警告

此标志是试验性的，可能随时发生变化。

当 PyTorch 运行 CUDA 线性代数操作时，它通常会使用 cuSOLVER 或 MAGMA 库。如果两者都可用，PyTorch 会通过启发式方法来决定使用哪个库。此标志（一个 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()

返回类型

_LinalgBackend

类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)[源代码]

警告

此标志为测试版，可能随时更改。

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

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

警告

此标志为测试版，可能随时更改。

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

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

警告

此标志为测试版，可能随时更改。

开启或关闭数学缩放点积注意力。

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

警告

此标志为测试版，可能随时更改。

返回是否启用了数学缩放点积注意力中的 fp16/bf16 减少功能。

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

警告

此标志为测试版，可能随时更改。

启用或禁用数学缩放点积注意力中的 fp16/bf16 减量。

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

警告

此标志为测试版，可能随时更改。

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

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

警告

此标志为测试版，可能随时更改。

开启或关闭cuDNN缩放点积注意力。

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

检查 PyTorch 是否在构建时启用了 FlashAttention 来支持 scaled_dot_product_attention。

返回值

如果是构建并可用的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) – 一个包含查询、键和值张量，可选的注意力掩码，dropout 率以及表示是否为因果关系标志的 SDPAParams 实例。

• debug (bool) – 是否记录调试信息，以说明为何无法运行FlashAttention。默认值为False。

返回值

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

返回类型

bool

注意

此函数需要在支持CUDA的PyTorch环境下运行。在不支持CUDA的环境中，该函数将返回False。

torch.backends.cuda.can_use_efficient_attention(params, debug=False)[源代码]

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

参数

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

• debug (bool) – 是否通过 logging.warn 记录无法运行 efficient_attention 的原因。默认值为 False。

返回值

如果给定的参数可以与efficient_attention一起使用，则返回True，否则返回False。

返回类型

bool

注意

此函数需要在支持CUDA的PyTorch环境下运行。在不支持CUDA的环境中，该函数将返回False。

torch.backends.cuda.can_use_cudnn_attention(params, debug=False)[源代码]

检查是否可以使用cudnn_attention来替代scaled_dot_product_attention。

参数

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

• debug (bool) – 是否记录警告信息以说明为何无法运行cuDNN注意力机制。默认值为False。

返回值

如果给定的参数可以与cuDNN一起使用，则返回真；否则返回假。

返回类型

bool

注意

此函数需要在支持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，用于控制在安培或更新的GPU上cuDNN卷积中是否可以使用TensorFloat-32张量核心。参见安培（及以后）设备上的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为真时尝试的最大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。

提示:

即使 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 可用；只有在配备正常工作 MPS 驱动程序和设备的机器上运行此 PyTorch 二进制文件时，我们才能使用它。

返回类型

bool

torch.backends.mkl

torch.backends.mkl.is_available()[源代码]

返回 PyTorch 是否支持 MKL（数学核心函数库）。

类torch.backends.mkl.verbose(enable)[源代码]

按需的一MKL详尽输出功能。

为了方便调试性能问题，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 构建。

类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