torch.export¶

警告

此功能是一个正在积极开发的原型，未来**将会**有**重大变更**。

概述¶

torch.export.export() 接受任意 Python 可调用对象（torch.nn.Module、函数或方法），并以预先编译 (AOT) 的方式生成一个跟踪图，该图仅表示该函数的张量计算，随后可以使用不同的输出执行或序列化。

import torch
from torch.export import export

class 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 + b

example_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 + b
            add: 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，它使用名为帧求值 API 的 CPython 功能来安全地跟踪 PyTorch 图。这提供了大大改进的图捕获体验，为了完全跟踪 PyTorch 代码，需要更少的重写。
AOT Autograd 提供了一个函数式 PyTorch 图，并确保该图被分解/降低到 ATen 运算符集。
Torch FX (torch.fx) 是图的底层表示，允许灵活的基于 Python 的转换。

现有框架¶

torch.compile() 也使用了与 torch.export 相同的 PT2 堆栈，但略有不同

JIT vs. AOT：torch.compile() 是一个 JIT 编译器，而 torch.export 是一个 AOT 编译器。JIT 编译器旨在用于部署，而 AOT 编译器不旨在用于生成部署之外的编译产物。
部分图捕获 vs. 完全图捕获：当 torch.compile() 遇到模型中无法跟踪的部分时，它会“图中断”并回退到在即时 Python 运行时中运行程序。相比之下，torch.export 旨在获得 PyTorch 模型的完整图表示，因此当遇到无法跟踪的内容时，它会报错。由于 torch.export 生成的完整图与任何 Python 功能或运行时分离，因此该图可以被保存、加载并在不同的环境和语言中运行。
可用性权衡：由于 torch.compile() 能够在遇到无法跟踪的内容时回退到 Python 运行时，因此它更加灵活。torch.export 则需要用户提供更多信息或重写其代码以使其可跟踪。

与 torch.fx.symbolic_trace() 相比，torch.export 使用 TorchDynamo 进行跟踪，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_。
参数（conv 的权重和偏置）被提升为图的输入，从而在图中没有 get_attr 节点，而这些节点以前存在于 torch.fx.symbolic_trace() 的结果中。
torch.export.ExportGraphSignature 模拟输入和输出签名，并指定哪些输入是参数。
记录了图中每个节点生成的张量的结果形状和 dtype。例如，convolution 节点将生成 dtype 为 torch.float32 和形状为 (1, 16, 256, 256) 的张量。

非严格导出¶

在 PyTorch 2.3 中，我们引入了一种新的跟踪模式，称为**非严格模式**。它仍在经历强化阶段，因此如果您遇到任何问题，请在 Github 上使用 “oncall: export” 标签提交。

在*非严格模式*中，我们使用 Python 解释器跟踪程序。您的代码将完全像在即时模式下一样执行；唯一的区别是所有 Tensor 对象将被 ProxyTensor 对象替换，ProxyTensor 对象将将其所有操作记录到图中。

在*严格*模式（目前是默认模式）中，我们首先使用 TorchDynamo（一种字节码分析引擎）跟踪程序。TorchDynamo 实际上不执行您的 Python 代码。相反，它符号化地分析它并基于结果构建图。这种分析使 torch.export 能够提供更强的安全保证，但并非所有 Python 代码都受支持。

一个可能希望使用非严格模式的情况示例是，如果您遇到一个不受支持的 TorchDynamo 功能，该功能可能不容易解决，并且您知道 python 代码不是计算所必需的。例如

import contextlib
import torch

class ContextManager():
    def __init__(self):
        self.count = 0
    def __enter__(self):
        self.count += 1
    def __exit__(self, exc_type, exc_value, traceback):
        self.count -= 1

