torch.export

警告

此功能目前处于积极开发阶段的原型，未来可能会有重大变化。

概述

torch.export.export() 接受任意 Python 可调用对象（如 torch.nn.Module、函数或方法），并以 Ahead-of-Time (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, arg0_1: f32[10, 10], arg1_1: f32[10, 10]):
            # code: a = torch.sin(x)
            sin: f32[10, 10] = torch.ops.aten.sin.default(arg0_1);

            # code: b = torch.cos(y)
            cos: f32[10, 10] = torch.ops.aten.cos.default(arg1_1);

            # code: return a + b
            add: f32[10, 10] = torch.ops.aten.add.Tensor(sin, cos);
            return (add,)

    Graph signature: ExportGraphSignature(
        parameters=[],
        buffers=[],
        user_inputs=['arg0_1', 'arg1_1'],
        user_outputs=['add'],
        inputs_to_parameters={},
        inputs_to_buffers={},
        buffers_to_mutate={},
        backward_signature=None,
        assertion_dep_token=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 的变换。

导出 PyTorch 模型

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, arg0_1: f32[16, 3, 3, 3], arg1_1: f32[16], arg2_1: f32[1, 3, 256, 256], arg3_1: f32[1, 16, 256, 256]):

            # code: a = self.conv(x)
            convolution: f32[1, 16, 256, 256] = torch.ops.aten.convolution.default(
                arg2_1, arg0_1, arg1_1, [1, 1], [1, 1], [1, 1], False, [0, 0], 1
            );

            # code: a.add_(constant)
            add: f32[1, 16, 256, 256] = torch.ops.aten.add.Tensor(convolution, arg3_1);

            # code: return self.maxpool(self.relu(a))
            relu: f32[1, 16, 256, 256] = torch.ops.aten.relu.default(add);
            max_pool2d_with_indices = torch.ops.aten.max_pool2d_with_indices.default(
                relu, [3, 3], [3, 3]
            );
            getitem: f32[1, 16, 85, 85] = max_pool2d_with_indices[0];
            return (getitem,)

    Graph signature: ExportGraphSignature(
        parameters=['L__self___conv.weight', 'L__self___conv.bias'],
        buffers=[],
        user_inputs=['arg2_1', 'arg3_1'],
        user_outputs=['getitem'],
        inputs_to_parameters={
            'arg0_1': 'L__self___conv.weight',
            'arg1_1': 'L__self___conv.bias',
        },
        inputs_to_buffers={},
        buffers_to_mutate={},
        backward_signature=None,
        assertion_dep_token=None,
    )
    Range constraints: {}

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

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, arg0_1: f32[32, 64], arg1_1: f32[32], arg2_1: f32[64, 128], arg3_1: f32[64], arg4_1: f32[32], arg5_1: f32[s0, 64], arg6_1: f32[s0, 128]):

            # code: out1 = self.branch1(x1)
            permute: f32[64, 32] = torch.ops.aten.permute.default(arg0_1, [1, 0]);
            addmm: f32[s0, 32] = torch.ops.aten.addmm.default(arg1_1, arg5_1, permute);
            relu: f32[s0, 32] = torch.ops.aten.relu.default(addmm);

            # code: out2 = self.branch2(x2)
            permute_1: f32[128, 64] = torch.ops.aten.permute.default(arg2_1, [1, 0]);
            addmm_1: f32[s0, 64] = torch.ops.aten.addmm.default(arg3_1, arg6_1, permute_1);
            relu_1: f32[s0, 64] = torch.ops.aten.relu.default(addmm_1);  addmm_1 = None

            # code: return (out1 + self.buffer, out2)
            add: f32[s0, 32] = torch.ops.aten.add.Tensor(relu, arg4_1);
            return (add, relu_1)

    Graph signature: ExportGraphSignature(
        parameters=[
            'branch1.0.weight',
            'branch1.0.bias',
            'branch2.0.weight',
            'branch2.0.bias',
        ],
        buffers=['L__self___buffer'],
        user_inputs=['arg5_1', 'arg6_1'],
        user_outputs=['add', 'relu_1'],
        inputs_to_parameters={
            'arg0_1': 'branch1.0.weight',
            'arg1_1': 'branch1.0.bias',
            'arg2_1': 'branch2.0.weight',
            'arg3_1': 'branch2.0.bias',
        },
        inputs_to_buffers={'arg4_1': 'L__self___buffer'},
        buffers_to_mutate={},
        backward_signature=None,
        assertion_dep_token=None,
    )
    Range constraints: {s0: RangeConstraint(min_val=2, max_val=9223372036854775806)}

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, arg0_1: "f32[s0]", arg1_1: "f32[s0 + 1]"):
        # code: return x + y[1:]
        slice_1: "f32[s0]" = torch.ops.aten.slice.Tensor(arg1_1, 0, 1, 9223372036854775807);  arg1_1 = None
        add: "f32[s0]" = torch.ops.aten.add.Tensor(arg0_1, slice_1);  arg0_1 = slice_1 = None
        return (add,)

