torch.export

警告

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

概述

采用任意 Python 可调用对象（、函数或方法）并生成跟踪图 仅表示 Ahead-of-Time 中函数的 Tensor 计算 （AOT） 方式执行，随后可以使用不同的输出执行，或者 序列 化。

import torch
from torch.export import export

def f(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(
    f, 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: {}
    Equality constraints: []

torch.export生成一个干净的中间表示 （IR），其中 遵循不变量。有关 IR 的更多规格，请点击此处（即将推出 很快！

• 健全性：保证是原件的合理表示 program 的调用，并保持与原始程序相同的调用约定。

• 规范化：图中没有 Python 语义。子模块 从原始程序内联以形成一个完全展平的程序 计算图。

• 定义的运算符集：生成的图形仅包含一个小的已定义核心 ATen IR 操作集和注册的自定义 运维。

• 图形属性：图形是纯函数式的，这意味着它不是 包含具有副作用（如突变或别名）的操作。它确实 不改变任何中间值、参数或缓冲区。

• 元数据：该图包含在跟踪期间捕获的元数据，例如 stacktrace 从用户的代码中获取。

在后台，利用以下最新技术：torch.export

• TorchDynamo （torch._dynamo） 是一个使用 CPython 功能的内部 API 调用了帧评估 API 以安全地跟踪 PyTorch 图形。这 提供大幅改进的图形捕获体验，而 需要重写才能完全跟踪 PyTorch 代码。

• AOT Autograd 提供功能化的 PyTorch 图，并确保该图 分解/降低到较小的已定义 Core ATen 运算符集。

• Torch FX （torch.fx） 是图形的底层表示， 允许基于 Python 的灵活转换。

现有框架

也使用与 相同的 PT2 堆栈，但 略有不同：torch.export

• JIT 与 AOT： 是 JIT 编译器，而 它不打算用于在 部署。

• 部分图形捕获与完整图形捕获：当遇到 无法追踪的部分，它将 “graph break” 并回退到正在运行的 急切的 Python 运行时中的程序。相比之下，目标 来获取 PyTorch 模型的完整图形表示，因此它会出错 当到达无法追踪的东西时。由于会生成一个完整的 graph 与任何 Python 功能或运行时不相交，那么这个图形可以是 在不同的环境和语言中保存、加载和运行。torch.exporttorch.export

• 可用性权衡：由于能够回退到 每当 Python 运行时达到无法追踪的程度时，它就会多得多 灵活。 将要求用户提供更多 信息或重写其代码以使其可跟踪。torch.export

与 ， 使用 TorchDynamo 在 Python 字节码级别运行，使其能够 跟踪不受 Python 运算符限制的任意 Python 构造 超载支持。此外，还可以对 Tensor 元数据，因此 Tensor 形状等内容上的条件不会 失败跟踪。一般来说，预期会对更多的用户起作用 程序生成较低级别的图形（在运算符 级别）。请注意，用户仍然可以将 预处理步骤 。torch.exporttorch.exporttorch.exporttorch.ops.atentorch.export

与 相比，不捕获 Python 控制流或数据结构，但它支持更多的 Python 语言功能 比 TorchScript 多（因为它更容易全面覆盖 Python 字节码）。生成的图形更简单，并且只有直线控制 flow （显式控制流运算符除外）。torch.export

与 相比，是声音：它能够 跟踪代码，该代码对 sizes 执行整数计算并记录所有 side-条件，以表明特定跟踪对其他 输入。torch.export

导出 PyTorch 模型

示例

主入口点是通过 ，它采用 callable （、 function 或 method） 和示例输入，以及 将计算图捕获到 .一 例：

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: {}
    Equality constraints: []

检查 ，我们可以注意到以下内容：ExportedProgram

• 这包含原始 程序，以及原始代码的记录，以便于调试。

• 该图仅包含在 Core ATen IR opset 中找到的运算符和自定义运算符，并且 功能齐全，无需任何就地运算符，例如 .torch.ops.atentorch.add_

• 参数（权重和对卷积的偏差）作为图形的输入被提升， 导致图中没有节点，而这些节点以前存在于 的结果。get_attr

• 对输入和输出进行建模 签名，并指定哪些输入是参数。

• 图中每个节点生成的张量的最终形状和 dtype 为 著名的。例如，该节点将产生 dtype 和 shape （1， 16， 256， 256） 的张量。convolutiontorch.float32

表达活力

默认情况下，将跟踪程序，假设所有输入形状都是静态的，并将导出的程序专门化到这些维度。然而 某些维度（如批次维度）可以是动态的，并且会因 Run 到 而异 跑。此类维度必须使用 API 标记为动态，并通过参数传入。一个例子：torch.exportconstraints

import torch
from torch.export import export, dynamic_dim

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))
constraints = [
    # First dimension of each input is a dynamic batch size
    dynamic_dim(example_args[0], 0),
    dynamic_dim(example_args[1], 0),
    # The dynamic batch size between the inputs are equal
    dynamic_dim(example_args[0], 0) == dynamic_dim(example_args[1], 0),
]