class 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 中，我们生成最通用的 IR，其中包含所有 ATen 运算符（包括函数式和非函数式运算符），这些运算符可用于在即时 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,)

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.PARAMETER: 2>,
                arg=TensorArgument(name='p_bn_weight'),
                target='bn.weight',
                persistent=None
            ),
            InputSpec(
                kind=<InputKind.PARAMETER: 2>,
                arg=TensorArgument(name='p_bn_bias'),
                target='bn.bias',
                persistent=None
            ),
            InputSpec(
                kind=<InputKind.BUFFER: 3>,
                arg=TensorArgument(name='b_bn_running_mean'),
                target='bn.running_mean',
                persistent=True
            ),
            InputSpec(
                kind=<InputKind.BUFFER: 3>,
                arg=TensorArgument(name='b_bn_running_var'),
                target='bn.running_var',
                persistent=True
            ),
            InputSpec(
                kind=<InputKind.BUFFER: 3>,
                arg=TensorArgument(name='b_bn_num_batches_tracked'),
                target='bn.num_batches_tracked',
                persistent=True
            ),
            InputSpec(
                kind=<InputKind.USER_INPUT: 1>,
                arg=TensorArgument(name='x'),
                target=None,
                persistent=None
            )
        ],
        output_specs=[
            OutputSpec(
                kind=<OutputKind.USER_OUTPUT: 1>,
                arg=TensorArgument(name='batch_norm'),
                target=None
            )
        ]
    )
Range constraints: {}

从上面的输出中，您可以看到 export_for_training() 生成的 ExportedProgram 与 export() 生成的 ExportedProgram 几乎相同，除了图中的运算符。您可以看到我们以最通用的形式捕获了 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)

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.PARAMETER: 2>,
                arg=TensorArgument(name='p_bn_weight'),
                target='bn.weight',
                persistent=None
            ),
            InputSpec(
                kind=<InputKind.PARAMETER: 2>,
                arg=TensorArgument(name='p_bn_bias'),
                target='bn.bias',
                persistent=None
            ),
            InputSpec(
                kind=<InputKind.BUFFER: 3>,
                arg=TensorArgument(name='b_bn_running_mean'),
                target='bn.running_mean',
                persistent=True
            ),
            InputSpec(
                kind=<InputKind.BUFFER: 3>,
                arg=TensorArgument(name='b_bn_running_var'),
                target='bn.running_var',
                persistent=True
            ),
            InputSpec(
                kind=<InputKind.BUFFER: 3>,
                arg=TensorArgument(name='b_bn_num_batches_tracked'),
                target='bn.num_batches_tracked',
                persistent=True
            ),
            InputSpec(
                kind=<InputKind.USER_INPUT: 1>,
                arg=TensorArgument(name='x'),
                target=None,
                persistent=None
            )
        ],
        output_specs=[
            OutputSpec(
                kind=<OutputKind.BUFFER_MUTATION: 3>,
                arg=TensorArgument(name='getitem_3'),
                target='bn.running_mean'
            ),
            OutputSpec(
                kind=<OutputKind.BUFFER_MUTATION: 3>,
                arg=TensorArgument(name='getitem_4'),
                target='bn.running_var'
            ),
            OutputSpec(
                kind=<OutputKind.BUFFER_MUTATION: 3>,
                arg=TensorArgument(name='add'),
                target='bn.num_batches_tracked'
            ),
            OutputSpec(
                kind=<OutputKind.USER_OUTPUT: 1>,
                arg=TensorArgument(name='getitem'),
                target=None
            )
        ]
    )
Range constraints: {}

在这里您可以看到我们在 IR 中保留了 conv2d 操作，同时分解了其余部分。现在 IR 是一个函数式 IR，其中包含核心 aten 运算符，除了 conv2d。

您可以通过直接注册您选择的分解行为来执行更多自定义。

您可以通过直接注册自定义分解行为来执行更多自定义