Graph signature: ExportGraphSignature(
    input_specs=[
        InputSpec(kind=<InputKind.USER_INPUT: 1>, arg=TensorArgument(name='arg0_1'), target=None, persistent=None),
        InputSpec(kind=<InputKind.USER_INPUT: 1>, arg=TensorArgument(name='arg1_1'), target=None, persistent=None)
    ],
    output_specs=[
        OutputSpec(kind=<OutputKind.USER_OUTPUT: 1>, arg=TensorArgument(name='add'), target=None)]
)
Range constraints: {s0: ValueRanges(lower=3, upper=6, is_bool=False), s0 + 1: ValueRanges(lower=4, upper=7, is_bool=False)}

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

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, arg0_1: f32[10, 2]):
            add: f32[10, 2] = torch.ops.aten.add.Tensor(arg0_1, 1);
            return (add,)

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, arg0_1: f32[2, 2], arg1_1, arg2_1):
            add: f32[2, 2] = torch.ops.aten.add.Tensor(arg0_1, 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.export 的局限性

了解更多

API参考

torch.export.export(mod, args, kwargs=None, *, dynamic_shapes=None, strict=True, preserve_module_call_signature=())[源代码]

export() 接受任意的 Python 可调用对象（如 nn.Module、函数或方法）以及示例输入，并生成一个仅表示张量计算的追踪图，以编译时（AOT）的方式进行处理。该追踪图可以使用不同的输入执行或序列化。(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的第一个维度必须小于或等于5才有效。你可以查看并直接将建议的修复方案复制到你的代码中，而无需修改export()调用中的dynamic_shapes参数。

参数

• mod (Module) – 我们将追踪该模块的前向方法。

• args (Tuple[Any, ...]) – 位置参数示例。

• kwargs (Optional[Dict[str, Any]]) – 可选的示例关键字参数。

• dynamic_shapes (可选[联合类型[字典[str, Any], 元组[Any], 列表[Any]]]) –

  一个可选参数，其类型应为以下两种之一：1) 从 f 的参数名到它们的动态形状规范的字典；2) 按照原始顺序指定每个输入的动态形状规范的元组。如果你要在关键字参数上指定动态性，则需要按照原始函数签名中定义的顺序传递这些参数。

  张量参数的动态形状可以以以下两种方式之一来指定：(1) 一个从动态维度索引到 Dim() 类型的字典，其中静态维度索引可选包含，并且当它们存在时应映射为 None；或者 (2) 由 Dim() 类型或 None 组成的元组/列表，其中 Dim() 类型表示动态维度，静态维度则用 None 表示。对于字典或包含张量的元组/列表参数，则通过递归地使用嵌套规格的映射或序列来指定。

• strict (bool) – 当启用时（默认情况下），导出函数将通过 TorchDynamo 追踪程序，确保生成的图的有效性。否则，导出的程序不会验证图中的隐含假设，并可能导致原始模型和导出模型之间的行为差异。这在用户需要绕过 tracer 中的问题或逐步增强模型安全性时很有用。请注意，此选项不会影响最终的 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)[源代码]

警告

该项目正处于积极开发阶段，保存的文件在 PyTorch 新版本中可能会失效。

将一个ExportedProgram保存到类似文件的对象中。之后可以使用Python API torch.export.load进行加载。

参数

• ep (ExportedProgram) - 需要保存的导出程序。