exported_program: torch.export.ExportedProgram = export(
  M(), args=example_args, constraints=constraints
)
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)}
    Equality constraints: [(InputDim(input_name='arg5_1', dim=0), InputDim(input_name='arg6_1', dim=0))]

一些需要注意的其他事项：

• 通过 API，我们指定了第一个 维度设置为动态的。查看输入 和 ，它们的符号形状为 （s0， 64） 和 （s0， 128），而不是 我们作为示例输入传入的 （32， 64） 和 （32， 128） 形状的张量。 是一个符号，表示此维度可以是范围 的值。arg5_1arg6_1s0

• exported_program.range_constraints描述每个元件的范围 显示在图表中。在本例中，我们看到 具有 [2， inf].由于此处难以解释的技术原因，它们是 假定不是 0 或 1。这不是一个错误，也不一定意味着 导出的程序将不适用于维度 0 或 1。有关此主题的深入讨论，请参阅 0/1 特化问题。s0

• exported_program.equality_constraints描述哪些维度是 要求相等。由于我们在 constraints 中指定了第一个 dimension 是等效的， (), 我们在相等约束中看到指定维度 0 和维度 0 相等的元组。dynamic_dim(example_args[0], 0) == dynamic_dim(example_args[1], 0)arg5_1arg6_1

序列化

要保存 ，用户可以使用 和 API。惯例是使用文件扩展名保存。ExportedProgramExportedProgram.pt2

一个例子：

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

专业化

输入形状

如前所述，默认情况下，将跟踪程序 专门研究输入张量的形状，除非维度指定为 dynamic 的 API 实现。这意味着，如果存在 存在形状相关的控制流，将专门用于 分支。例如：torch.exporttorch.export

import torch
from torch.export import export

def fn(x):
    if x.shape[0] > 5:
        return x + 1
    else:
        return x - 1

example_inputs = (torch.rand(10, 2),)
exported_program = export(fn, 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,)

（） 的条件不会出现在 中，因为示例输入具有静态 的形状 （10， 2）。Since 专门研究输入的 static shapes 时，将永远不会到达 else 分支 （）。要保留动态 需要使用基于跟踪图中张量形状的分支行为来指定维度 的 API 设置为 dynamic 的，源代码将 需要重写。x.shape[0] > 5ExportedProgramtorch.exportx - 1x.shape[0]

非 Tensor 输入

torch.export还根据 inputs 的值专门化 traced graph 不是 ，例如 、 、 和 。 但是，我们可能会在不久的将来更改此设置，使其不再专注于 基元类型的输入。torch.Tensorintfloatboolstr

例如：

import torch
from torch.export import export