# 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)

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.PARAMETER: 2>,
                arg=TensorArgument(name='p_bn_weight'),
                target='bn.weight',
                persistent=None
            ),
            InputSpec(
                kind=<InputKind.PARAMETER: 2>,
                arg=TensorArgument(name='p_bn_bias'),
                target='bn.bias',
                persistent=None
            ),
            InputSpec(
                kind=<InputKind.BUFFER: 3>,
                arg=TensorArgument(name='b_bn_running_mean'),
                target='bn.running_mean',
                persistent=True
            ),
            InputSpec(
                kind=<InputKind.BUFFER: 3>,
                arg=TensorArgument(name='b_bn_running_var'),
                target='bn.running_var',
                persistent=True
            ),
            InputSpec(
                kind=<InputKind.BUFFER: 3>,
                arg=TensorArgument(name='b_bn_num_batches_tracked'),
                target='bn.num_batches_tracked',
                persistent=True
            ),
            InputSpec(
                kind=<InputKind.USER_INPUT: 1>,
                arg=TensorArgument(name='x'),
                target=None,
                persistent=None
            )
        ],
        output_specs=[
            OutputSpec(
                kind=<OutputKind.BUFFER_MUTATION: 3>,
                arg=TensorArgument(name='getitem_3'),
                target='bn.running_mean'
            ),
            OutputSpec(
                kind=<OutputKind.BUFFER_MUTATION: 3>,
                arg=TensorArgument(name='getitem_4'),
                target='bn.running_var'
            ),
            OutputSpec(
                kind=<OutputKind.BUFFER_MUTATION: 3>,
                arg=TensorArgument(name='add'),
                target='bn.num_batches_tracked'
            ),
            OutputSpec(
                kind=<OutputKind.USER_OUTPUT: 1>,
                arg=TensorArgument(name='getitem'),
                target=None
            )
        ]
)
Range constraints: {}

表达动态性¶

默认情况下，torch.export 将跟踪程序，假设所有输入形状都是**静态**的，并将导出的程序专门用于这些维度。但是，某些维度（例如批次维度）可以是动态的，并且每次运行都可能不同。必须使用 torch.export.Dim() API 创建此类维度，并通过 dynamic_shapes 参数将其传递到 torch.export.export() 中。一个示例

import torch
from torch.export import Dim, export

class 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)

Graph signature:
    ExportGraphSignature(
        input_specs=[
            InputSpec(
                kind=<InputKind.PARAMETER: 2>,
                arg=TensorArgument(name='p_branch1_0_weight'),
                target='branch1.0.weight',
                persistent=None
            ),
            InputSpec(
                kind=<InputKind.PARAMETER: 2>,
                arg=TensorArgument(name='p_branch1_0_bias'),
                target='branch1.0.bias',
                persistent=None
            ),
            InputSpec(
                kind=<InputKind.PARAMETER: 2>,
                arg=TensorArgument(name='p_branch2_0_weight'),
                target='branch2.0.weight',
                persistent=None
            ),
            InputSpec(
                kind=<InputKind.PARAMETER: 2>,
                arg=TensorArgument(name='p_branch2_0_bias'),
                target='branch2.0.bias',
                persistent=None
            ),
            InputSpec(
                kind=<InputKind.CONSTANT_TENSOR: 4>,
                arg=TensorArgument(name='c_buffer'),
                target='buffer',
                persistent=True
            ),
            InputSpec(
                kind=<InputKind.USER_INPUT: 1>,
                arg=TensorArgument(name='x1'),
                target=None,
                persistent=None
            ),
            InputSpec(
                kind=<InputKind.USER_INPUT: 1>,
                arg=TensorArgument(name='x2'),
                target=None,
                persistent=None
            )
        ],
        output_specs=[
            OutputSpec(
                kind=<OutputKind.USER_OUTPUT: 1>,
                arg=TensorArgument(name='add'),
                target=None
            ),
            OutputSpec(
                kind=<OutputKind.USER_OUTPUT: 1>,
                arg=TensorArgument(name='relu_1'),
                target=None
            )
        ]
    )
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 + 1