• f (Union[str, os.PathLike, io.BytesIO) – 类文件对象（需要实现 write 和 flush 方法）或包含文件名的字符串。

• extra_files (Optional[Dict[str, Any]]) – 文件名到内容的映射，这些内容将作为 f 的一部分进行存储。

• opset_version (Optional[Dict[str, int]]) – 操作集名称到其版本的可选字典映射

示例：

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

警告

该项目正处于积极开发阶段，保存的文件在 PyTorch 新版本中可能会失效。

加载一个之前使用torch.export.save保存的ExportedProgram。

参数

• ep (ExportedProgram) - 需要保存的导出程序。

• f (Union[str, os.PathLike, io.BytesIO) – 类文件对象（需要实现 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 的内容，这是必需的) –

• dataclass -

示例：

@dataclass
class InputDataClass:
    feature: torch.Tensor
    bias: int

class OutputDataClass:
    res: torch.Tensor

torch.export.register_dataclass(InputDataClass)
torch.export.register_dataclass(OutputDataClass)

def fn(o: InputDataClass) -> torch.Tensor:
    res = res=o.feature + o.bias
    return OutputDataClass(res=res)

ep = torch.export.export(fn, (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.dims(*names, min=None, max=None)[源代码]

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

类torch.export.dynamic_shapes.ShapesCollection[源代码]

动态形状的构建器，用于为输入中出现的张量分配动态形状规格。

示例:

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

生成动态形状。

torch.export.dynamic_shapes.refine_dynamic_shapes_from_suggested_fixes(msg, dynamic_shapes)[源代码]

处理导出动态形状的建议修复和/或自动动态形状。根据 ConstraintViolation 错误消息和原始动态形状，精炼给定的动态形状规范。

在大多数情况下，行为很简单：对于那些专门化或细化Dim范围的建议修复，或是提出派生关系的修复，新的动态形状规范将会进行相应的更新。

例如：建议的修复方法：

dim = Dim('dim', min=3, max=6) -> 这只是精炼了 dim 的范围
dim = 4 -> 这将其专门化为常量 4
dy = dx + 1 -> dy 被指定为一个独立的 dim，但实际上与 dx 存在这样的关系

然而，与派生维度相关的建议修复可能会更复杂。例如，如果为根维度提供了建议修复，那么新的派生维度值会根据根维度来评估。

例如：dx = Dim('dx')
dy = dx + 2
dynamic_shapes = {"x": (dx,), "y": (dy,)}

建议的解决方法：

dx = 4 # 特化会使得 dy 也进行特化，即 dy = 6 dx = Dim('dx', max=6) # 此时 dy 的最大值为 8

派生维度的建议修复还可以用来表示可除性约束。这涉及到创建不依赖于特定输入形状的新根维度。在这种情况下，这些根维度不会直接出现在新的规范中，而是会成为其他某个维度的根。

例如：建议的修复方法：

_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]

classtorch.export.ExportedProgram(root, graph, graph_signature, state_dict, range_constraints, module_call_graph, example_inputs=None, constants=None, *, verifiers=None)[源代码]

这是由 export() 函数生成的程序包。它包含了表示张量计算的 torch.fx.Graph，一个包含所有提升参数和缓冲区张量值的状态字典，以及各种元数据。

你可以像调用最初被 export() 追踪的原始可调用对象一样，使用相同的调用约定来调用一个 ExportedProgram。

要对图进行变换，请使用.module属性访问一个torch.fx.GraphModule。然后可以使用FX变换重写图。之后，你可以简单地再次调用export()来构建一个正确的ExportedProgram。

模块()[源代码]

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

返回类型

Module

缓冲区()[源代码]

返回一个遍历原始模块缓冲区的迭代器。

警告

此 API 是实验性的，且不支持向下兼容。

返回类型

Iterator [ Tensor ]

named_buffers()[源代码]

返回一个原始模块缓冲区的迭代器，并同时提供缓冲区的名字和其内容。

警告

此 API 是实验性的，且不支持向下兼容。

返回类型

Iterator[Tuple[str, Tensor]]

参数()[源代码]

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

警告

此 API 是实验性的，且不支持向下兼容。

返回类型

Iterator[Parameter]

named_parameters()[源代码]

返回原始模块参数的迭代器，同时提供参数的名称和其本身的值。

警告

此 API 是实验性的，且不支持向下兼容。

返回类型

Iterator[Tuple[str, Parameter]]

run_decompositions(decomp_table=None, _preserve_ops=())[源代码]

对导出的程序进行一系列分解，并返回一个新的导出程序。默认情况下，我们会使用核心ATen分解来获取核心ATen操作集中的操作。

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

返回类型

ExportedProgram

类torch.export.ExportBackwardSignature(gradients_to_parameters: Dict[str, str], gradients_to_user_inputs: Dict[str, str], loss_output: str)[源代码]

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

类torch.export.ModuleCallSignature(inputs:List[Union[torch.export.graph_signature.TensorArgument, torch.export.graph_signature.SymIntArgument, 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.ConstantArgument, torch.export.graph_signature.CustomObjArgument, torch.export.graph_signature.TokenArgument], in_spec:torch.utils._pytree.TreeSpec, out_spec:torch.utils._pytree.TreeSpec)[源代码]

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

类torch.export.graph_signature.InputKind(value)[源代码]

一个枚举类型。

类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.ConstantArgument,torch.export.graph_signature.CustomObjArgument,torch.export.graph_signature.TokenArgument], target:Optional[str], persistent:Optional[bool]=None)[源代码]

类torch.export.graph_signature.OutputKind(value)[源代码]

一个枚举类型。

类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.ConstantArgument, torch.export.graph_signature.CustomObjArgument, torch.export.graph_signature.TokenArgument], target: Optional[str])[源代码]

类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()[源代码]

类torch.export.graph_signature.CustomObjArgument(name: str, class_fqn: str, fake_val: Optional[torch._library.fake_class_registry.FakeScriptObject] = None)[源代码]

类torch.export.unflatten.FlatArgsAdapter[源代码]

将输入参数适配为符合 target_spec，通过使用 input_spec。

abstractadapt(target_spec, input_spec, input_args)[源代码]

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

返回类型

`List[Any]`

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

这是一个使用 torch.fx.Interpreter 进行执行的模块，而不用通常 GraphModule 中使用的代码生成。这种方式提供更详细的堆栈跟踪信息，使得调试变得更加简单。

torch.export.unflatten.unflatten(module, flat_args_adapter=None)[源代码]

将导出的程序展开，生成一个具有与原始即时执行模块相同层次结构的新模块。如果你尝试使用 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)[源代码]

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

参数

• ep (ExportedProgram) - 需要移动的导出程序。

• location (Union[torch.device, str, Dict[str, str]]) – 导出程序要移动到的设备。如果是一个字符串，它会被解释为设备名称；如果是字典，则表示从现有设备到目标设备的映射。

返回值

已移动的导出程序。

返回类型

ExportedProgram