def fn(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(fn, 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,)

由于整数是专用的，因此运算 都是使用内联常量 而不是 计算的。 此外，循环中使用的迭代器也是 “inlined” 在图中，通过 3 个重复调用，以及 从不使用 input。torch.ops.aten.add.Tensor1arg1_1timesfortorch.ops.aten.add.Tensorarg2_1

torch.export 的限制

图形中断

从 PyTorch 程序中，它最终可能会遇到程序中无法追踪的部分，如 几乎不可能支持跟踪所有 PyTorch 和 Python 功能。在 的情况下，不支持的操作将导致“图 break“，则不支持的操作将使用默认 Python 评估运行。 相反，将要求用户提供额外的 信息或重写其代码的某些部分以使其可跟踪。由于 跟踪基于 TorchDynamo，它在 Python bytecode 级别，则与 以前的跟踪框架。torch.exporttorch.compiletorch.export

当遇到图形中断时，ExportDB 非常有用 用于了解支持的程序类型的资源，以及 unsupported，以及重写程序以使其可跟踪的方法。

数据/形状相关控制流

当形状没有专用化时，在数据依赖的控制流 （） 上也会遇到图形中断，因为跟踪编译器不能 可能无需为组合爆炸生成代码即可处理 路径数。在这种情况下，用户需要使用 特殊控制流运算符（即将推出！if x.shape[0] > 2

数据依赖访问

数据依赖行为，例如使用张量内部的值来构造 另一个张量，或使用一个张量的值切入另一个张量，是 也是示踪剂无法完全确定的东西。用户将需要重写 他们的代码使用内联约束 API 和 .

缺少 Operator 的 Meta Kernels

跟踪时，所有 运维。这用于推断此 算子。

注意，为自定义 op 注册自定义元内核的官方 API 是 目前正在开发中。在优化最终 API 时，您可以 请参阅此处的文档。

不幸的是，如果您的模型使用了 ATen 运算符，而 则 没有 还没有 meta kernel 实现，请提交一个 issue。

阅读更多

Export 用户的其他链接

• 在 ATen IR 上编写图形转换
• IRs
• 导出数据库

面向 PyTorch 开发人员的深入探讨

• TorchDynamo 深入探究
• 动态形状
• 假张量

API 参考

torch.export 中。export（f， args， kwargs=无， *， constraints=无）

采用任意 Python 可调用对象（nn.Module、函数或 一个方法）并生成一个仅表示 Tensor 的跟踪图 以预先 （AOT） 方式计算函数，它可以 随后使用不同的输出执行或序列化。被追踪的 图 （1） 生成一个标准化的运算符集，该运算符集仅包含功能性 Core ATen 运算符集和用户指定的自定义运算符，（2） 已消除所有 Python 控件 流和数据结构（某些 conditions） 和 （3） 具有一组所需的形状约束，以证明 这种规范化和控制流消除对未来来说是合理的 输入。

稳健性保证

追踪时，注意与形状相关的假设 由用户程序和底层 PyTorch 算子内核制作。 只有在以下情况下，输出才被视为有效 假设是正确的。

在跟踪过程中，会做出 2 种类型的假设

• 输入张量的形状（而不是值）。

• 通过或直接索引从中间张量中提取的值的范围（下限和上限）。.item()

所有假设都必须在图形捕获时进行验证才能成功。具体说来：

• 对输入张量的静态形状的假设会自动验证，无需额外工作。

• 对输入张量动态形状的假设需要使用 API 构造的显式 Input Constraint

• 对中间值范围的假设需要显式的 Inline Constraint， 构造使用和 API。constraint_as_value()

如果任何假设无法验证，则会引发致命错误。当这种情况发生时， 错误消息将包含构建必要代码所需的建议代码 constraints 来验证假设，例如，会建议 以下 input constraints 代码：

def specify_constraints(x):
    return [
        # x:
        dynamic_dim(x, 0) <= 5,
    ]

此示例意味着程序要求 input 的 dim 0 小于 大于或等于 5 才有效。您可以检查所需的约束，并且 然后将这个确切的函数复制到你的代码中，以生成所需的 constraints 传递给 argument。xconstraints

参数

• f （Callable） – 要跟踪的可调用对象。

• args （Tuple[Any， ...]） – 位置输入示例。

• kwargs （Optional[Dict[str， Any]]） – 可选的示例关键字输入。

• constraints （Optional[List[Constraint]]） – 动态参数的可选约束列表 指定其可能的形状范围。默认情况下，形状为 输入 Torch。张量被假定为静态的。如果输入手电筒。张肌 预期具有动态形状，请使用 to define 对象来指定动态和可能的 的形状范围。有关以下示例，请参阅 docstring 如何使用它。

返回

An 包含跟踪的可调用对象。

返回类型

导出程序

可接受的输入/输出类型

可接受的输入类型（for 和 ）和输出包括：argskwargs

• 基元类型，即 、 、 和 。torch.Tensorintfloatboolstr

• （嵌套）由 、 、 、 包含上述所有类型的数据结构。dictlisttuplenamedtupleOrderedDict

torch.export 中。dynamic_dim（t， index）

构造一个对象，该对象描述 张量 的维度 。对象应传递给 的参数。indextconstraints

参数

• t （Torch。Tensor） – 具有动态维度大小的示例输入张量

• index （int） – 动态维度的索引

返回

描述形状动态的对象。它可以传递给 so 这不假设指定 Tensor 的静态大小，即保持其动态 作为符号大小，而不是根据示例跟踪输入的大小进行专门化。

具体来说，可以用来表达以下类型的活力。

• 维度的大小是动态的且不受限制的：

  t0 = torch.rand(2, 3)
  t1 = torch.rand(3, 4)
  
  # First dimension of t0 can be dynamic size rather than always being static size 2
  constraints = [dynamic_dim(t0, 0)]
  ep = export(fn, (t0, t1), constraints=constraints)
  

• 维度的大小是动态的，具有下限：

  t0 = torch.rand(10, 3)
  t1 = torch.rand(3, 4)
  
  # First dimension of t0 can be dynamic size with a lower bound of 5 (inclusive)
  # Second dimension of t1 can be dynamic size with a lower bound of 2 (exclusive)
  constraints = [
      dynamic_dim(t0, 0) >= 5,
      dynamic_dim(t1, 1) > 2,
  ]
  ep = export(fn, (t0, t1), constraints=constraints)
  

• 维度的大小是动态的，具有上限：

  t0 = torch.rand(10, 3)
  t1 = torch.rand(3, 4)
  
  # First dimension of t0 can be dynamic size with a upper bound of 16 (inclusive)
  # Second dimension of t1 can be dynamic size with a upper bound of 8 (exclusive)
  constraints = [
      dynamic_dim(t0, 0) <= 16,
      dynamic_dim(t1, 1) < 8,
  ]
  ep = export(fn, (t0, t1), constraints=constraints)
  

• 维度的大小是动态的，并且它始终等于另一个动态维度的大小：

  t0 = torch.rand(10, 3)
  t1 = torch.rand(3, 4)
  
  # Sizes of second dimension of t0 and first dimension are always equal
  constraints = [
      dynamic_dim(t0, 1) == dynamic_dim(t1, 0),
  ]
  ep = export(fn, (t0, t1), constraints=constraints)
  

• 混合并匹配上述所有类型，只要它们不表达冲突的要求

torch.export 中。constrain_as_size（符号，min=无，max=None）

关于中间标量值的约束的提示，该值 表示张量的形状，以便后续的张量构造函数可以是 跟踪正确，因为许多运算符需要对范围进行假设 的大小。

参数

• symbol – 要应用范围约束的中间标量值 （现在仅为 int）。

• min （Optional[int]） – 给定交易品种的最小可能值 （含）

• max （Optional[int]） - 给定交易品种的最大可能值 （含）

返回

没有

例如，如果没有 用于提供有关形状范围的提示，则无法清楚地跟踪以下程序：

def fn(x):
    d = x.max().item()
    return torch.ones(v)

将给出以下错误：

torch._dynamo.exc.Unsupported: guard on data-dependent symbolic int/float

假设 的实际范围可以介于 [3， 10] 之间，您可以在源代码中添加对的调用，如下所示：d

def fn(x):
    d = x.max().item()
    torch.export.constrain_as_size(d, min=3, max=10)
    return torch.ones(d)

通过额外的提示，将能够通过获取 分支，结果如下：else

graph():
    %arg0_1 := placeholder[target=arg0_1]

    # d = x.max().item()
    %max_1 := call_function[target=torch.ops.aten.max.default](args = (%arg0_1,))
    %_local_scalar_dense := call_function[target=torch.ops.aten._local_scalar_dense.default](args = (%max_1,))

    # Asserting 3 <= d <= 10
    %ge := call_function[target=operator.ge](args = (%_local_scalar_dense, 3))
    %scalar_tensor := call_function[target=torch.ops.aten.scalar_tensor.default](args = (%ge,))
    %_assert_async := call_function[target=torch.ops.aten._assert_async.msg](
        args = (%scalar_tensor, _local_scalar_dense is outside of inline constraint [3, 10].))
    %le := call_function[target=operator.le](args = (%_local_scalar_dense, 10))
    %scalar_tensor_1 := call_function[target=torch.ops.aten.scalar_tensor.default](args = (%le,))
    %_assert_async_1 := call_function[target=torch.ops.aten._assert_async.msg](
        args = (%scalar_tensor_1, _local_scalar_dense is outside of inline constraint [3, 10].))
    %sym_constrain_range_for_size := call_function[target=torch.ops.aten.sym_constrain_range_for_size.default](
        args = (%_local_scalar_dense,), kwargs = {min: 3, max: 10})

    # Constructing new tensor with d
    %full := call_function[target=torch.ops.aten.full.default](
        args = ([%_local_scalar_dense], 1),
        kwargs = {dtype: torch.float32, layout: torch.strided, device: cpu, pin_memory: False})

    ......

警告

如果您的尺寸是动态的，请不要测试尺寸是否等于 0 或 1， 这些将静默报告 false 并被绕过

torch.export 中。constrain_as_value（符号，min=无，max=None）

关于中间标量值的约束的提示，以便后续的 检查上述标量值范围的分支行为可以是 追踪清晰。

警告

（请注意，如果中间标量值将像大小一样使用，包括 作为 size arg 传递给 Tensor Factory 或 View，请改为调用

参数

• symbol – 要应用范围约束的中间标量值 （现在仅为 int）。

• min （Optional[int]） – 给定交易品种的最小可能值 （含）

• max （Optional[int]） - 给定交易品种的最大可能值 （含）

返回

没有

例如，无法合理追踪以下程序：

def fn(x):
    v = x.max().item()
    if v > 1024:
        return x
    else:
        return x * 2

v是一个与数据相关的值，假定其范围为 （-inf， inf）。关于要采用哪个分支的提示将无法确定 跟踪的分支决策是否正确。因此会给出以下错误：

torch._dynamo.exc.UserError: Consider annotating your code using
torch.export.constrain_as_size() or torch.export().constrain_as_value() APIs.
It appears that you're trying to get a value out of symbolic int/float whose value
is data-dependent (and thus we do not know the true value.)  The expression we were
trying to evaluate is f0 > 1024 (unhinted: f0 > 1024).

假设 的实际范围可以介于 [10， 200] 之间，您可以在源代码中添加对的调用，如下所示：v

def fn(x):
    v = x.max().item()

    # Give export() a hint
    torch.export.constrain_as_value(v, min=10, max=200)

    if v > 1024:
        return x
    else:
        return x * 2

通过额外的提示，将能够通过获取 分支，结果如下：else

graph():
    %arg0_1 := placeholder[target=arg0_1]

    # v = x.max().item()
    %max_1 := call_function[target=torch.ops.aten.max.default](args = (%arg0_1,))
    %_local_scalar_dense := call_function[target=torch.ops.aten._local_scalar_dense.default](args = (%max_1,))

    # Asserting 10 <= v <= 200
    %ge := call_function[target=operator.ge](args = (%_local_scalar_dense, 10))
    %scalar_tensor := call_function[target=torch.ops.aten.scalar_tensor.default](args = (%ge,))
    %_assert_async := call_function[target=torch.ops.aten._assert_async.msg](
        args = (%scalar_tensor, _local_scalar_dense is outside of inline constraint [10, 200].))
    %le := call_function[target=operator.le](args = (%_local_scalar_dense, 200))
    %scalar_tensor_1 := call_function[target=torch.ops.aten.scalar_tensor.default](args = (%le,))
    %_assert_async_1 := call_function[target=torch.ops.aten._assert_async.msg](
        args = (%scalar_tensor_1, _local_scalar_dense is outside of inline constraint [10, 200].))
    %sym_constrain_range := call_function[target=torch.ops.aten.sym_constrain_range.default](
        args = (%_local_scalar_dense,), kwargs = {min: 10, max: 200})

    # Always taking `else` branch to multiply elements `x` by 2 due to hints above
    %mul := call_function[target=torch.ops.aten.mul.Tensor](args = (%arg0_1, 2), kwargs = {})
    return (mul,)

torch.export 中。save（ep， f， *， extra_files=无， opset_version=无）

警告

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

将 保存到类似文件的对象。然后它可以是 使用 Python API 加载。

参数

• ep （ExportedProgram） – 要保存的导出程序。

• f （Union[str， pathlib.路径，io.BytesIO） – 一个类似文件的对象 （必须 implement write 和 flush）或包含文件名的字符串。

• 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'}
torch.export.save(ep, 'exported_program.pt2', extra_files=extra_files)

torch.export 中。load（f， *， extra_files=无， expected_opset_version=无）

警告

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

加载以前使用 保存的 。

参数

• ep （ExportedProgram） – 要保存的导出程序。

• f （Union[str， pathlib.路径，io.BytesIO） – 一个类似文件的对象 （必须 implement write 和 flush）或包含文件名的字符串。

• extra_files （Optional[Dict[str， Any]]） – 在 此映射将被加载，其内容将存储在 提供的地图。

• expected_opset_version （Optional[Dict[str， int]]） – opset 名称的映射 到预期的 Opset 版本

返回

对象

返回类型

导出程序

例：

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

类 torch.export 中。Constraint（*args， **kwargs）

警告

不要直接构造，而是使用。

这表示对 input tensor 维度的约束，例如，要求 它们是完全多态的或在某个范围内。

类 torch.export 中。ExportedProgram（根、图形、graph_signature、call_spec、state_dict、range_constraints、equality_constraints、module_call_graph、example_inputs=None）

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

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

要对图形执行转换，请使用 property 访问 一个 .然后，您可以使用 FX 转换来重写图形。之后，您可以简单地再次使用来构建正确的 ExportedProgram。.module

module（）

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

返回类型

模块

类 torch.export 中。ExportBackwardSignature（gradients_to_parameters： Dict[str， str]， gradients_to_user_inputs： Dict[str， str]， loss_output： str）

类 torch.export 中。ExportGraphSignature（参数、缓冲区、user_inputs、user_outputs、inputs_to_parameters、inputs_to_buffers、buffers_to_mutate、backward_signature、 assertion_dep_token=无）

对 Export Graph 的输入/输出签名进行建模， 这就是 FX。具有更强不变量保证的图形。

Export Graph 是功能性的，并且不访问参数之类的“状态” 或通过节点在图中的缓冲区。相反，可以保证将参数和缓冲区作为输入从图形中提取出来。 同样，对 buffers 的任何更改也不包含在图中， 相反，突变缓冲区的更新值被建模为附加输出 导出图形。getattr

所有输入和输出的顺序为：

Inputs = [*parameters_buffers, *flattened_user_inputs]
Outputs = [*mutated_inputs, *flattened_user_outputs]

例如，如果导出了以下模块：

class CustomModule(nn.Module):
    def __init__(self):
        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 将为：

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(
    # Indicates that there is one parameter named `my_parameter`
    parameters=['L__self___my_parameter'],

    # Indicates that there are two buffers, `my_buffer1` and `my_buffer2`
    buffers=['L__self___my_buffer1', 'L__self___my_buffer2'],

    # Indicates that the nodes `arg3_1` and `arg4_1` in produced graph map to
    # original user inputs, ie. x1 and x2
    user_inputs=['arg3_1', 'arg4_1'],

    # Indicates that the node `add_tensor_1` maps to output of original program
    user_outputs=['add_tensor_1'],

    # Indicates that there is one parameter (self.my_parameter) captured,
    # its name is now mangled to be `L__self___my_parameter`, which is now
    # represented by node `arg0_1` in the graph.
    inputs_to_parameters={'arg0_1': 'L__self___my_parameter'},

    # Indicates that there are two buffers (self.my_buffer1, self.my_buffer2) captured,
    # their name are now mangled to be `L__self___my_my_buffer1` and `L__self___my_buffer2`.
    # They are now represented by nodes `arg1_1` and `arg2_1` in the graph.
    inputs_to_buffers={'arg1_1': 'L__self___my_buffer1', 'arg2_1': 'L__self___my_buffer2'},

    # Indicates that one buffer named `L__self___my_buffer2` is mutated during execution,
    # its new value is output from the graph represented by the node named `add_tensor_2`
    buffers_to_mutate={'add_tensor_2': 'L__self___my_buffer2'},

    # Backward graph not captured
    backward_signature=None,

    # Work in progress feature, please ignore now.
    assertion_dep_token=None
)

类 torch.export 中。ArgumentKind（值）

一个枚举。

类 torch.export 中。ArgumentSpec（种类：torch.export.ArgumentKind，值：Any）

类 torch.export 中。ModuleCallSignature（inputs： List[torch.export.ArgumentSpec]， 输出： List[torch.export.ArgumentSpec]， in_spec： torch.utils._pytree。TreeSpec，out_spec：torch.utils._pytree。TreeSpec）

类 torch.export 中。ModuleCallEntry（fqn： str， signature： Union[torch.export.ModuleCallSignature， NoneType] = 无）