exported_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,)

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: {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 io

class MyModule(torch.nn.Module):
    def forward(self, x):
        return x + 10

exported_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 指定维度为动态。这意味着如果存在依赖于形状的控制流，torch.export 将专门针对使用给定示例输入所采用的分支。例如

import torch
from torch.export import export

class Mod(torch.nn.Module):
    def forward(self, x):
        if x.shape[0] > 5:
            return x + 1
        else:
            return x - 1

example_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 + 1
        add: "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 export

class Mod(torch.nn.Module):
    def forward(self, x: torch.Tensor, const: int, times: int):
        for i in range(times):
            x = x + const
        return x

example_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 + const
            add: "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 传递与导出时使用的值 1 不同的值（例如 2），则会导致错误。此外，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，TorchDynamo 在 Python 字节码级别进行求值，因此与以前的跟踪框架相比，所需的重写将大大减少。

当遇到图中断时，ExportDB 是一个很好的资源，可以了解支持和不支持的程序类型，以及重写程序以使其可跟踪的方法。

解决此图中断问题的一种方法是使用非严格导出

数据/形状依赖的控制流¶

当形状未被专门化时，在数据依赖的控制流 (if x.shape[0] > 2) 上也可能遇到图中断，因为跟踪编译器不可能在不为组合爆炸数量的路径生成代码的情况下处理这种情况。在这种情况下，用户将需要使用特殊的控制流运算符重写他们的代码。目前，我们支持 torch.cond 来表达类似 if-else 的控制流（更多即将推出！）。

运算符缺少 Fake/Meta/Abstract Kernels¶

跟踪时，所有运算符都需要 FakeTensor 内核（又名元内核，抽象实现）。这用于推断此运算符的输入/输出形状。

请参阅 torch.library.register_fake() 了解更多详细信息。

在不幸的情况下，您的模型使用的 ATen 运算符还没有 FakeTensor 内核实现，请提交 issue。

API 参考¶

torch.export.export(mod, args, kwargs=None, *, dynamic_shapes=None, strict=True, preserve_module_call_signature=())[source][source]¶

export() 接受任何 nn.Module 以及示例输入，并以预先编译 (AOT) 的方式生成一个跟踪图，该图仅表示函数的 Tensor 计算，随后可以使用不同的输入执行或序列化。跟踪图 (1) 在功能性 ATen 运算符集中生成规范化的运算符（以及任何用户指定的自定义运算符），(2) 消除了所有 Python 控制流和数据结构（某些例外情况除外），以及 (3) 记录了证明这种规范化和控制流消除对于未来输入是合理所需的形状约束集。

健全性保证

在跟踪期间，export() 会记录用户程序和底层 PyTorch 运算符内核所做的形状相关假设。只有当这些假设成立时，输出 ExportedProgram 才被认为是有效的。

跟踪会对输入张量的形状（而非值）做出假设。必须在图捕获时验证这些假设，export() 才能成功。具体来说，

关于输入张量静态形状的假设会自动验证，无需额外操作。
关于输入张量动态形状的假设需要通过使用 Dim() API 构建动态维度，并通过 dynamic_shapes 参数将它们与示例输入关联来显式指定。

如果任何假设无法验证，将引发致命错误。发生这种情况时，错误消息将包含验证假设所需的规范建议修复。例如，export() 可能会建议对动态维度 dim0_x 的定义进行以下修复，例如出现在与输入 x 关联的形状中，该形状之前定义为 Dim("dim0_x")

dim = Dim("dim0_x", max=5)

此示例意味着生成的代码需要输入 x 的维度 0 小于或等于 5 才能有效。您可以检查对动态维度定义的建议修复，然后逐字复制到您的代码中，而无需更改 dynamic_shapes 参数到您的 export() 调用。

参数

mod (Module) – 我们将跟踪此模块的 forward 方法。
args (Tuple[Any, ...]) – 示例位置输入。
kwargs (Optional[Dict[str, Any]]) – 可选的示例关键字输入。
dynamic_shapes (Optional[Union[Dict[str, Any], Tuple[Any], List[Any]]]) –
一个可选参数，其类型应为：1) 从 f 的参数名称到其动态形状规范的字典，2) 一个元组，用于指定原始顺序中每个输入的动态形状规范。如果您要指定关键字参数的动态性，则需要按照原始函数签名中定义的顺序传递它们。

张量参数的动态形状可以指定为 (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)[source][source]¶

警告

正在积极开发中，保存的文件可能无法在较新版本的 PyTorch 中使用。

将 ExportedProgram 保存到类似文件的对象。然后可以使用 Python API torch.export.load 加载它。

参数

ep (ExportedProgram) – 要保存的导出程序。
f (Union[str, os.PathLike, io.BytesIO) – 类似文件的对象（必须实现写入和刷新）或包含文件名的字符串。
extra_files (Optional[Dict[str, Any]]) – 从文件名到内容的映射，这些内容将作为 f 的一部分存储。
opset_version (Optional[Dict[str, int]]) – opset 名称到此 opset 版本的映射

示例

import torch
import io

class MyModule(torch.nn.Module):
    def forward(self, x):
        return x + 10

ep = 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)[source][source]¶

