torch.backends

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

这些后端包括：

• torch.backends.cpu

• torch.backends.cuda

• torch.backends.cudnn

• torch.backends.mps

• torch.backends.mkl

• torch.backends.mkldnn

• 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

A bool that controls whether TensorFloat-32 tensor cores may be used in matrix multiplications on Ampere or newer GPUs. See TensorFloat-32(TF32) on Ampere devices.

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_linalg_library(backend=None)

警告

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

当 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.SDPBackend(value)

用于缩放点积注意力后端的枚举类。

警告

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

此类需要与以下定义的枚举保持一致： pytorch/aten/src/ATen/native/transformers/sdp_utils_cpp.h

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.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=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

A bool that controls where TensorFloat-32 tensor cores may be used in cuDNN convolutions on Ampere or newer GPUs. See TensorFloat-32(TF32) on Ampere devices.

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.mps

torch.backends.mps.is_available()

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

Return type

布尔

torch.backends.mps.is_built()

返回 PyTorch 是否以 MPS 支持构建。请注意，这并不一定意味着 MPS 是可用的；这只是表示如果在此 PyTorch 二进制文件运行在具有正常 MPS 驱动程序和设备的机器上时，我们可以使用它。

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.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