PyTorch 后端 ¶

PyTorch 后端控制着 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()[source][source]

返回 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

一个控制是否可以在 Ampere 或更新的 GPU 上使用 TensorFloat-32 矩阵乘法的 bool 。有关 Ampere（及以后）设备上的 TensorFloat-32（TF32）的更多信息，请参阅。

torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction

一个控制是否允许使用 fp16 累积类型进行精度降低的 bool 。

torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction

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

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 和 CK（ROCm 专用）。

警告

此标志为实验性，可能会更改。

当 PyTorch 运行 CUDA BLAS 操作时，默认使用 cuBLAS，即使 cuBLAS 和 cuBLASLt 都可用。对于为 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 的启发式库选择对您的应用程序输入不正确，此标志可能实现更好的性能。

当前支持的 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()

警告

此标志处于测试阶段，可能随时更改。

返回是否启用了闪存缩放点积注意力。

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()[source][source]

警告

此标志处于测试阶段，可能随时更改。

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

torch.backends.cuda.enable_cudnn_sdp(enabled)

警告

此标志处于测试阶段，可能随时更改。

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

torch.backends.cuda.is_flash_attention_available()[source][source]

检查 PyTorch 是否是用 FlashAttention 构建的缩放点积注意力。

返回:

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

返回类型:

bool 类型

备注

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

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

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

参数:

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

• debug（布尔值）- 是否记录.warn 调试信息，说明为什么 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) – 一个包含查询、键、值张量、可选的注意力掩码、dropout 率以及指示注意力是否因果的标志的 SDPAParams 实例。

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

返回值:

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

返回类型:

布尔型

备注

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

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

检查是否可以在 scaled_dot_product_attention 中使用 cudnn_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

控制是否启用 cuDNN 的 bool 。

torch.backends.cudnn.allow_tf32

一个 bool ，用于控制 TensorFloat-32 张量核心在 Ampere 或更新的 GPU 上用于 cuDNN 卷积的位置。请参阅 Ampere（及以后）设备上的 TensorFloat-32（TF32）。

torch.backends.cudnn.deterministic

一个 bool ，如果为 True，将导致 cuDNN 仅使用确定性卷积算法。另请参阅 torch.are_deterministic_algorithms_enabled() 和 torch.use_deterministic_algorithms() 。

torch.backends.cudnn.benchmark

如果设置为 True，将导致 cuDNN 基准测试多个卷积算法并选择最快的。

torch.backends.cudnn.benchmark_limit

当 torch.backends.cudnn.benchmark 为 True 时，指定尝试的最大 cuDNN 卷积算法数量。将 benchmark_limit 设置为 0 以尝试每个可用的算法。请注意，此设置仅影响通过 cuDNN v8 API 分发的卷积。

torch.backends.cusparselt

torch.backends.cusparselt.version()[source][source]

返回 cuSPARSELt 的版本

返回类型:

可选[int]

torch.backends.cusparselt.is_available()[来源][来源]

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

返回类型:

bool

torch.backends.mha

torch.backends.mha.get_fastpath_enabled()[source][source]

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

备注

即使 get_fastpath_enabled 返回 True ，也可能不会运行 fastpath，除非所有输入条件都满足。

返回类型:

bool

torch.backends.mha.set_fastpath_enabled(value)[source][source]

设置是否启用快速路径

torch.backends.mps

torch.backends.mps.is_available()[source][source]

返回一个布尔值，指示 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)

参数:

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

参数:

级别 - 详细输出级别 - VERBOSE_OFF : 禁用详细输出 - VERBOSE_ON : 启用详细输出 - VERBOSE_ON_CREATION : 启用详细输出，包括 oneDNN 内核创建

torch.backends.nnpack

torch.backends.nnpack.is_available()[source][source]

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

torch.backends.nnpack.flags(enabled=False)[source][source]

全局启用 nnpack 的上下文管理器

torch.backends.nnpack.set_flags(_enabled)[source][source]

设置是否全局启用 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 正确导入。

返回类型:

布尔值

torch.backends.opt_einsum.get_opt_einsum()

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

返回类型:

任何

torch.backends.opt_einsum.enabled

控制是否启用 opt_einsum（默认启用）。如果是，torch.einsum 将使用 opt_einsum（https://optimized-einsum.readthedocs.io/en/stable/path_finding.html）来计算最优收缩路径，以实现更快的性能。

如果 opt_einsum 不可用，torch.einsum 将回退到默认的从左到右的收缩路径。

torch.backends.opt_einsum.strategy

指定当 torch.backends.opt_einsum.enabled 为 True 时尝试哪些策略的 str 。默认情况下，torch.einsum 将尝试“auto”策略，但还支持“greedy”和“optimal”策略。请注意，“optimal”策略与输入数量呈阶乘关系，因为它尝试所有可能的路径。更多详细信息请参阅 opt_einsum 的文档（https://optimized-einsum.readthedocs.io/en/stable/path_finding.html）。

torch.backends.xeon