警告

正在积极开发中，保存的文件可能无法在较新版本的 PyTorch 中使用。

加载先前使用 torch.export.save 保存的 ExportedProgram。

参数

ep (ExportedProgram) – 要保存的导出程序。
f (Union[str, os.PathLike, io.BytesIO) – 类似文件的对象（必须实现写入和刷新）或包含文件名的字符串。
extra_files (Optional[Dict[str, Any]]) – 此映射中给出的额外文件名将被加载，其内容将存储在提供的映射中。
expected_opset_version (Optional[Dict[str, int]]) – opset 名称到预期 opset 版本的映射

返回

一个 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)[source][source]¶

将数据类注册为 torch.export.export() 的有效输入/输出类型。

参数

cls (Type[Any]) – 要注册的数据类类型
serialized_type_name (Optional[str]) – 数据类的序列化名称。这是
this (如果您要序列化包含 pytree TreeSpec，则为必需) –
dataclass. –

示例

import torch
from dataclasses import dataclass

@dataclass
class InputDataClass:
    feature: torch.Tensor
    bias: 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.bias
        return 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)[source][source]¶

Dim() 构建一个类似于具有范围的命名符号整数的类型。它可以用于描述动态张量维度的多个可能值。请注意，同一张量的不同动态维度或不同张量的动态维度可以用相同的类型描述。

参数

name (str) – 用于调试的人类可读名称。
min (Optional[int]) – 给定符号的最小可能值（包括）
max (Optional[int]) – 给定符号的最大可能值（包括）

返回

一种可用于张量的动态形状规范中的类型。

torch.export.exported_program.default_decompositions()[source][source]¶

这是默认的分解表，其中包含将所有 ATEN 运算符分解为核心 aten opset 的分解。将此 API 与 run_decompositions() 一起使用

返回类型: CustomDecompTable

torch.export.dims(*names, min=None, max=None)[source][source]¶

用于创建多个 Dim() 类型的实用工具。

class torch.export.dynamic_shapes.ShapesCollection[source][source]¶

dynamic_shapes 的构建器。用于将动态形状规范分配给输入中出现的张量。

示例：

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} # 这等效于以下内容（现在自动生成）： # 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)[source][source]¶: 生成 dynamic_shapes。

torch.export.dynamic_shapes.refine_dynamic_shapes_from_suggested_fixes(msg, dynamic_shapes)[source][source]¶

用于处理 export 的动态形状建议修复和/或自动动态形状。根据给定的 ConstraintViolation 错误消息和原始动态形状，优化给定的动态形状规范。

在大多数情况下，行为很简单 - 即，对于专门化或优化 Dim 范围的建议修复，或建议派生关系的修复，新的动态形状规范将按此方式更新。

例如，建议修复

dim = Dim(‘dim’, min=3, max=6) -> 这只是优化 dim 的范围 dim = 4 -> 这专门化为常数 dy = dx + 1 -> dy 被指定为独立的 dim，但实际上通过此关系与 dx 关联

但是，与派生 dim 关联的建议修复可能更复杂。例如，如果为根 dim 提供了建议修复，则新的派生 dim 值将根据根进行评估。

例如，dx = Dim(‘dx’) dy = dx + 2 dynamic_shapes = {“x”: (dx,), “y”: (dy,)}

建议修复

dx = 4 # 专门化将导致 dy 也专门化 = 6 dx = Dim(‘dx’, max=6) # dy 现在具有 max = 8

派生 dims 建议修复也可以用于表达可除性约束。这涉及创建与特定输入形状无关的新根 dims。在这种情况下，根 dims 不会直接出现在新规范中，而是作为 dims 之一的根出现。

例如，建议修复

_dx = Dim(‘_dx’, max=1024) # 这不会出现在返回结果中，但 dx 会 dx = 4*_dx # dx 现在可以被 4 整除，最大值为 4096

返回类型: 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)[source][source]¶

来自 export() 的程序包。它包含一个 torch.fx.Graph，表示 Tensor 计算，一个 state_dict，包含所有提升的参数和缓冲区的张量值，以及各种元数据。

