PyTorch torch.backends 模块详解
一、模块概述
torch.backends 是 PyTorch 中用于控制各种后端行为的模块，支持 CPU、CUDA、cuDNN、MPS、MKL、OpenMP 等多种硬件和库后端。
二、各后端详细介绍
1. torch.backends.cpu - CPU 后端
| 函数/属性 | 说明 |
|---|---|
get_cpu_capability() |
获取 CPU 能力，返回值：DEFAULT、VSX、Z VECTOR、NO AVX、AVX2、AVX512、SVE256 |
2. torch.backends.cuda - CUDA 后端
| 函数/属性 | 说明 |
|---|---|
is_built() |
检查 PyTorch 是否以 CUDA 支持编译 |
matmul.allow_tf32 |
控制是否使用 TensorFloat-32（Ampere 及以上 GPU） |
matmul.allow_fp16_reduced_precision_reduction |
控制 fp16 GEMM 是否允许降低精度归约 |
matmul.allow_bf16_reduced_precision_reduction |
控制 bf16 GEMM 是否允许降低精度归约 |
cufft_plan_cache |
cuFFT 计划缓存管理（size/max_size/clear） |
preferred_blas_library() |
选择 BLAS 库：cuBLAS/cuBLASLt/CK |
preferred_rocm_fa_library() |
ROCm 下选择 Flash Attention 后端 |
preferred_linalg_library() |
选择线性代数库：cuSOLVER/MAGMA |
SDPA (Scaled Dot Product Attention) 相关：
| 函数 | 功能 |
|---|---|
flash_sdp_enabled() / enable_flash_sdp() |
Flash Attention 开关 |
mem_efficient_sdp_enabled() / enable_mem_efficient_sdp() |
内存高效注意力开关 |
math_sdp_enabled() / enable_math_sdp() |
数学实现注意力开关 |
cudnn_sdp_enabled() / enable_cudnn_sdp() |
cuDNN 注意力开关 |
sdp_kernel() |
上下文管理器，临时启用/禁用指定后端 |
is_flash_attention_available() |
检查 FlashAttention 是否可用 |
can_use_flash_attention() |
检查给定参数是否能使用 FlashAttention |
can_use_efficient_attention() |
检查是否能使用 efficient attention |
can_use_cudnn_attention() |
检查是否能使用 cuDNN attention |
3. torch.backends.cudnn - cuDNN 后端
| 函数/属性 | 说明 |
|---|---|
version() |
返回 cuDNN 版本 |
is_available() |
检查 cuDNN 是否可用 |
enabled |
控制 cuDNN 是否启用 |
allow_tf32 |
控制 TensorFloat-32 使用 |
deterministic |
是否仅使用确定性卷积算法 |
benchmark |
是否对多种卷积算法进行基准测试 |
benchmark_limit |
基准测试时尝试的最大算法数量 |
4. torch.backends.cusparselt - cuSPARSELt 后端
| 函数 | 说明 |
|---|---|
version() |
返回 cuSPARSELt 版本 |
is_available() |
检查 cuSPARSELt 是否可用 |
5. torch.backends.mha - 多头注意力后端
| 函数 | 说明 |
|---|---|
get_fastpath_enabled() |
获取 TransformerEncoder/MHA 快速路径是否启用 |
set_fastpath_enabled() |
设置快速路径启用状态 |
6. torch.backends.miopen - MIOpen 后端 (ROCm)
| 属性 | 说明 |
|---|---|
immediate |
控制是否使用 Immediate Mode |
7. torch.backends.mps - Apple MPS 后端
| 函数 | 说明 |
|---|---|
is_available() |
检查 MPS 是否可用 |
is_built() |
检查是否以 MPS 支持编译 |
8. torch.backends.mkl - Intel MKL 后端
| 函数 | 说明 |
|---|---|
is_available() |
检查是否以 MKL 支持编译 |
verbose(level) |
按需启用 MKL 详细日志（上下文管理器） |
日志级别： VERBOSE_OFF、VERBOSE_ON
9. torch.backends.mkldnn - oneDNN 后端
| 函数 | 说明 |
|---|---|
is_available() |
检查是否以 MKL-DNN 支持编译 |
verbose(level) |
按需启用 oneDNN 详细日志 |
日志级别： VERBOSE_OFF、VERBOSE_ON、VERBOSE_ON_CREATION
10. torch.backends.nnpack - NNPACK 后端
| 函数 | 说明 |
|---|---|
is_available() |
检查是否以 NNPACK 支持编译 |
flags() |
上下文管理器设置 NNPACK 启用状态 |
set_flags() |
全局设置 NNPACK 启用状态 |
11. torch.backends.openmp - OpenMP 后端
| 函数 | 说明 |
|---|---|
is_available() |
检查是否以 OpenMP 支持编译 |
12. torch.backends.opt_einsum - opt_einsum 优化后端
| 函数/属性 | 说明 |
|---|---|
is_available() |
检查 opt_einsum 是否可用 |
get_opt_einsum() |
获取 opt_einsum 包（如果可用） |
enabled |
控制是否启用 opt_einsum（默认 True） |
strategy |
设置收缩策略：auto（默认）、greedy、optimal |
安装方式：
pip install torch[opt-einsum]或pip install opt-einsum
13. torch.backends.xeon - Intel Xeon 后端
文档中未列出具体 API。
三、实用代码示例
python
import torch
# 1. 查看 CPU 能力
print(torch.backends.cpu.get_cpu_capability()) # 如: "AVX512"
# 2. 检查 CUDA 相关
if torch.backends.cuda.is_built():
print(f"CUDA 可用: {torch.backends.cuda.is_available()}")
print(f"cuDNN 版本: {torch.backends.cudnn.version()}")
# 3. 控制 cuDNN 行为
torch.backends.cudnn.benchmark = True # 自动寻找最优算法
torch.backends.cudnn.deterministic = True # 使用确定性算法
# 4. 控制 TF32（Ampere 及以上 GPU）
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
# 5. SDPA 后端控制
torch.backends.cuda.enable_flash_sdp(True)
torch.backends.cuda.enable_mem_efficient_sdp(True)
# 使用上下文管理器临时切换
with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False):
# 在此范围内只使用 Flash Attention
output = torch.nn.functional.scaled_dot_product_attention(q, k, v)
# 6. MKL/oneDNN 详细日志
with torch.backends.mkl.verbose(torch.backends.mkl.VERBOSE_ON):
# 只在此范围内输出 MKL 详细日志
model(data)
with torch.backends.mkldnn.verbose(torch.backends.mkldnn.VERBOSE_ON):
# 只在此范围内输出 oneDNN 详细日志
model(data)
# 7. 优化 einsum
torch.backends.opt_einsum.enabled = True
torch.backends.opt_einsum.strategy = "greedy" # 或 "optimal", "auto"四、环境变量速查
| 环境变量 | 作用 |
|---|---|
TORCH_BLAS_PREFER_CUBLASLT=1 |
优先使用 cuBLASLt |
TORCH_LINALG_PREFER_CUSOLVER=1 |
优先使用 cuSOLVER |
TORCH_ROCM_FA_PREFER_CK=1 |
ROCm 下优先使用 CK 作为 Flash Attention 后端 |
MKL_VERBOSE |
全局启用 MKL 详细日志 |
DNNL_VERBOSE |
全局启用 oneDNN 详细日志 |
五、总结
torch.backends 模块是 PyTorch 性能调优的重要工具，通过它可以：
- 查询硬件能力 - 了解 CPU/GPU 支持特性
- 控制计算精度 - TF32、fp16/bf16 归约精度
- 选择计算后端 - cuBLAS/cuBLASLt、cuSOLVER/MAGMA 等
- 优化注意力计算 - Flash Attention、Memory Efficient Attention 等
- 调试性能问题 - MKL/oneDNN 按需详细日志！