torch.backends¶
torch.backends 控制 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.cuda¶
- torch.backends.cuda.is_built()[源代码]¶
返回 PyTorch 是否使用 CUDA 支持构建。
请注意,这并不一定意味着 CUDA 可用;只是如果此 PyTorch 二进制文件在具有正常工作的 CUDA 驱动程序和设备的机器上运行,我们将能够使用它。
- 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 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)[源代码]¶
覆盖 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)[source]¶
覆盖 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.eighvals()
- 返回类型
_LinalgBackend
- class torch.backends.cuda.SDPAParams¶
- torch.backends.cuda.enable_mem_efficient_sdp(enabled)[source]¶
警告
此标志处于测试阶段,可能会发生变化。
启用或禁用内存高效缩放点积注意力。
- torch.backends.cuda.fp16_bf16_reduction_math_sdp_allowed()[source]¶
警告
此标志处于测试阶段,可能会发生变化。
返回数学缩放点积注意力中是否启用 fp16/bf16 缩减。
- torch.backends.cuda.allow_fp16_bf16_reduction_math_sdp(enabled)[source]¶
警告
此标志处于测试阶段,可能会发生变化。
启用或禁用数学缩放点积注意力中的 fp16/bf16 缩减。
- torch.backends.cuda.is_flash_attention_available()[source]¶
检查 PyTorch 是否使用 FlashAttention 构建了 scaled_dot_product_attention。
- 返回值
如果 FlashAttention 已构建并可用,则返回 True;否则返回 False。
- 返回类型
注意
此函数依赖于支持 CUDA 的 PyTorch 版本。在非 CUDA 环境中,它将返回 False。
- torch.backends.cuda.can_use_flash_attention(params, debug=False)[source]¶
检查 FlashAttention 是否可以在 scaled_dot_product_attention 中使用。
- 参数
params (_SDPAParams) – 一个 SDPAParams 实例,其中包含查询、键、值张量、可选的注意力掩码、dropout 率以及指示注意力是否因果的标志。
debug (bool) – 是否记录调试信息以说明为什么无法运行 FlashAttention。默认为 False。
- 返回值
如果 FlashAttention 可以使用给定的参数,则返回 True;否则返回 False。
- 返回类型
注意
此函数依赖于支持 CUDA 的 PyTorch 版本。在非 CUDA 环境中,它将返回 False。
- torch.backends.cuda.can_use_efficient_attention(params, debug=False)[source]¶
检查 efficient_attention 是否可以在 scaled_dot_product_attention 中使用。
- 参数
params (_SDPAParams) – 一个 SDPAParams 实例,其中包含查询、键、值张量、可选的注意力掩码、dropout 率以及指示注意力是否因果的标志。
debug (bool) – 是否使用信息记录警告以说明为什么无法运行 efficient_attention。默认为 False。
- 返回值
如果 efficient_attention 可以使用给定的参数,则返回 True;否则返回 False。
- 返回类型
注意
此函数依赖于支持 CUDA 的 PyTorch 版本。在非 CUDA 环境中,它将返回 False。
- torch.backends.cuda.can_use_cudnn_attention(params, debug=False)[source]¶
检查 cudnn_attention 是否可以在 scaled_dot_product_attention 中使用。
- 参数
params (_SDPAParams) – 一个 SDPAParams 实例,其中包含查询、键、值张量、可选的注意力掩码、dropout 率以及指示注意力是否因果的标志。
debug (bool) – 是否记录有关无法运行 cuDNN attention 的原因的信息。默认为 False。
- 返回值
如果 cuDNN 可以与给定的参数一起使用,则为 True;否则为 False。
- 返回类型
注意
此函数依赖于支持 CUDA 的 PyTorch 版本。在非 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)
- 参数
级别 – 详细输出级别 -
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)
- 参数
级别 – 详细输出级别 -
VERBOSE_OFF
:禁用详细输出 -VERBOSE_ON
:启用详细输出 -VERBOSE_ON_CREATION
:启用详细输出,包括 oneDNN 内核创建
torch.backends.nnpack¶
torch.backends.opt_einsum¶
- torch.backends.opt_einsum.get_opt_einsum()[source]¶
如果 opt_einsum 目前可用,则返回 opt_einsum 包,否则返回 None。
- 返回类型
- torch.backends.opt_einsum.enabled¶
一个 :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¶
一个 :class:
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)。