您可以像调用 export() 跟踪的原始可调用对象一样调用 ExportedProgram，并使用相同的调用约定。

要在图上执行转换，请使用 .module 属性访问 torch.fx.GraphModule。然后，您可以使用 FX 转换来重写图。之后，您可以简单地再次使用 export() 来构造正确的 ExportedProgram。

module()[source][source]¶

返回一个自包含的 GraphModule，其中包含所有内联的参数/缓冲区。

返回类型

Module

buffers()[source][source]¶

返回原始模块缓冲区的迭代器。

警告

此 API 是实验性的，不向后兼容。

返回类型: Iterator[Tensor]

named_buffers()[source][source]¶

返回原始模块缓冲区的迭代器，同时产生缓冲区的名称和缓冲区本身。

警告

此 API 是实验性的，不向后兼容。

返回类型: Iterator[Tuple[str, Tensor]]

parameters()[源代码][源代码]¶

返回原始模块参数的迭代器。

警告

此 API 是实验性的，不向后兼容。

返回类型: Iterator[Parameter]

named_parameters()[源代码][源代码]¶

返回原始模块参数的迭代器，同时产生参数的名称和参数本身。

警告

此 API 是实验性的，不向后兼容。

返回类型: Iterator[Tuple[str, Parameter]]

run_decompositions(decomp_table=None)[源代码][源代码]¶

在导出的程序上运行一组分解，并返回一个新的导出的程序。默认情况下，我们将运行 Core ATen 分解，以获取 Core ATen 运算符集中的运算符。

目前，我们不分解联合图。

