PyTorch API 3 - mps、xpu、backends、导出
文章目录
- torch.mps
- MPS 性能分析器
- MPS 事件
- torch.xpu
- 随机数生成器
- 流与事件
- 内存管理
- torch.mtia
- 流与事件
- torch.mtia.memory
- 元设备
- 元张量操作惯用法
- torch.backends
- 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.export
- 概述
- 现有框架
- 导出 PyTorch 模型
- 示例
- 非严格导出模式
- 训练与推理导出功能
- 表达动态性
- 序列化
- 特化机制
- 输入张量形状
- Python 基本类型
- Python 容器
- torch.export 的局限性
- 图中断问题
- 数据/形状相关的控制流
- 算子缺少 Fake/Meta/Abstract 内核实现
- 扩展阅读
- API 参考
torch.mps
该包提供了在Python中访问MPS（Metal Performance Shaders）后端的接口。Metal是苹果公司用于编程金属GPU（图形处理器）的API。使用MPS意味着可以通过在金属GPU上运行工作来实现更高的性能。详情请参阅https://developer.apple.com/documentation/metalperformanceshaders。
device_count | 返回可用的MPS设备数量 |
|---|---|
synchronize | 等待MPS设备上所有流中的所有内核完成 |
get_rng_state | 以ByteTensor形式返回随机数生成器状态 |
set_rng_state | 设置随机数生成器状态 |
manual_seed | 设置生成随机数的种子 |
seed | 将生成随机数的种子设置为随机数 |
empty_cache | 释放缓存分配器当前持有的所有未占用缓存内存，以便其他GPU应用程序使用 |
set_per_process_memory_fraction | 设置MPS设备上限制进程内存分配的内存比例 |
current_allocated_memory | 返回当前张量占用的GPU内存（字节） |
driver_allocated_memory | 返回Metal驱动为该进程分配的总GPU内存（字节） |
recommended_max_memory | 返回GPU内存工作集大小的推荐最大值（字节） |
compile_shader | 从源代码编译计算着色器，并允许从Python运行时轻松调用其中定义的内核示例 |
MPS 性能分析器
profiler.start | 启动 MPS 后端的 OS Signpost 追踪功能 |
|---|---|
profiler.stop | 停止 MPS 后端的 OS Signpost 追踪功能 |
profiler.profile | 上下文管理器，用于启用 MPS 后端的 OS Signpost 追踪 |
profiler.is_capturing_metal | 检查 Metal 捕获是否正在进行 |
profiler.is_metal_capture_enabled | 检查 metal_capture 上下文管理器是否可用。要启用 Metal 捕获，需设置 MTL_CAPTURE_ENABLED 环境变量 |
profiler.metal_capture | 上下文管理器，用于将 Metal 调用捕获到 gputrace 中 |
MPS 事件
event.Event | MPS 事件的封装类 |
|---|
torch.xpu
该包提供了对XPU后端的支持，专门针对英特尔GPU进行了优化。
该包采用延迟初始化机制，因此您可以随时导入它，并通过is_available()方法检测当前系统是否支持XPU。
StreamContext | 用于选择指定流的上下文管理器 |
|---|---|
current_device | 返回当前选定设备的索引 |
current_stream | 返回指定设备当前选定的Stream |
device | 用于变更选定设备的上下文管理器 |
device_count | 返回可用的XPU设备数量 |
device_of | 将当前设备切换为给定对象所在设备的上下文管理器 |
get_arch_list | 返回本库编译时支持的XPU架构列表 |
get_device_capability | 获取设备的XPU计算能力 |
get_device_name | 获取设备名称 |
get_device_properties | 获取设备属性 |
get_gencode_flags | 返回本库编译时使用的XPU预编译构建标志 |
get_stream_from_external | 从外部SYCL队列返回Stream |
init | 初始化PyTorch的XPU状态 |
is_available | 返回布尔值表示XPU当前是否可用 |
is_initialized | 返回PyTorch的XPU状态是否已初始化 |
set_device | 设置当前设备 |
set_stream | 设置当前流（该API是设置流的封装接口） |
stream | 封装了用于选择指定流的StreamContext上下文管理器 |
synchronize | 等待XPU设备上所有流中的所有内核执行完成 |
随机数生成器
get_rng_state | 以ByteTensor形式返回指定GPU的随机数生成器状态。 |
|---|---|
get_rng_state_all | 返回一个ByteTensor列表，表示所有设备的随机数状态。 |
initial_seed | 返回当前GPU的随机种子值。 |
manual_seed | 为当前GPU设置随机数生成种子。 |
manual_seed_all | 为所有GPU设置随机数生成种子。 |
seed | 将当前GPU的随机数生成种子设置为一个随机数。 |
seed_all | 将所有GPU的随机数生成种子设置为随机数。 |
set_rng_state | 设置指定GPU的随机数生成器状态。 |
set_rng_state_all | 设置所有设备的随机数生成器状态。 |
流与事件
Event | XPU事件的封装类 |
|---|---|
Stream | XPU流的封装类 |
内存管理
empty_cache | 释放缓存分配器当前持有的所有未占用缓存内存，以便其他XPU应用程序可以使用这些内存。 |
|---|---|
max_memory_allocated | 返回给定设备上张量占用的最大GPU内存（以字节为单位）。 |
max_memory_reserved | 返回给定设备上缓存分配器管理的最大GPU内存（以字节为单位）。 |
mem_get_info | 返回给定设备上全局可用的和总的GPU内存。 |
memory_allocated | 返回给定设备上张量当前占用的GPU内存（以字节为单位）。 |
memory_reserved | 返回给定设备上缓存分配器当前管理的GPU内存（以字节为单位）。 |
memory_stats | 返回给定设备上XPU内存分配器统计信息的字典。 |
memory_stats_as_nested_dict | 以嵌套字典的形式返回memory_stats()的结果。 |
reset_accumulated_memory_stats | 重置XPU内存分配器跟踪的"累计"（历史）统计信息。 |
reset_peak_memory_stats | 重置XPU内存分配器跟踪的"峰值"统计信息。 |
torch.mtia
MTIA后端实现位于外部代码库，此处仅定义接口。
该包提供了在Python中访问MTIA后端的接口。
StreamContext | 用于选择指定流的上下文管理器 |
|---|---|
current_device | 返回当前选定设备的索引 |
current_stream | 返回指定设备当前选定的Stream |
default_stream | 返回指定设备的默认Stream |
device_count | 返回可用的MTIA设备数量 |
init | |
is_available | 如果MTIA设备可用则返回true |
is_initialized | 返回PyTorch的MTIA状态是否已初始化 |
memory_stats | 返回指定设备的MTIA内存分配器统计字典 |
get_device_capability | 以(主版本号,次版本号)元组形式返回指定设备的计算能力 |
empty_cache | 清空MTIA设备缓存 |
record_memory_history | 启用/禁用MTIA分配器的内存分析器 |
snapshot | 返回MTIA内存分配器历史记录字典 |
set_device | 设置当前设备 |
set_stream | 设置当前流。这是设置流的封装API |
stream | 封装了选择指定流的StreamContext上下文管理器 |
synchronize | 等待MTIA设备上所有流中的所有任务完成 |
device | 用于更改选定设备的上下文管理器 |
set_rng_state | 设置随机数生成器状态 |
get_rng_state | 以ByteTensor形式返回随机数生成器状态 |
DeferredMtiaCallError |
流与事件
Event | 查询和记录流状态，用于识别或控制跨流依赖关系以及测量时间。 |
|---|---|
Stream | 一个按先进先出(FIFO)顺序异步执行相应任务的顺序队列。 |
torch.mtia.memory
MTIA后端实现位于外部代码库中，此处仅定义接口。
该包提供了MTIA实现的设备内存管理支持。
memory_stats | 返回指定设备的MTIA内存分配器统计信息字典。 |
|---|
元设备
"元"设备是一种抽象设备，它表示仅记录元数据而不存储实际数据的张量。元张量主要有两个使用场景：
- 模型可以加载到元设备上，这样您可以在不将实际参数加载到内存的情况下获取模型的表示形式。如果您需要在加载真实数据之前对模型进行转换，这会非常有用。
- 大多数操作都可以在元张量上执行，生成新的元张量来描述如果在真实张量上执行该操作会得到什么结果。您可以用这种方式进行抽象分析，而无需花费计算时间或存储空间来表示实际张量。由于元张量没有真实数据，因此无法执行数据依赖的操作，如
torch.nonzero()或item()。在某些情况下，并非所有设备类型(例如CPU和CUDA)对同一操作都能产生完全相同的输出元数据；我们通常倾向于在这种情况下准确表示CUDA的行为。
警告：虽然原则上元张量计算应该总是比等效的CPU/CUDA计算更快，但许多元张量实现是用Python编写的，尚未移植到C++以提升速度，因此您可能会发现使用小型CPU张量时框架的绝对延迟更低。
元张量操作惯用法
可以通过指定 map_location='meta' 将对象加载到元设备上，使用 torch.load() 方法实现。
>>> torch.save(torch.randn(2), 'foo.pt')
>>> torch.load('foo.pt', map_location='meta')
tensor(..., device='meta', size=(2,))
如果你有一段任意代码，它在没有明确指定设备的情况下执行张量构建操作，你可以通过使用 torch.device() 上下文管理器来覆盖该行为，改为在元设备(meta device)上进行构建：
>>> with torch.device('meta'):
... print(torch.randn(30, 30))
...
tensor(..., device='meta', size=(30, 30))
这在神经网络模块构建中特别有用，因为通常无法显式传入设备进行初始化。
>>> from torch.nn.modules import Linear
>>> with torch.device('meta'):
... print(Linear(20, 30))
...
Linear(in_features=20, out_features=30, bias=True)
无法直接将元张量转换为CPU/CUDA张量，因为元张量不存储数据，我们无法确定新张量的正确数据值。
>>> torch.ones(5, device='meta').to("cpu")
Traceback (most recent call last):File "<stdin>", line 1, in <module>
NotImplementedError: Cannot copy out of meta tensor; no data!
使用类似 torch.empty_like() 的工厂函数来明确指定缺失数据的填充方式。
神经网络模块提供了一个便捷方法 torch.nn.Module.to_empty()，允许将模块转移到其他设备并保持所有参数未初始化状态。开发者需要手动显式地重新初始化这些参数。
>>> from torch.nn.modules import Linear
>>> with torch.device('meta'):
... m = Linear(20, 30)
>>> m.to_empty(device="cpu")
Linear(in_features=20, out_features=30, bias=True)
torch._subclasses.meta_utils 包含一系列未公开的工具函数，能够以高保真度将任意 Tensor 转换为等价的元数据 Tensor。这些 API 目前处于实验阶段，可能会随时做出不兼容的破坏性变更。
torch.backends
torch.backends 用于控制 PyTorch 所支持的各种后端的行为。
这些后端包括：
torch.backends.cputorch.backends.cudatorch.backends.cudnntorch.backends.cusparselttorch.backends.mhatorch.backends.mpstorch.backends.mkltorch.backends.mkldnntorch.backends.nnpacktorch.backends.openmptorch.backends.opt_einsumtorch.backends.xeon
torch.backends.cpu
torch.backends.cpu.get_cpu_capability()
返回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 GEMM运算中使用降低精度归约（例如采用fp16累加类型）的bool值。
torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction
A boolthat controls whether reduced precision reductions are allowed with bf16 GEMMs.
torch.backends.cuda.cufft_plan_cache
cufft_plan_cache contains the cuFFT plan caches for each CUDA device.
Query a specific device i’s cache via torch.backends.cuda.cufft_plan_cache[i].
torch.backends.cuda.cufft_plan_cache.size
A readonly int that shows the number of plans currently in a cuFFT plan cache.
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 和 cuBLASLt 都可用，默认仍会使用 cuBLAS。
针对 ROCm 构建的 PyTorch 中，hipBLAS、hipBLASLt 和 CK 可能提供不同的性能表现。
此标志（str 类型）允许覆盖要使用的 BLAS 库：
- 设为 “cublas” 时，将尽可能使用 cuBLAS
- 设为 “cublaslt” 时，将尽可能使用 cuBLASLt
- 设为 “ck” 时，将尽可能使用 CK
- 设为 “default”（默认值）时，将通过启发式方法在其他选项间选择
- 无输入时，此函数返回当前首选库
用户可通过环境变量 TORCH_BLAS_PREFER_CUBLASLT=1 全局设置首选库为 cuBLASLt。
注意：
1、此标志仅设置首选库的初始值，后续仍可能被脚本中的函数调用覆盖
2、当某个库被设为首选时，若该库未实现所调用的运算，仍可能使用其他库
3、若 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”（默认值）时，若两个库均可用则通过启发式规则选择
- 无输入参数时，函数返回当前优先使用的库
- 用户可通过环境变量
TORCH_LINALG_PREFER_CUSOLVER=1全局设置优先使用 cuSOLVER
注意：
1、此标志仅设置初始优先库，后续仍可通过脚本中的函数调用覆盖
2、即使设置了优先库，若该库未实现特定运算，仍可能使用其他库
3、当 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
class 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)– 包含查询(query)、键(key)、值(value)张量的 SDPAParams 实例，可选注意力掩码(attention mask)、dropout率，以及指示注意力是否为因果(causal)的标志。debug (bool)– 是否通过logging.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)- 包含query、key、value张量、可选attention掩码、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)
检查是否可以在 scaled_dot_product_attention 中使用 cudnn_attention。
参数
params (_SDPAParams)- 一个包含查询(query)、键(key)、值(value)张量的 SDPAParams 实例，可选注意力掩码(attention mask)、dropout率，以及一个指示注意力是否为因果(causal)的标志。debug (bool)- 是否通过 logging.warn 记录无法运行 cuDNN 注意力的原因信息。默认为 False。
返回
如果可以使用 cuDNN 处理给定参数则返回 True；否则返回 False。
返回类型：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
一个控制是否在Ampere或更新款GPU的cuDNN卷积中使用TensorFloat-32张量核心的bool值。详见Ampere（及后续）设备上的TensorFloat-32 (TF32)。
（说明：严格遵循核心翻译原则，保留代码块bool和链接格式，技术术语"TensorFloat-32/TF32/Ampere/cuDNN"不翻译，被动语态转为主动语态"是否使用"，并合并了原文因换行中断的连贯内容）
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 为 True 时，尝试的 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 一定可用；仅表示如果该 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)
参数
level– 详细级别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)
参数
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
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正确导入。
返回类型：bool
torch.backends.opt_einsum.get_opt_einsum()
如果当前可用 opt_einsum 包则返回该包，否则返回 None。
返回类型：Any
torch.backends.opt_einsum.enabled
一个控制是否启用 opt_einsum 的 bool 值（默认为 True）。如果启用，torch.einsum 将使用 opt_einsum（https://optimized-einsum.readthedocs.io/en/stable/path_finding.html）来计算最优的收缩路径以获得更快的性能（前提是 opt_einsum 可用）。
如果 opt_einsum 不可用，torch.einsum 将回退到默认的从左到右收缩路径。
torch.backends.opt_einsum.strategy
一个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)。
torch.backends.xeon
torch.export
警告：此功能为正在积极开发中的原型，未来将会有破坏性变更。
概述
torch.export.export() 接收一个 torch.nn.Module 并生成一个跟踪图，该图以提前编译（AOT）的方式仅表示函数的张量计算过程。生成的跟踪图随后可以用不同的输出执行或进行序列化。
import torch
from torch.export import exportclass Mod(torch.nn.Module):def forward(self, x: torch.Tensor, y: torch.Tensor) -torch.Tensor:a = torch.sin(x)b = torch.cos(y)return a + bexample_args = (torch.randn(10, 10), torch.randn(10, 10))exported_program: torch.export.ExportedProgram = export(Mod(), args=example_args
)
print(exported_program)ExportedProgram:class GraphModule(torch.nn.Module):def forward(self, x: "f32[10, 10]", y: "f32[10, 10]"):# code: a = torch.sin(x)sin: "f32[10, 10]" = torch.ops.aten.sin.default(x)# code: b = torch.cos(y)cos: "f32[10, 10]" = torch.ops.aten.cos.default(y)# code: return a + badd: f32[10, 10] = torch.ops.aten.add.Tensor(sin, cos)return (add,)Graph signature:ExportGraphSignature(input_specs=[InputSpec(kind=<InputKind.USER_INPUT: 1>, arg=TensorArgument(name='x'), target=None, persistent=None), InputSpec(kind=<InputKind.USER_INPUT: 1>, arg=TensorArgument(name='y'), target=None, persistent=None)], output_specs=[OutputSpec(kind=<OutputKind.USER_OUTPUT: 1>, arg=TensorArgument(name='add'), target=None)])Range constraints: {}torch.export 生成一个符合以下不变量的简洁中间表示（IR）。关于该 IR 的更多规范可查阅此处。
- 正确性：保证是原始程序的准确表示，并保持原始程序的调用约定。
- 规范化：图中不包含 Python 语义。原始程序中的子模块会被内联，形成一个完全扁平化的计算图。
- 图属性：该图是纯函数式的，意味着不包含具有副作用（如突变或别名）的操作。它不会改变任何中间值、参数或缓冲区。
- 元数据：图中包含在跟踪期间捕获的元数据，例如来自用户代码的堆栈跟踪。
在底层，torch.export 利用了以下最新技术：
- TorchDynamo (torch._dynamo) 是一个内部 API，它使用 CPython 的 Frame Evaluation API 功能来安全地跟踪 PyTorch 图。这极大地改善了图捕获体验，减少了完全跟踪 PyTorch 代码所需的改写工作。
- AOT Autograd 提供了一个功能化的 PyTorch 图，并确保图被分解/降级为 ATen 操作符集。
- Torch FX (torch.fx) 是图的底层表示，支持基于 Python 的灵活转换。
现有框架
torch.compile() 同样使用了与 torch.export 相同的 PT2 技术栈，但存在以下差异：
- JIT 与 AOT：
torch.compile()是一个即时（JIT）编译器，其设计目的并非用于生成可部署的编译产物。 - 部分图与完整图捕获：当
torch.compile()遇到模型不可追踪部分时，会触发"图中断"并回退到 Python 即时运行模式。相比之下，torch.export旨在获取 PyTorch 模型的完整计算图表示，因此遇到不可追踪内容时会直接报错。由于torch.export生成的完整图与 Python 特性及运行时完全解耦，该计算图可被保存、加载并在不同环境和语言中运行。 - 可用性权衡：
torch.compile()在遇到不可追踪内容时可回退至 Python 运行时，因此灵活性更高。而torch.export需要用户提供更多信息或重写代码以确保可追踪性。
与 torch.fx.symbolic_trace() 相比，torch.export 通过 TorchDynamo 在 Python 字节码层面进行追踪，使其能够处理不受 Python 运算符重载限制的任意 Python 结构。此外，torch.export 会精细追踪张量元数据，因此基于张量形状等条件的操作不会导致追踪失败。总体而言，torch.export 适用于更多用户程序，并能生成更底层的计算图（基于 torch.ops.aten 算子级别）。用户仍可将 torch.fx.symbolic_trace() 作为 torch.export 的预处理步骤。
与 torch.jit.script() 相比，torch.export 不捕获 Python 控制流或数据结构，但支持比 TorchScript 更多的 Python 语言特性（因其对 Python 字节码的覆盖更全面）。生成的计算图更简洁，仅包含直线控制流（显式控制流算子除外）。
与 torch.jit.trace() 相比，torch.export 具备可靠性：它能追踪对张量尺寸进行整数运算的代码，并记录所有必要的边界条件，以证明特定追踪结果对其他输入的有效性。
导出 PyTorch 模型
示例
主入口是通过 torch.export.export()，它接收一个可调用对象（torch.nn.Module、函数或方法）和示例输入，并将计算图捕获到 torch.export.ExportedProgram 中。示例如下：
import torch
from torch.export import export# Simple module for demonstration
class M(torch.nn.Module):def __init__(self) -None:super().__init__()self.conv = torch.nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1)self.relu = torch.nn.ReLU()self.maxpool = torch.nn.MaxPool2d(kernel_size=3)def forward(self, x: torch.Tensor, *, constant=None) -torch.Tensor:a = self.conv(x)a.add_(constant)return self.maxpool(self.relu(a))example_args = (torch.randn(1, 3, 256, 256),)
example_kwargs = {"constant": torch.ones(1, 16, 256, 256)}exported_program: torch.export.ExportedProgram = export(M(), args=example_args, kwargs=example_kwargs
)
print(exported_program)
ExportedProgram:class GraphModule(torch.nn.Module):def forward(self, p_conv_weight: "f32[16, 3, 3, 3]", p_conv_bias: "f32[16]", x: "f32[1, 3, 256, 256]", constant: "f32[1, 16, 256, 256]"):# code: a = self.conv(x)conv2d: "f32[1, 16, 256, 256]" = torch.ops.aten.conv2d.default(x, p_conv_weight, p_conv_bias, [1, 1], [1, 1])# code: a.add_(constant)add_: "f32[1, 16, 256, 256]" = torch.ops.aten.add_.Tensor(conv2d, constant)# code: return self.maxpool(self.relu(a))relu: "f32[1, 16, 256, 256]" = torch.ops.aten.relu.default(add_)max_pool2d: "f32[1, 16, 85, 85]" = torch.ops.aten.max_pool2d.default(relu, [3, 3], [3, 3])return (max_pool2d,)Graph signature:ExportGraphSignature(input_specs=[InputSpec(kind=<InputKind.PARAMETER: 2>,arg=TensorArgument(name='p_conv_weight'),target='conv.weight',persistent=None),InputSpec(kind=<InputKind.PARAMETER: 2>,arg=TensorArgument(name='p_conv_bias'),target='conv.bias',persistent=None),InputSpec(kind=<InputKind.USER_INPUT: 1>,arg=TensorArgument(name='x'),target=None,persistent=None),InputSpec(kind=<InputKind.USER_INPUT: 1>,arg=TensorArgument(name='constant'),target=None,persistent=None)],output_specs=[OutputSpec(kind=<OutputKind.USER_OUTPUT: 1>,arg=TensorArgument(name='max_pool2d'),target=None)])
Range constraints: {}
在检查ExportedProgram时，我们可以注意到以下几点：
torch.fx.Graph包含了原始程序的计算图，同时保留了原始代码记录以便于调试。- 图中仅包含此处列出的
torch.ops.aten运算符和自定义运算符，且完全可运行，不包含任何原地操作符（如torch.add_）。 - 参数（如卷积层的权重和偏置）被提升为图的输入节点，因此图中不再存在
torch.fx.symbolic_trace()结果中曾出现的get_attr节点。 torch.export.ExportGraphSignature对输入输出签名进行建模，并明确指定哪些输入是参数。- 图中每个节点输出的张量形状和数据类型都有标注。例如，
convolution节点将产生一个数据类型为torch.float32、形状为(1, 16, 256, 256)的张量。
非严格导出模式
在 PyTorch 2.3 版本中，我们引入了一种新的追踪模式——非严格模式。该功能目前仍处于完善阶段，如果您遇到任何问题，请在 GitHub 上提交 issue 并标记 “oncall: export” 标签。
非严格模式下，我们通过 Python 解释器执行程序追踪。您的代码会像在即时执行模式（eager mode）中一样运行，唯一的区别是所有 Tensor 对象都会被替换为 ProxyTensor，这些代理张量会将其所有操作记录到计算图中。
当前默认使用的是严格模式，该模式下我们首先使用 TorchDynamo（一个字节码分析引擎）进行程序追踪。TorchDynamo 实际上不会执行您的 Python 代码，而是对其进行符号化分析，并根据分析结果构建计算图。这种分析方式使得 torch.export 能够提供更强的安全性保证，但并非所有 Python 代码都受支持。
当您遇到 TorchDynamo 不支持的特性且难以解决时，如果确定相关 Python 代码并非计算所必需，就可以考虑使用非严格模式。例如：
import contextlib
import torchclass ContextManager():def __init__(self):self.count = 0def __enter__(self):self.count += 1def __exit__(self, exc_type, exc_value, traceback):self.count -= 1class M(torch.nn.Module):def forward(self, x):with ContextManager():return x.sin() + x.cos()export(M(), (torch.ones(3, 3),), strict=False) # Non-strict traces successfully
export(M(), (torch.ones(3, 3),)) # Strict mode fails with torch._dynamo.exc.Unsupported: ContextManager
在这个示例中，首次调用使用非严格模式（通过strict=False标志）能成功追踪，而第二次采用默认严格模式的调用则失败了，因为TorchDynamo无法支持上下文管理器。一种解决方案是重写代码（参见torch.export的限制），但考虑到上下文管理器不会影响模型中的张量计算，我们可以直接采用非严格模式的结果。
训练与推理导出功能
在 PyTorch 2.5 中，我们推出了名为 export_for_training() 的新 API。该功能目前仍在强化阶段，如果您遇到任何问题，请在 Github 上提交问题并标记 “oncall: export” 标签。
此 API 会生成包含所有 ATen 算子（包括功能性和非功能性）的最通用中间表示（IR），可用于 PyTorch Autograd 的即时训练模式。该 API 主要面向即时训练场景，例如 PT2 量化，并将很快成为 torch.export.export 的默认 IR。要深入了解这一变更背后的动机，请参阅 https://dev-discuss.pytorch.org/t/why-pytorch-does-not-need-a-new-standardized-operator-set/2206。
当此 API 与 run_decompositions() 结合使用时，您应该能够获得具有任何所需分解行为的推理 IR。
以下是一些示例：
class ConvBatchnorm(torch.nn.Module):def __init__(self) -None:super().__init__()self.conv = torch.nn.Conv2d(1, 3, 1, 1)self.bn = torch.nn.BatchNorm2d(3)def forward(self, x):x = self.conv(x)x = self.bn(x)return (x,)mod = ConvBatchnorm()
inp = torch.randn(1, 1, 3, 3)ep_for_training = torch.export.export_for_training(mod, (inp,))
print(ep_for_training)
ExportedProgram:class GraphModule(torch.nn.Module):def forward(self, p_conv_weight: "f32[3, 1, 1, 1]", p_conv_bias: "f32[3]", p_bn_weight: "f32[3]", p_bn_bias: "f32[3]", b_bn_running_mean: "f32[3]", b_bn_running_var: "f32[3]", b_bn_num_batches_tracked: "i64[]", x: "f32[1, 1, 3, 3]"):conv2d: "f32[1, 3, 3, 3]" = torch.ops.aten.conv2d.default(x, p_conv_weight, p_conv_bias)add_: "i64[]" = torch.ops.aten.add_.Tensor(b_bn_num_batches_tracked, 1)batch_norm: "f32[1, 3, 3, 3]" = torch.ops.aten.batch_norm.default(conv2d, p_bn_weight, p_bn_bias, b_bn_running_mean, b_bn_running_var, True, 0.1, 1e-05, True)return (batch_norm,)
从上述输出可以看出，export_for_training()生成的ExportedProgram与export()几乎相同，除了图中的运算符不同。可以看到我们以最通用的形式捕获了batch_norm操作。该操作是非功能性的，在运行推理时会被降级为不同的操作。
你还可以通过run_decompositions()从这个中间表示(IR)转换到推理IR，并进行任意自定义。
# Lower to core aten inference IR, but keep conv2d
decomp_table = torch.export.default_decompositions()
del decomp_table[torch.ops.aten.conv2d.default]
ep_for_inference = ep_for_training.run_decompositions(decomp_table)print(ep_for_inference)
ExportedProgram:class GraphModule(torch.nn.Module):def forward(self, p_conv_weight: "f32[3, 1, 1, 1]", p_conv_bias: "f32[3]", p_bn_weight: "f32[3]", p_bn_bias: "f32[3]", b_bn_running_mean: "f32[3]", b_bn_running_var: "f32[3]", b_bn_num_batches_tracked: "i64[]", x: "f32[1, 1, 3, 3]"):conv2d: "f32[1, 3, 3, 3]" = torch.ops.aten.conv2d.default(x, p_conv_weight, p_conv_bias)add: "i64[]" = torch.ops.aten.add.Tensor(b_bn_num_batches_tracked, 1)_native_batch_norm_legit_functional = torch.ops.aten._native_batch_norm_legit_functional.default(conv2d, p_bn_weight, p_bn_bias, b_bn_running_mean, b_bn_running_var, True, 0.1, 1e-05)getitem: "f32[1, 3, 3, 3]" = _native_batch_norm_legit_functional[0]getitem_3: "f32[3]" = _native_batch_norm_legit_functional[3]getitem_4: "f32[3]" = _native_batch_norm_legit_functional[4]return (getitem_3, getitem_4, add, getitem)
可以看到，我们在保持 IR 中 conv2d 算子不变的同时，分解了其余部分。现在该 IR 已成为一个功能型中间表示，仅保留 conv2d 以外的核心 aten 算子。
通过直接注册您选择的分解行为，您可以实现更深入的定制化。您还可以通过直接注册自定义分解行为来获得更灵活的定制能力。
# Lower to core aten inference IR, but customize conv2d
decomp_table = torch.export.default_decompositions()def my_awesome_custom_conv2d_function(x, weight, bias, stride=[1, 1], padding=[0, 0], dilation=[1, 1], groups=1):return 2 * torch.ops.aten.convolution(x, weight, bias, stride, padding, dilation, False, [0, 0], groups)decomp_table[torch.ops.aten.conv2d.default] = my_awesome_conv2d_function
ep_for_inference = ep_for_training.run_decompositions(decomp_table)print(ep_for_inference)
ExportedProgram:class GraphModule(torch.nn.Module):def forward(self, p_conv_weight: "f32[3, 1, 1, 1]", p_conv_bias: "f32[3]", p_bn_weight: "f32[3]", p_bn_bias: "f32[3]", b_bn_running_mean: "f32[3]", b_bn_running_var: "f32[3]", b_bn_num_batches_tracked: "i64[]", x: "f32[1, 1, 3, 3]"):convolution: "f32[1, 3, 3, 3]" = torch.ops.aten.convolution.default(x, p_conv_weight, p_conv_bias, [1, 1], [0, 0], [1, 1], False, [0, 0], 1)mul: "f32[1, 3, 3, 3]" = torch.ops.aten.mul.Tensor(convolution, 2)add: "i64[]" = torch.ops.aten.add.Tensor(b_bn_num_batches_tracked, 1)_native_batch_norm_legit_functional = torch.ops.aten._native_batch_norm_legit_functional.default(mul, p_bn_weight, p_bn_bias, b_bn_running_mean, b_bn_running_var, True, 0.1, 1e-05)getitem: "f32[1, 3, 3, 3]" = _native_batch_norm_legit_functional[0]getitem_3: "f32[3]" = _native_batch_norm_legit_functional[3]getitem_4: "f32[3]" = _native_batch_norm_legit_functional[4];return (getitem_3, getitem_4, add, getitem)表达动态性
默认情况下，torch.export 会假设所有输入形状都是静态的来追踪程序，并将导出的程序特化到这些维度。然而，某些维度（例如批次维度）可以是动态的，每次运行都可能变化。必须通过使用 torch.export.Dim() API 创建这些维度，并通过 dynamic_shapes 参数将它们传递给 torch.export.export() 来指定这些维度。示例如下：
import torch
from torch.export import Dim, exportclass M(torch.nn.Module):def __init__(self):super().__init__()self.branch1 = torch.nn.Sequential(torch.nn.Linear(64, 32), torch.nn.ReLU())self.branch2 = torch.nn.Sequential(torch.nn.Linear(128, 64), torch.nn.ReLU())self.buffer = torch.ones(32)def forward(self, x1, x2):out1 = self.branch1(x1)out2 = self.branch2(x2)return (out1 + self.buffer, out2)example_args = (torch.randn(32, 64), torch.randn(32, 128))# Create a dynamic batch size
batch = Dim("batch")
# Specify that the first dimension of each input is that batch size
dynamic_shapes = {"x1": {0: batch}, "x2": {0: batch}}exported_program: torch.export.ExportedProgram = export(M(), args=example_args, dynamic_shapes=dynamic_shapes
)
print(exported_program)ExportedProgram:
class GraphModule(torch.nn.Module):def forward(self, p_branch1_0_weight: "f32[32, 64]", p_branch1_0_bias: "f32[32]", p_branch2_0_weight: "f32[64, 128]", p_branch2_0_bias: "f32[64]", c_buffer: "f32[32]", x1: "f32[s0, 64]", x2: "f32[s0, 128]"):# code: out1 = self.branch1(x1)linear: "f32[s0, 32]" = torch.ops.aten.linear.default(x1, p_branch1_0_weight, p_branch1_0_bias)relu: "f32[s0, 32]" = torch.ops.aten.relu.default(linear)# code: out2 = self.branch2(x2)linear_1: "f32[s0, 64]" = torch.ops.aten.linear.default(x2, p_branch2_0_weight, p_branch2_0_bias)relu_1: "f32[s0, 64]" = torch.ops.aten.relu.default(linear_1)# code: return (out1 + self.buffer, out2)add: "f32[s0, 32]" = torch.ops.aten.add.Tensor(relu, c_buffer)return (add, relu_1)Range constraints: {s0: VR[0, int_oo]}
需要注意的几点补充事项：
- 通过
torch.export.Dim()API 和dynamic_shapes参数，我们指定了每个输入的第一个维度为动态维度。观察输入x1和x2，它们的符号形状分别为 (s0, 64) 和 (s0, 128)，而非我们传入的示例输入中 (32, 64) 和 (32, 128) 形状的张量。
s0 是一个符号，表示该维度可以接受一定范围内的数值。
exported_program.range_constraints描述了图中每个符号的取值范围。在本例中，我们看到s0的范围是 [0, int_oo]。由于某些技术原因（此处难以详述），系统假定这些取值不为 0 或 1。这并非程序错误，也不意味着导出的程序一定无法处理维度为 0 或 1 的情况。关于此话题的深入讨论，请参阅文档 0/1 特化问题。
我们还可以指定输入形状之间更具表现力的关系，例如：两个形状可能相差 1，某个形状可能是另一个的两倍，或者某个形状是偶数。示例如下：
class M(torch.nn.Module):def forward(self, x, y):return x + y[1:]x, y = torch.randn(5), torch.randn(6)
dimx = torch.export.Dim("dimx", min=3, max=6)
dimy = dimx + 1exported_program = torch.export.export(M(), (x, y), dynamic_shapes=({0: dimx}, {0: dimy}), )
print(exported_program)ExportedProgram:
class GraphModule(torch.nn.Module):def forward(self, x: "f32[s0]", y: "f32[s0 + 1]"):# code: return x + y[1:]slice_1: "f32[s0]" = torch.ops.aten.slice.Tensor(y, 0, 1, 9223372036854775807)add: "f32[s0]" = torch.ops.aten.add.Tensor(x, slice_1)return (add,)Range constraints: {s0: VR[3, 6], s0 + 1: VR[4, 7]}
需要注意以下几点：
- 当为第一个输入指定
{0: dimx}时，可以看到第一个输入的结果形状变为动态的[s0]。接着为第二个输入指定{0: dimy}时，第二个输入的结果形状也变为动态。但由于我们定义了dimy = dimx + 1，y的形状并未引入新符号，而是沿用x中的符号s0来表示。可以看到dimy = dimx + 1的关系通过s0 + 1体现。 - 观察范围约束条件，
s0的初始范围是 [3, 6]，而s0 + 1的解算范围是 [4, 7]。
序列化
要保存 ExportedProgram，用户可以使用 torch.export.save() 和 torch.export.load() API。通常建议使用 .pt2 文件扩展名来保存 ExportedProgram。
示例：
import torch
import ioclass MyModule(torch.nn.Module):def forward(self, x):return x + 10exported_program = torch.export.export(MyModule(), torch.randn(5))torch.export.save(exported_program, 'exported_program.pt2')saved_exported_program = torch.export.load('exported_program.pt2')
特化机制
理解torch.export行为的一个核心概念在于区分静态值与动态值。
动态值指每次运行可能发生变化的值。这类值的行为类似于Python函数的常规参数——你可以为同一参数传递不同值，并期望函数能正确执行。张量的数据就被视为动态值。
静态值则是在导出时固定，且在导出程序多次执行间保持不变的值。当跟踪过程中遇到静态值时，导出器会将其视为常量并硬编码到计算图中。
当执行某个操作（例如x + y）且所有输入均为静态值时，该操作的输出会直接硬编码到计算图中，该操作将不会显式出现（即被常量折叠优化）。
当一个值被硬编码到计算图中时，我们称该计算图已针对该值进行了特化。
以下类型的值属于静态值：
输入张量形状
默认情况下，torch.export 会根据输入张量的具体形状进行程序追踪，除非通过 dynamic_shapes 参数将某个维度指定为动态。这意味着如果存在依赖形状的控制流，torch.export 将根据给定示例输入所触发的分支进行特化处理。例如：
import torch
from torch.export import exportclass Mod(torch.nn.Module):def forward(self, x):if x.shape[0] > 5:return x + 1else:return x - 1example_inputs = (torch.rand(10, 2),)
exported_program = export(Mod(), example_inputs)
print(exported_program)
ExportedProgram:
class GraphModule(torch.nn.Module):def forward(self, x: "f32[10, 2]"):# code: return x + 1add: "f32[10, 2]" = torch.ops.aten.add.Tensor(x, 1)return (add,)
条件判断 (x.shape[0] > 5) 不会出现在 ExportedProgram 中，因为示例输入的静态形状是 (10, 2)。由于 torch.export 会针对输入的静态形状进行特化处理，else分支 (x - 1) 将永远不会被执行。若要在追踪图中保留基于张量形状的动态分支行为，需要使用 torch.export.Dim() 来指定输入张量的维度 (x.shape[0]) 为动态维度，同时需要重写源代码。
请注意：作为模块状态一部分的张量（如参数和缓冲区）始终具有静态形状。
Python 基本类型
torch.export 同样支持对 Python 基本类型的特化处理，例如 int、float、bool 和 str。不过它们也有对应的动态变体，如 SymInt、SymFloat 和 SymBool。
例如：
import torch
from torch.export import exportclass Mod(torch.nn.Module):def forward(self, x: torch.Tensor, const: int, times: int):for i in range(times):x = x + constreturn xexample_inputs = (torch.rand(2, 2), 1, 3)
exported_program = export(Mod(), example_inputs)
print(exported_program)
ExportedProgram:class GraphModule(torch.nn.Module):def forward(self, x: "f32[2, 2]", const, times):# code: x = x + constadd: "f32[2, 2]" = torch.ops.aten.add.Tensor(x, 1)add_1: "f32[2, 2]" = torch.ops.aten.add.Tensor(add, 1)add_2: "f32[2, 2]" = torch.ops.aten.add.Tensor(add_1, 1)return (add_2,)
由于整数是特化的，torch.ops.aten.add.Tensor 操作都会使用硬编码的常量 1 进行计算，而非变量 const。如果用户在运行时传入与导出时不同的 const 值（例如 2 而非 1），就会导致错误。
此外，for 循环中使用的 times 迭代器也通过 3 次重复的 torch.ops.aten.add.Tensor 调用被"内联"到计算图中，而输入参数 times 实际上从未被使用。
Python 容器
Python 容器（List、Dict、NamedTuple 等）被认为具有静态结构。
torch.export 的局限性
图中断问题
由于torch.export是一个从PyTorch程序中捕获计算图的一次性过程，它最终可能会遇到程序无法追踪的部分，因为几乎不可能支持追踪所有PyTorch和Python特性。在torch.compile的情况下，遇到不支持的操作会导致"图中断"，该操作将通过默认的Python解释器执行。相比之下，torch.export会要求用户提供额外信息或重写部分代码使其可追踪。由于追踪基于TorchDynamo（在Python字节码级别进行评估），与之前的追踪框架相比，所需的代码重写将显著减少。
当遇到图中断时，ExportDB是了解支持/不支持程序类型以及如何重写程序使其可追踪的绝佳资源。
解决图中断问题的一个选项是使用非严格导出模式。
数据/形状相关的控制流
当形状未被特化时，在数据依赖的控制流（如if x.shape[0] > 2）中也可能遇到图中断问题。这是因为追踪编译器无法在不生成组合爆炸路径数量的代码的情况下处理这种情况。此时，用户需要使用特定的控制流运算符重写代码。目前，我们支持使用torch.cond来表达类似if-else的控制流（更多功能即将推出！）。
算子缺少 Fake/Meta/Abstract 内核实现
在进行追踪时，所有算子都必须具备 FakeTensor 内核（也称为元内核或抽象实现）。该内核用于推导该算子的输入/输出形状。
更多详情请参阅 torch.library.register_fake()。
若您的模型使用了尚未实现 FakeTensor 内核的 ATen 算子，请提交问题报告。
扩展阅读
面向导出用户的附加链接
- torch.export 编程模型
- torch.export IR 规范
- 在ATen IR上编写图转换
- 中间表示层
- 导出数据库
- 控制流 - Cond
PyTorch开发者深度指南
- Dynamo概览
- Dynamo深度解析
- 动态形状
- 伪张量
API 参考
torch.export.export(mod, args, kwargs=None, *, dynamic_shapes=None, strict=True, preserve_module_call_signature=())
export() 接收任意 nn.Module 及示例输入，并以预先编译（AOT）的方式生成一个仅表示函数张量计算过程的追踪图。该追踪图具有以下特性：(1) 生成符合功能化 ATen 算子集的标准化算子（以及用户指定的任何自定义算子）；(2) 消除了所有 Python 控制流和数据结构（特定例外情况除外）；(3) 记录了所需的形状约束集合，以证明这种标准化和控制流消除对未来输入是可靠的。
可靠性保证
在追踪过程中，export() 会记录用户程序及底层 PyTorch 算子内核对形状相关的假设。只有当这些假设成立时，输出的 ExportedProgram 才被视为有效。
追踪过程会对输入张量的形状（而非数值）作出假设。这些假设必须在图捕获阶段完成验证，export() 才能成功执行。具体而言：
- 对输入张量静态形状的假设会自动验证，无需额外操作。
- 对输入张量动态形状的假设需要通过
Dim()API 显式声明动态维度，并通过dynamic_shapes参数将其与示例输入关联。
若任何假设无法验证，将触发致命错误。此时，错误信息会包含验证假设所需的规范修改建议。例如，export() 可能针对输入 x 关联形状中出现的动态维度 dim0_x（假设先前定义为 Dim("dim0_x")）给出如下修正建议：
dim = Dim("dim0_x", max=5)
这个示例意味着生成的代码要求输入x的第0维必须小于或等于5才有效。您可以检查针对动态维度定义的建议修正方案，然后原封不动地复制到代码中，而无需修改export()调用中的dynamic_shapes参数。
参数说明
-
mod (Module)– 我们将追踪该模块的forward方法。 -
args (tuple[Any, ...])– 示例位置输入。 -
kwargs (Optional[dict[str, Any]])– 可选的示例关键字输入。 -
dynamic_shapes (Optional[Union[dict[str, Any], tuple[Any], list[Any]]])– 可选参数，其类型应为以下之一：- 从
f的参数名到其动态形状规范的字典； - 按原始顺序为每个输入指定动态形状规范的元组。
若要对关键字参数指定动态性，需按照原始函数签名中定义的顺序传递它们。
张量参数的动态形状可通过以下方式指定：
(1) 从动态维度索引到Dim()类型的字典（静态维度索引无需包含在此字典中，若包含则应映射为None）；
(2)Dim()类型或None组成的元组/列表，其中Dim()类型对应动态维度，静态维度用None表示。
对于字典或张量元组/列表类型的参数，可通过递归使用包含规范的映射或序列来指定。 - 从
-
strict (bool)– 启用时（默认），导出函数将通过TorchDynamo追踪程序以确保结果图的正确性。否则，导出程序不会验证图中隐含的假设，可能导致原始模型与导出模型的行为差异。这在用户需要绕过追踪器错误或逐步启用模型安全性时很有用。注意这不会导致最终IR规范不同，无论此处传递何值，模型都将以相同方式序列化。
警告：此选项为实验性功能，使用时需自行承担风险。
返回值：返回包含被追踪可调用对象的ExportedProgram。
返回类型：ExportedProgram
可接受的输入/输出类型
args和kwargs的输入及输出可接受类型包括：
- 基本类型：
torch.Tensor、int、float、bool和str - 数据类（需先通过调用
register_dataclass()注册） - 包含上述所有类型的（嵌套）数据结构：
dict、list、tuple、namedtuple和OrderedDict
torch.export.save(ep, f, *, extra_files=None, opset_version=None, pickle_protocol=2)
警告： 当前功能处于积极开发阶段，保存的文件可能无法在 PyTorch 新版本中使用。
将 ExportedProgram 保存到类文件对象中，后续可通过 Python API torch.export.load 加载。
参数说明
ep (ExportedProgram)– 待保存的导出程序f (str | os.PathLike[str] | *IO[bytes ])– 需实现 write 和 flush 方法的类文件对象，或包含文件名的字符串extra_files (Optional[Dict[str, Any]])– 文件名到内容的映射，这些内容将作为文件的一部分存储opset_version (Optional[Dict[str, int ]])– 操作集名称到其版本的映射pickle_protocol ( int )– 可指定以覆盖默认协议
示例：
import torch
import ioclass MyModule(torch.nn.Module):def forward(self, x):return x + 10ep = torch.export.export(MyModule(), (torch.randn(5),))# Save to file
torch.export.save(ep, 'exported_program.pt2')# Save to io.BytesIO buffer
buffer = io.BytesIO()
torch.export.save(ep, buffer)# Save with extra files
extra_files = {'foo.txt': b'bar'.decode('utf-8')}
torch.export.save(ep, 'exported_program.pt2', extra_files=extra_files)
torch.export.load(f, *, extra_files=None, expected_opset_version=None)
警告：当前功能处于积极开发阶段，保存的文件可能无法在较新版本的PyTorch中使用。
加载先前通过 torch.export.save 保存的 ExportedProgram。
参数
f (str | os.PathLike[str] | *IO[bytes ])– 文件类对象（需实现write和flush方法）或包含文件名的字符串。extra_files (Optional[Dict[str, Any]])– 此映射中提供的额外文件名将被加载，其内容会存储到给定的映射中。expected_opset_version (Optional[Dict[str, int ]])– 操作集名称到预期版本号的映射
返回值：一个 ExportedProgram 对象
返回类型：ExportedProgram
示例：
import torch
import io# Load ExportedProgram from file
ep = torch.export.load('exported_program.pt2')# Load ExportedProgram from io.BytesIO object
with open('exported_program.pt2', 'rb') as f:buffer = io.BytesIO(f.read())
buffer.seek(0)
ep = torch.export.load(buffer)# Load with extra files.
extra_files = {'foo.txt': ''} # values will be replaced with data
ep = torch.export.load('exported_program.pt2', extra_files=extra_files)
print(extra_files['foo.txt'])
print(ep(torch.randn(5)))
torch.export.register_dataclass(cls, *, serialized_type_name=None)
将数据类注册为 torch.export.export() 的有效输入/输出类型。
参数
cls (type[Any])- 要注册的数据类类型serialized_type_name (Optional[str])- 数据类的序列化名称。这是*this (当需要序列化包含数据类的 pytree TreeSpec 时为必填项)
示例：
import torchfrom dataclasses import dataclass@dataclass
class InputDataClass:feature: torch.Tensorbias: int@dataclass
class OutputDataClass:res: torch.Tensor
torch.export.register_dataclass(InputDataClass)
torch.export.register_dataclass(OutputDataClass)class Mod(torch.nn.Module):def forward(self, x: InputDataClass) -OutputDataClass:res = x.feature + x.biasreturn OutputDataClass(res=res)ep = torch.export.export(Mod(), (InputDataClass(torch.ones(2, 2), 1), ))
print(ep)torch.export.dynamic_shapes.Dim(name, *, min=None, max=None)
Dim() 构造了一个类似于带范围命名符号整数的类型。
该类型可用于描述动态张量维度的多种可能取值。
注意：同一张量的不同动态维度，或不同张量的动态维度，都可以用同一类型来描述。
参数说明
name (str)- 用于调试的可读名称min (Optional[int])- 符号的最小可能值（包含）max (Optional[int])- 符号的最大可能值（包含）
返回值
返回一个可用于张量动态形状规范的类型。
torch.export.exported_program.default_decompositions()
这是默认的分解表，包含了所有 ATEN 算子向核心 aten 算子集的分解映射。请将此 API 与 run_decompositions() 配合使用。
返回类型 : CustomDecompTable
torch.export.dims(*names, min=None, max=None)
用于创建多个 Dim() 类型的工具函数。
返回值：返回由 Dim() 类型组成的元组。
返回类型：tuple[torch.export.dynamic_shapes._Dim, …]
class torch.export.dynamic_shapes.ShapesCollection
动态形状构建器。
用于为输入张量分配动态形状规格。
当args()是嵌套输入结构时特别有用，相比在dynamic_shapes()规范中复制args()的结构，直接索引输入张量会更方便。
示例：
args = ({"x": tensor_x, "others": [tensor_y, tensor_z]})dim = torch.export.Dim(...)
dynamic_shapes = torch.export.ShapesCollection()
dynamic_shapes[tensor_x] = (dim, dim + 1, 8)
dynamic_shapes[tensor_y] = {0: dim * 2}
# This is equivalent to the following (now auto-generated):
# dynamic_shapes = {"x": (dim, dim + 1, 8), "others": [{0: dim * 2}, None]}torch.export(..., args, dynamic_shapes=dynamic_shapes)dynamic_shapes(m, args, kwargs=None)
生成与 args() 和 kwargs() 对应的 dynamic_shapes() pytree 结构。
torch.export.dynamic_shapes.refine_dynamic_shapes_from_suggested_fixes(msg, dynamic_shapes)
使用 dynamic_shapes() 导出时，如果规格与模型追踪推断出的约束条件不匹配，可能会因 ConstraintViolation 错误导致导出失败。错误信息通常会提供修正建议——即对 dynamic_shapes() 进行哪些修改才能成功导出。
ConstraintViolation 错误示例信息：
Suggested fixes:dim = Dim('dim', min=3, max=6) # this just refines the dim's rangedim = 4 # this specializes to a constantdy = dx + 1 # dy was specified as an independent dim, but is actually tied to dx with this relation这是一个辅助函数，它接收 ConstraintViolation 错误信息和原始的 dynamic_shapes() 规格参数，返回一个整合了建议修复方案的新 dynamic_shapes() 规格。
使用示例：
try:ep = export(mod, args, dynamic_shapes=dynamic_shapes)
except torch._dynamo.exc.UserError as exc:new_shapes = refine_dynamic_shapes_from_suggested_fixes(exc.msg, dynamic_shapes)ep = export(mod, args, dynamic_shapes=new_shapes)返回类型：Union[dict[str, Any], tuple[Any], list[Any]]
torch.export.Constraint
Union[_Constraint, _DerivedConstraint, _RelaxedConstraint]的别名
class torch.export.ExportedProgram(root, graph, graph_signature, state_dict, range_constraints, module_call_graph, example_inputs=None, constants=None, *, verifiers=None)
由 export() 导出的程序包。它包含：
- 一个表示张量计算的
torch.fx.Graph - 包含所有提升参数和缓冲区张量值的 state_dict
- 各种元数据
可以像调用原始可调用对象那样调用 ExportedProgram，其调用约定与 export() 追踪的版本相同。
要对计算图进行转换时，可通过 .module 属性访问 torch.fx.GraphModule，然后使用 FX 变换重写计算图。完成后，只需再次调用 export() 即可构建正确的 ExportedProgram。
module()
返回一个自包含的 GraphModule，其中内联了所有参数/缓冲区。
返回类型：Module
buffers()
返回一个遍历原始模块缓冲区的迭代器。
警告：此 API 是实验性的，且不向后兼容。
返回类型：Iterator[Tensor]
named_buffers()
返回一个遍历原始模块缓冲区的迭代器，同时生成缓冲区的名称和缓冲区本身。
警告：此 API 为实验性质，且不向后兼容。
返回类型：Iterator[tuple[str, torch.Tensor]]
parameters()
返回一个遍历原始模块参数的迭代器。
警告：此 API 为实验性质，且不向后兼容。
返回类型：迭代器 [Parameter]
named_parameters()
返回一个遍历原始模块参数的迭代器，同时生成参数名称和参数本身。
警告：此 API 为实验性质，且不向后兼容。
返回类型：Iterator[tuple[str, torch.nn.parameter.Parameter]]
run_decompositions(decomp_table=None, decompose_custom_triton_ops=False)
对导出的程序运行一系列分解操作，并返回一个新的导出程序。默认情况下，我们会运行 Core ATen 分解来获取 Core ATen Operator Set 中的运算符。
目前暂不支持分解联合图。
参数
decomp_table (Optional[dict[torch._ops.OperatorBase, Callable]])- 可选参数，用于指定 Aten 运算符的分解行为：(1) 如果为 None，则分解至核心 aten 分解；(2) 如果为空字典，则不分解任何运算符
返回类型：ExportedProgram
使用示例：
如果不想分解任何操作
ep = torch.export.export(model, ...)
ep = ep.run_decompositions(decomp_table={})
如果你想获取核心 aten 操作符集合（排除特定操作符），可以执行以下操作：
ep = torch.export.export(model, ...)
decomp_table = torch.export.default_decompositions()
decomp_table[your_op] = your_custom_decomp
ep = ep.run_decompositions(decomp_table=decomp_table)
class torch.export.ExportBackwardSignature(gradients_to_parameters: dict[str, str], gradients_to_user_inputs: dict[str, str], loss_output: str)
class torch.export.ExportGraphSignature(input_specs, output_specs)
ExportGraphSignature 定义了导出图的输入/输出签名，这是一个具有更强不变性保证的 fx.Graph。
导出图是纯函数式的，不会通过 getattr 节点访问图内的"状态"（如参数或缓冲区）。相反，export() 确保将参数、缓冲区和常量张量都作为输入从图中提取出来。同样，对缓冲区的任何修改也不会包含在图中，而是将修改后的缓冲区值建模为导出图的额外输出。
所有输入和输出的顺序如下：
Inputs = [*parameters_buffers_constant_tensors, *flattened_user_inputs]
Outputs = [*mutated_inputs, *flattened_user_outputs]
例如，如果导出以下模块：
class CustomModule(nn.Module):def __init__(self) -> None:super(CustomModule, self).__init__()# Define a parameterself.my_parameter = nn.Parameter(torch.tensor(2.0))# Define two buffersself.register_buffer('my_buffer1', torch.tensor(3.0))self.register_buffer('my_buffer2', torch.tensor(4.0))def forward(self, x1, x2):# Use the parameter, buffers, and both inputs in the forward methodoutput = (x1 + self.my_parameter) * self.my_buffer1 + x2 * self.my_buffer2# Mutate one of the buffers (e.g., increment it by 1)self.my_buffer2.add_(1.0) # In-place additionreturn output
生成的图表如下：
graph():%arg0_1 := placeholder[target=arg0_1]%arg1_1 := placeholder[target=arg1_1]%arg2_1 := placeholder[target=arg2_1]%arg3_1 := placeholder[target=arg3_1]%arg4_1 := placeholder[target=arg4_1]%add_tensor := call_function[target=torch.ops.aten.add.Tensor](args = (%arg3_1, %arg0_1), kwargs = {})%mul_tensor := call_function[target=torch.ops.aten.mul.Tensor](args = (%add_tensor, %arg1_1), kwargs = {})%mul_tensor_1 := call_function[target=torch.ops.aten.mul.Tensor](args = (%arg4_1, %arg2_1), kwargs = {})%add_tensor_1 := call_function[target=torch.ops.aten.add.Tensor](args = (%mul_tensor, %mul_tensor_1), kwargs = {})%add_tensor_2 := call_function[target=torch.ops.aten.add.Tensor](args = (%arg2_1, 1.0), kwargs = {})return (add_tensor_2, add_tensor_1)
生成的 ExportGraphSignature 将是：
ExportGraphSignature(input_specs=[InputSpec(kind=<InputKind.PARAMETER: 2>, arg=TensorArgument(name='arg0_1'), target='my_parameter'),InputSpec(kind=<InputKind.BUFFER: 3>, arg=TensorArgument(name='arg1_1'), target='my_buffer1'),InputSpec(kind=<InputKind.BUFFER: 3>, arg=TensorArgument(name='arg2_1'), target='my_buffer2'),InputSpec(kind=<InputKind.USER_INPUT: 1>, arg=TensorArgument(name='arg3_1'), target=None),InputSpec(kind=<InputKind.USER_INPUT: 1>, arg=TensorArgument(name='arg4_1'), target=None)],output_specs=[OutputSpec(kind=<OutputKind.BUFFER_MUTATION: 3>, arg=TensorArgument(name='add_2'), target='my_buffer2'),OutputSpec(kind=<OutputKind.USER_OUTPUT: 1>, arg=TensorArgument(name='add_1'), target=None)]
)
replace_all_uses(old, new)
将签名中所有旧名称的使用替换为新名称。
get_replace_hook(replace_inputs=False
class torch.export.graph_signature.CustomObjArgument(name: str, class_fqn: str, fake_val: Optional[torch._library.fake_class_registry.FakeScriptObject] = None)
class torch.export.unflatten.FlatArgsAdapter
根据输入参数 input_spec 进行调整，使其与 target_spec 对齐。
abstract adapt(target_spec, input_spec, input_args, metadata=None)
注意：此适配器可能会修改给定的 input_args_with_path。
返回类型：list[Any]
class torch.export.unflatten.InterpreterModule(graph, ty=None)
一个使用 torch.fx.Interpreter 而非 GraphModule 常规代码生成来执行的模块。这种方式能提供更清晰的堆栈跟踪信息，使执行调试更加便捷。
class torch.export.unflatten.InterpreterModuleDispatcher(attrs, call_modules)
一个包含一系列对应模块调用序列的InterpreterModule的模块。每次对该模块的调用都会分派给下一个InterpreterModule，并在最后一个之后循环回第一个。
torch.export.unflatten.unflatten(module, flat_args_adapter=None)
展开一个ExportedProgram，生成与原始eager模块具有相同模块层级的模块。当你尝试将torch.export与其他期望模块层级（而非torch.export通常生成的扁平图）的系统结合使用时，这会很有用。
注意：展开后模块的args/kwargs不一定与eager模块匹配，因此直接进行模块替换（例如self.submod = new_mod）可能无效。如需替换模块，必须设置torch.export.export()的preserve_module_call_signature参数。
参数
module (ExportedProgram)– 待展开的ExportedProgramflat_args_adapter (Optional[FlatArgsAdapter])– 当输入TreeSpec与导出模块不匹配时，用于适配扁平参数
返回值
返回UnflattenedModule实例，其模块层级与导出前的原始eager模块相同。
返回类型：UnflattenedModule
torch.export.passes.move_to_device_pass(ep, location)
将导出的程序移动到指定设备。
参数
ep ([ExportedProgram](https://pytorch.org/docs/stable/data.html#torch.export.ExportedProgram "torch.export.ExportedProgram"))– 需要移动的导出程序。location (Union[torch.device ,* str, Dict[str,* str]])– 目标设备。
如果是字符串，会被解析为设备名称。
如果是字典，会被解析为从现有设备到目标设备的映射关系。
返回
移动后的导出程序。
返回类型：ExportedProgram
2025-05-10（六）