参数: decomp_table (Optional[Dict[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 parameter
        self.my_parameter = nn.Parameter(torch.tensor(2.0))

        # Define two buffers
        self.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 method
        output = (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 addition

        return 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)
    ]
)

class torch.export.ModuleCallSignature(inputs: List[Union[torch.export.graph_signature.TensorArgument, torch.export.graph_signature.SymIntArgument, torch.export.graph_signature.SymFloatArgument, torch.export.graph_signature.SymBoolArgument, torch.export.graph_signature.ConstantArgument, torch.export.graph_signature.CustomObjArgument, torch.export.graph_signature.TokenArgument]], outputs: List[Union[torch.export.graph_signature.TensorArgument, torch.export.graph_signature.SymIntArgument, torch.export.graph_signature.SymFloatArgument, torch.export.graph_signature.SymBoolArgument, torch.export.graph_signature.ConstantArgument, torch.export.graph_signature.CustomObjArgument, torch.export.graph_signature.TokenArgument]], in_spec: torch.utils._pytree.TreeSpec, out_spec: torch.utils._pytree.TreeSpec, forward_arg_names: Optional[List[str]] = None)[源代码][源代码]¶

class torch.export.ModuleCallEntry(fqn: str, signature: Optional[torch.export.exported_program.ModuleCallSignature] = None)[源代码][源代码]¶

class torch.export.decomp_utils.CustomDecompTable[源代码][源代码]¶

这是一个自定义字典，专门用于处理导出中的 decomp_table。我们需要它的原因是，在新的体系中，您只能从分解表中删除一个操作以保留它。这对于自定义操作来说是有问题的，因为我们不知道自定义操作何时真正加载到调度器中。因此，我们需要记录自定义操作，直到我们真正需要实现它（即当我们运行分解过程时）。

我们保持的不变性是

所有 aten 分解都在初始化时加载
当用户从表中读取时，我们会实现所有操作，以使调度器更有可能拾取自定义操作。
如果是写入操作，我们不一定实现
我们在导出期间的最后一次加载，就在调用 run_decompositions() 之前

copy()[源代码][源代码]¶

返回类型: CustomDecompTable

items()[源代码][源代码]¶

keys()[源代码][源代码]¶

materialize()[源代码][源代码]¶

返回类型: Dict[OperatorBase, Callable]

pop(*args)[源代码][源代码]¶

update(other_dict)[源代码][源代码]¶

class torch.export.graph_signature.InputKind(value)[源代码][源代码]¶: 一个枚举。

class torch.export.graph_signature.InputSpec(kind: torch.export.graph_signature.InputKind, arg: Union[torch.export.graph_signature.TensorArgument, torch.export.graph_signature.SymIntArgument, torch.export.graph_signature.SymFloatArgument, torch.export.graph_signature.SymBoolArgument, torch.export.graph_signature.ConstantArgument, torch.export.graph_signature.CustomObjArgument, torch.export.graph_signature.TokenArgument], target: Optional[str], persistent: Optional[bool] = None)[源代码][源代码]¶

class torch.export.graph_signature.OutputKind(value)[源代码][源代码]¶: 一个枚举。

class torch.export.graph_signature.OutputSpec(kind: torch.export.graph_signature.OutputKind, arg: Union[torch.export.graph_signature.TensorArgument, torch.export.graph_signature.SymIntArgument, torch.export.graph_signature.SymFloatArgument, torch.export.graph_signature.SymBoolArgument, torch.export.graph_signature.ConstantArgument, torch.export.graph_signature.CustomObjArgument, torch.export.graph_signature.TokenArgument], target: Optional[str])[源代码][源代码]¶

class torch.export.graph_signature.SymIntArgument(name: str)[源代码][源代码]¶

class torch.export.graph_signature.SymBoolArgument(name: str)[源代码][源代码]¶

class torch.export.graph_signature.SymFloatArgument(name: str)[源代码][源代码]¶

class torch.export.graph_signature.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 parameter
        self.my_parameter = nn.Parameter(torch.tensor(2.0))

        # Define two buffers
        self.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 method
        output = (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 addition

        return 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)[源代码][源代码]¶

注意：此适配器可能会改变给定的 input_args_with_path。

返回类型: List[Any]

class torch.export.unflatten.InterpreterModule(graph)[源代码][源代码]¶

一个模块，它使用 torch.fx.Interpreter 执行，而不是 GraphModule 使用的常用代码生成。这提供了更好的堆栈跟踪信息，并使调试执行更容易。

class torch.export.unflatten.InterpreterModuleDispatcher(attrs, call_modules)[源代码][源代码]¶

一个模块，它携带一系列 InterpreterModule，对应于该模块的一系列调用。每次调用该模块都会调度到下一个 InterpreterModule，并在最后一个模块之后循环返回。

torch.export.unflatten.unflatten(module, flat_args_adapter=None)[source][source]¶

解展平一个 ExportedProgram，生成一个模块，该模块具有与原始 eager 模块相同的模块层级结构。如果您尝试将 torch.export 与另一个期望模块层级结构而不是 torch.export 通常生成的扁平图的系统一起使用，这将非常有用。

注意

解展平模块的 args/kwargs 不一定与 eager 模块匹配，因此进行模块交换（例如 self.submod = new_mod）不一定有效。如果需要交换模块，则需要设置 torch.export.export() 的 preserve_module_call_signature 参数。

参数

module (ExportedProgram) – 要解展平的 ExportedProgram。
flat_args_adapter (Optional[FlatArgsAdapter]) – 如果输入 TreeSpec 与导出的模块不匹配，则适配扁平参数。

返回

UnflattenedModule 的一个实例，它具有与导出前原始 eager 模块相同的模块层级结构。

返回类型

UnflattenedModule

torch.export.passes.move_to_device_pass(ep, location)[source][source]¶

将导出的程序移动到给定的设备。

参数

ep (ExportedProgram) – 要移动的导出的程序。
location (Union[torch.device, str, Dict[str, str]]) – 要将导出的程序移动到的设备。如果是字符串，则将其解释为设备名称。如果是字典，则将其解释为从现有设备到目标设备的映射

返回

移动后的导出的程序。

返回类型

ExportedProgram

torch.export¶

概述¶

现有框架¶

导出 PyTorch 模型¶

示例¶

非严格导出¶

训练和推理的导出¶

表达动态性¶

序列化¶

特殊化¶

输入张量形状¶

Python 原始类型¶

Python 容器¶

torch.export 的限制¶

图中断¶

数据/形状依赖的控制流¶

运算符缺少 Fake/Meta/Abstract Kernels¶

阅读更多¶

API 参考¶

文档

教程

资源