评价此页

torch.export 教程#

创建于:2023 年 10 月 2 日 | 上次更新:2026 年 3 月 24 日 | 上次验证:2024 年 11 月 5 日

作者: William Wen, Zhengxu Chen, Angela Yi, Pian Pawakapan

警告

torch.export 及其相关功能处于原型状态,可能会发生向后兼容性的重大更改。本教程提供了截至 PyTorch 2.5 的 torch.export 使用快照。

torch.export() 是 PyTorch 2.X 将 PyTorch 模型导出为标准化模型表示的方法,旨在不同(即无 Python)的环境中运行。官方文档可以参考 这里

在本教程中,您将学习如何使用 torch.export() 从 PyTorch 程序中提取 ExportedProgram(即单图表示)。我们还将详细说明一些注意事项/修改,您可能需要执行这些操作才能使模型与 torch.export 兼容。

内容

基本用法#

torch.export 通过在给定示例输入的情况下追踪(tracing)目标函数,从 PyTorch 程序中提取单图表示。torch.export.export()torch.export 的主要入口点。

在本教程中,torch.exporttorch.export.export() 实际上几乎同义,尽管 torch.export 通常指 PyTorch 2.X 的导出过程,而 torch.export.export() 通常指实际的函数调用。

torch.export.export() 的签名是

export(
    mod: torch.nn.Module,
    args: Tuple[Any, ...],
    kwargs: Optional[Dict[str, Any]] = None,
    *,
    dynamic_shapes: Optional[Dict[str, Dict[int, Dim]]] = None
) -> ExportedProgram

torch.export.export() 追踪调用 mod(*args, **kwargs) 的张量计算图,并将其封装在一个 ExportedProgram 中,该程序可以序列化或稍后使用不同输入执行。要执行 ExportedProgram,我们可以对其调用 .module() 以返回一个可调用的 torch.nn.Module,就像原始程序一样。我们将在本教程后面详细介绍 dynamic_shapes 参数。

import torch
from torch.export import export

class MyModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.lin = torch.nn.Linear(100, 10)

    def forward(self, x, y):
        return torch.nn.functional.relu(self.lin(x + y), inplace=True)

mod = MyModule()
exported_mod = export(mod, (torch.randn(8, 100), torch.randn(8, 100)))
print(type(exported_mod))
print(exported_mod.module()(torch.randn(8, 100), torch.randn(8, 100)))
<class 'torch.export.exported_program.ExportedProgram'>
tensor([[0.0000, 0.5046, 0.3134, 1.5888, 0.0000, 1.5607, 0.9119, 0.5936, 0.0000,
         0.0000],
        [0.3313, 0.0000, 0.5634, 0.3717, 0.8447, 0.6885, 0.0000, 0.2318, 0.5655,
         0.0000],
        [0.3732, 0.0000, 1.9044, 1.2545, 1.1205, 0.1476, 1.7466, 0.5967, 0.6459,
         0.0000],
        [0.0000, 0.0000, 0.0000, 2.0516, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         0.0000],
        [0.0000, 0.3533, 0.1453, 0.0000, 0.0000, 0.5596, 1.5400, 0.0000, 0.0000,
         0.5601],
        [0.3208, 0.0000, 0.3426, 0.0000, 0.0652, 0.0000, 0.0000, 0.5805, 0.3502,
         1.0795],
        [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
         1.1037],
        [1.2570, 0.0000, 0.6683, 1.0902, 0.0000, 0.1108, 0.0387, 0.0000, 0.0000,
         0.0000]], grad_fn=<ReluBackward0>)

让我们回顾一下 ExportedProgram 中值得关注的一些属性。

graph 属性是从我们导出的函数中追踪到的 FX 图,即所有 PyTorch 操作的计算图。FX 图处于 “ATen IR” 中,意味着它只包含 “ATen 级别” 的操作。

graph_signature 属性对导出图中的输入和输出节点进行了更详细的描述,说明了哪些是参数、缓冲区、用户输入或用户输出。

range_constraints 属性将在稍后介绍。

print(exported_mod)
ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, p_lin_weight: "f32[10, 100]", p_lin_bias: "f32[10]", x: "f32[8, 100]", y: "f32[8, 100]"):
            # File: /var/lib/workspace/intermediate_source/torch_export_tutorial.py:71 in forward, code: return torch.nn.functional.relu(self.lin(x + y), inplace=True)
            add: "f32[8, 100]" = torch.ops.aten.add.Tensor(x, y);  x = y = None

            # File: /var/lib/ci-user/.local/lib/python3.10/site-packages/torch/nn/modules/linear.py:134 in forward, code: return F.linear(input, self.weight, self.bias)
            linear: "f32[8, 10]" = torch.ops.aten.linear.default(add, p_lin_weight, p_lin_bias);  add = p_lin_weight = p_lin_bias = None

            # File: /var/lib/workspace/intermediate_source/torch_export_tutorial.py:71 in forward, code: return torch.nn.functional.relu(self.lin(x + y), inplace=True)
            relu_: "f32[8, 10]" = torch.ops.aten.relu_.default(linear);  linear = None
            return (relu_,)

Graph signature:
    # inputs
    p_lin_weight: PARAMETER target='lin.weight'
    p_lin_bias: PARAMETER target='lin.bias'
    x: USER_INPUT
    y: USER_INPUT

    # outputs
    relu_: USER_OUTPUT

Range constraints: {}

更多详情请参阅 torch.export 文档

图断裂 (Graph Breaks)#

虽然 torch.exporttorch.compile 共享部分组件,但 torch.export 的关键限制(尤其是与 torch.compile 相比时)是它不支持图断裂。这是因为处理图断裂涉及使用默认的 Python 评估来解释不支持的操作,这与导出用例不兼容。因此,为了使您的模型代码与 torch.export 兼容,您需要修改代码以移除图断裂。

图断裂在以下情况下是必需的:

  • 数据依赖的控制流

class Bad1(torch.nn.Module):
    def forward(self, x):
        if x.sum() > 0:
            return torch.sin(x)
        return torch.cos(x)

import traceback as tb
try:
    export(Bad1(), (torch.randn(3, 3),))
except Exception:
    tb.print_exc()
def forward(self, arg0_1: "f32[3, 3]"):
    # File: /var/lib/workspace/intermediate_source/torch_export_tutorial.py:116 in forward, code: if x.sum() > 0:
    sum_1: "f32[]" = torch.ops.aten.sum.default(arg0_1);  arg0_1 = None
    gt: "b8[]" = torch.ops.aten.gt.Scalar(sum_1, 0);  sum_1 = None
    ne: "b8[]" = torch.ops.aten.ne.Scalar(gt, 0);  gt = None
    item: "Sym(Eq(u0, 1))" = torch.ops.aten.item.default(ne);  ne = item = None




def forward(self, arg0_1: "f32[3, 3]"):
    # File: /var/lib/workspace/intermediate_source/torch_export_tutorial.py:116 in forward, code: if x.sum() > 0:
    sum_1: "f32[]" = torch.ops.aten.sum.default(arg0_1);  arg0_1 = None
    gt: "b8[]" = torch.ops.aten.gt.Scalar(sum_1, 0);  sum_1 = None
    ne: "b8[]" = torch.ops.aten.ne.Scalar(gt, 0);  gt = None
    item: "Sym(Eq(u0, 1))" = torch.ops.aten.item.default(ne);  ne = item = None

Traceback (most recent call last):
  File "/var/lib/workspace/intermediate_source/torch_export_tutorial.py", line 122, in <module>
    export(Bad1(), (torch.randn(3, 3),))
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/__init__.py", line 205, in export
    raise e
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/__init__.py", line 171, in export
    return _export(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1344, in wrapper
    raise e
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1310, in wrapper
    ep = fn(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/exported_program.py", line 124, in wrapper
    return fn(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_utils_internal.py", line 96, in wrapper_function
    return function(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2516, in _export
    ep = _export_for_training(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1344, in wrapper
    raise e
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1310, in wrapper
    ep = fn(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/exported_program.py", line 124, in wrapper
    return fn(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2304, in _export_for_training
    export_artifact = export_func(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2233, in _non_strict_export
    aten_export_artifact = _to_aten_func(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2010, in _export_to_aten_ir_make_fx
    gm, graph_signature = transform(_make_fx_helper)(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2140, in _aot_export_non_strict
    gm, sig = aot_export(stack, wrapped_mod, args, kwargs=kwargs, **flags)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1916, in _make_fx_helper
    gm = make_fx(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 3060, in wrapped
    return make_fx_tracer.trace(f, *args)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 2962, in trace
    return self._trace_inner(f, *args)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 2923, in _trace_inner
    t = dispatch_trace(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_compile.py", line 54, in inner
    return disable_fn(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 1445, in _fn
    return fn(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 1691, in dispatch_trace
    graph = tracer.trace(root, concrete_args)  # type: ignore[arg-type]
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 2497, in trace
    res = super().trace(root, concrete_args)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 914, in trace
    (self.create_arg(fn(*args)),),
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 1761, in wrapped
    out = f(*tensors)  # type:ignore[call-arg]
  File "<string>", line 1, in <lambda>
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1798, in wrapped_fn
    return tuple(flat_fn(*args))
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/utils.py", line 192, in flat_fn
    tree_out = fn(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/graph_capture_wrappers.py", line 1536, in functional_call
    out = mod(*args[params_len:], **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 888, in module_call_wrapper
    return self.call_module(mod, forward, args, kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 2586, in call_module
    return Tracer.call_module(self, m, forward, args, kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 577, in call_module
    ret_val = forward(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 881, in forward
    return _orig_module_call(mod, *args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
    return forward_call(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2124, in forward
    tree_out = mod(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 888, in module_call_wrapper
    return self.call_module(mod, forward, args, kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 2586, in call_module
    return Tracer.call_module(self, m, forward, args, kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 577, in call_module
    ret_val = forward(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 881, in forward
    return _orig_module_call(mod, *args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
    return forward_call(*args, **kwargs)
  File "/var/lib/workspace/intermediate_source/torch_export_tutorial.py", line 116, in forward
    if x.sum() > 0:
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 1823, in __torch_function__
    return func(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 1909, in __torch_function__
    return func(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_export/non_strict_utils.py", line 1205, in __torch_function__
    return func(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/sym_node.py", line 576, in guard_bool
    r = self.evaluate()
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/sym_node.py", line 550, in evaluate
    return self.shape_env.evaluate_sym_node(self, size_oblivious)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/symbolic_shapes.py", line 8107, in evaluate_sym_node
    return self.evaluate_expr(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/symbolic_shapes.py", line 8203, in evaluate_expr
    return self._inner_evaluate_expr(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/recording.py", line 297, in wrapper
    return retlog(fn(*args, **kwargs))
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/symbolic_shapes.py", line 8226, in _inner_evaluate_expr
    return self._evaluate_expr(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/symbolic_shapes.py", line 8459, in _evaluate_expr
    raise self._make_data_dependent_error(
torch.fx.experimental.symbolic_shapes.GuardOnDataDependentSymNode: Could not guard on data-dependent expression Eq(u0, 1) (unhinted: Eq(u0, 1)).  (Size-like symbols: none)

consider using data-dependent friendly APIs such as guard_or_false, guard_or_true and statically_known_true.
Caused by: (_export/non_strict_utils.py:1205 in __torch_function__)
For more information, run with TORCH_LOGS="dynamic"
For extended logs when we create symbols, also add TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="u0"
If you suspect the guard was triggered from C++, add TORCHDYNAMO_EXTENDED_DEBUG_CPP=1
For more debugging help, see https://docs.google.com/document/d/1HSuTTVvYH1pTew89Rtpeu84Ht3nQEFTYhAX3Ypa_xJs/edit?usp=sharing

For C++ stack trace, run with TORCHDYNAMO_EXTENDED_DEBUG_CPP=1

The following call raised this error:
  File "/var/lib/workspace/intermediate_source/torch_export_tutorial.py", line 116, in forward
    if x.sum() > 0:


The error above occurred when calling torch.export.export. If you would like to view some more information about this error, and get a list of all other errors that may occur in your export call, you can replace your `export()` call with `draft_export()`.
  • 使用 .data 访问张量数据

class Bad2(torch.nn.Module):
    def forward(self, x):
        x.data[0, 0] = 3
        return x

try:
    export(Bad2(), (torch.randn(3, 3),))
except Exception:
    tb.print_exc()
Traceback (most recent call last):
  File "/var/lib/workspace/intermediate_source/torch_export_tutorial.py", line 135, in <module>
    export(Bad2(), (torch.randn(3, 3),))
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/__init__.py", line 205, in export
    raise e
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/__init__.py", line 171, in export
    return _export(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1344, in wrapper
    raise e
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1310, in wrapper
    ep = fn(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/exported_program.py", line 124, in wrapper
    return fn(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_utils_internal.py", line 96, in wrapper_function
    return function(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2516, in _export
    ep = _export_for_training(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1344, in wrapper
    raise e
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1310, in wrapper
    ep = fn(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/exported_program.py", line 124, in wrapper
    return fn(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2330, in _export_for_training
    raise RuntimeError(error_msg)
RuntimeError: We found a fake tensor in the exported program constant's list. This typically means our tracing system encountered an op that we can't trace through. For the potential source, you can refer to following model attribute: lifted_tensor_1. Please file an issue on github.
  • 调用不支持的函数(如许多内置函数)

class Bad3(torch.nn.Module):
    def forward(self, x):
        x = x + 1
        return x + id(x)

try:
    export(Bad3(), (torch.randn(3, 3),))
except Exception:
    tb.print_exc()

非严格模式导出#

为了追踪程序,torch.export 默认使用 TorchDynamo(一个字节码分析引擎)对 Python 代码进行符号分析,并根据结果构建图。这种分析允许 torch.export 提供更强的安全性保证,但并非所有 Python 代码都受支持,从而导致这些图断裂。

为了解决这个问题,在 PyTorch 2.3 中,我们引入了一种称为非严格模式的新导出方式,在这种模式下,我们使用 Python 解释器追踪程序,就像在 eager 模式下一样准确地执行它,允许我们跳过不支持的 Python 特性。这是通过添加 strict=False 标志来实现的。

观察之前一些导致图断裂的示例

  • 调用不支持的函数(如许多内置函数)时可以追踪

通过,但在这种情况下,id(x) 在图中会被特化为一个常数整数。这是因为 id(x) 不是一个张量操作,所以该操作不会记录在图中。

class Bad3(torch.nn.Module):
    def forward(self, x):
        x = x + 1
        return x + id(x)

bad3_nonstrict = export(Bad3(), (torch.randn(3, 3),), strict=False)
print(bad3_nonstrict)
print(bad3_nonstrict.module()(torch.ones(3, 3)))
ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, x: "f32[3, 3]"):
            # File: /var/lib/workspace/intermediate_source/torch_export_tutorial.py:179 in forward, code: x = x + 1
            add: "f32[3, 3]" = torch.ops.aten.add.Tensor(x, 1);  x = None

            # File: /var/lib/workspace/intermediate_source/torch_export_tutorial.py:180 in forward, code: return x + id(x)
            add_1: "f32[3, 3]" = torch.ops.aten.add.Tensor(add, 140255759236960);  add = None
            return (add_1,)

Graph signature:
    # inputs
    x: USER_INPUT

    # outputs
    add_1: USER_OUTPUT

Range constraints: {}

tensor([[1.4026e+14, 1.4026e+14, 1.4026e+14],
        [1.4026e+14, 1.4026e+14, 1.4026e+14],
        [1.4026e+14, 1.4026e+14, 1.4026e+14]])

然而,仍然有一些特性需要对原始模块进行重写

控制流操作 (Control Flow Ops)#

torch.export 实际上支持数据依赖的控制流。但这些需要使用控制流操作来表达。例如,我们可以使用 cond 操作来修复上面的控制流示例,如下所示

class Bad1Fixed(torch.nn.Module):
    def forward(self, x):
        def true_fn(x):
            return torch.sin(x)
        def false_fn(x):
            return torch.cos(x)
        return torch.cond(x.sum() > 0, true_fn, false_fn, [x])

exported_bad1_fixed = export(Bad1Fixed(), (torch.randn(3, 3),))
print(exported_bad1_fixed)
print(exported_bad1_fixed.module()(torch.ones(3, 3)))
print(exported_bad1_fixed.module()(-torch.ones(3, 3)))
ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, x: "f32[3, 3]"):
            # File: /var/lib/workspace/intermediate_source/torch_export_tutorial.py:205 in forward, code: return torch.cond(x.sum() > 0, true_fn, false_fn, [x])
            sum_1: "f32[]" = torch.ops.aten.sum.default(x)
            gt: "b8[]" = torch.ops.aten.gt.Scalar(sum_1, 0);  sum_1 = None

            # File: <eval_with_key>.36:9 in forward, code: cond = torch.ops.higher_order.cond(l_args_0_, cond_true_0, cond_false_0, (l_args_3_0_,));  l_args_0_ = cond_true_0 = cond_false_0 = l_args_3_0_ = None
            true_graph_0 = self.true_graph_0
            false_graph_0 = self.false_graph_0
            cond = torch.ops.higher_order.cond(gt, true_graph_0, false_graph_0, (x,));  gt = true_graph_0 = false_graph_0 = x = None
            getitem: "f32[3, 3]" = cond[0];  cond = None
            return (getitem,)

        class true_graph_0(torch.nn.Module):
            def forward(self, x: "f32[3, 3]"):
                # File: <eval_with_key>.37:6 in forward, code: sin = torch.sin(l_args_3_0__1);  l_args_3_0__1 = None
                sin: "f32[3, 3]" = torch.ops.aten.sin.default(x);  x = None
                return (sin,)

        class false_graph_0(torch.nn.Module):
            def forward(self, x: "f32[3, 3]"):
                # File: <eval_with_key>.38:6 in forward, code: cos = torch.cos(l_args_3_0__1);  l_args_3_0__1 = None
                cos: "f32[3, 3]" = torch.ops.aten.cos.default(x);  x = None
                return (cos,)

Graph signature:
    # inputs
    x: USER_INPUT

    # outputs
    getitem: USER_OUTPUT

Range constraints: {}

tensor([[0.8415, 0.8415, 0.8415],
        [0.8415, 0.8415, 0.8415],
        [0.8415, 0.8415, 0.8415]])
tensor([[0.5403, 0.5403, 0.5403],
        [0.5403, 0.5403, 0.5403],
        [0.5403, 0.5403, 0.5403]])

使用 cond 时需要注意一些局限性

  • 谓词(即 x.sum() > 0)必须产生一个布尔值或单元素张量。

  • 操作数(即 [x])必须是张量。

  • 分支函数(即 true_fnfalse_fn)的签名必须与操作数匹配,并且它们必须都返回一个具有相同元数据(例如 dtypeshape 等)的单张量。

  • 分支函数不能修改输入或全局变量。

  • 分支函数不能访问闭包变量,除非该函数是在方法范围内定义的且访问的是 self

有关 cond 的更多详情,请查看 cond 文档

我们还可以使用 map,它将一个函数应用于第一个张量参数的第一维度。

from torch._higher_order_ops.map import map as torch_map

class MapModule(torch.nn.Module):
    def forward(self, xs, y, z):
        def body(x, y, z):
            return x + y + z

        return torch_map(body, xs, y, z)

inps = (torch.ones(6, 4), torch.tensor(5), torch.tensor(4))
exported_map_example = export(MapModule(), inps)
print(exported_map_example)
print(exported_map_example.module()(*inps))
ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, xs: "f32[6, 4]", y: "i64[]", z: "i64[]"):
            # File: <eval_with_key>.96:9 in forward, code: map_impl = torch.ops.higher_order.map_impl(map_body_0, [l_flat_xs_0_], [l_flat_args_0_, l_flat_args_1_]);  map_body_0 = l_flat_xs_0_ = l_flat_args_0_ = l_flat_args_1_ = None
            body_graph_0 = self.body_graph_0
            map_impl = torch.ops.higher_order.map_impl(body_graph_0, [xs], [y, z]);  body_graph_0 = xs = y = z = None
            getitem: "f32[6, 4]" = map_impl[0];  map_impl = None
            return (getitem,)

        class body_graph_0(torch.nn.Module):
            def forward(self, xs: "f32[4]", y: "i64[]", z: "i64[]"):
                # File: <eval_with_key>.97:5 in forward, code: add = child + l_flat_args_0_;  child = l_flat_args_0_ = None
                add: "f32[4]" = torch.ops.aten.add.Tensor(xs, y);  xs = y = None

                # File: <eval_with_key>.97:6 in forward, code: add_1 = add + l_flat_args_1_;  add = l_flat_args_1_ = None
                add_1: "f32[4]" = torch.ops.aten.add.Tensor(add, z);  add = z = None
                return (add_1,)

Graph signature:
    # inputs
    xs: USER_INPUT
    y: USER_INPUT
    z: USER_INPUT

    # outputs
    getitem: USER_OUTPUT

Range constraints: {}

tensor([[10., 10., 10., 10.],
        [10., 10., 10., 10.],
        [10., 10., 10., 10.],
        [10., 10., 10., 10.],
        [10., 10., 10., 10.],
        [10., 10., 10., 10.]])

其他控制流操作包括 while_loopassociative_scanscan。有关各操作符的更多文档,请参阅此页面

约束/动态形状 (Constraints/Dynamic Shapes)#

本节涵盖已导出程序的动态行为和表示。动态行为取决于被导出的特定模型,因此在本教程的大部分内容中,我们将专注于这个特定的玩具模型(并附带产生的张量形状注释)

class DynamicModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.l = torch.nn.Linear(5, 3)

    def forward(
        self,
        w: torch.Tensor,  # [6, 5]
        x: torch.Tensor,  # [4]
        y: torch.Tensor,  # [8, 4]
        z: torch.Tensor,  # [32]
    ):
        x0 = x + y  # [8, 4]
        x1 = self.l(w)  # [6, 3]
        x2 = x0.flatten()  # [32]
        x3 = x2 + z  # [32]
        return x1, x3

默认情况下,torch.export 产生一个静态程序。其后果之一是,在运行时,该程序将无法在具有不同形状的输入上运行,即使它们在 eager 模式下是有效的。

w = torch.randn(6, 5)
x = torch.randn(4)
y = torch.randn(8, 4)
z = torch.randn(32)
model = DynamicModel()
ep = export(model, (w, x, y, z))
model(w, x, torch.randn(3, 4), torch.randn(12))
try:
    ep.module()(w, x, torch.randn(3, 4), torch.randn(12))
except Exception:
    tb.print_exc()
Traceback (most recent call last):
  File "/var/lib/workspace/intermediate_source/torch_export_tutorial.py", line 286, in <module>
    ep.module()(w, x, torch.randn(3, 4), torch.randn(12))
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/graph_module.py", line 1000, in call_wrapped
    return self._wrapped_call(self, *args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/graph_module.py", line 507, in __call__
    raise e
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/graph_module.py", line 493, in __call__
    return super(self.cls, obj).__call__(*args, **kwargs)  # type: ignore[misc]
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1884, in _call_impl
    return inner()
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1832, in inner
    result = forward_call(*args, **kwargs)
  File "<eval_with_key>.133", line 8, in forward
    _guards_fn = self._guards_fn(w, x, y, z);  _guards_fn = None
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
    return forward_call(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_dynamo/external_utils.py", line 215, in inner
    return func(*args, **kwargs)
  File "<string>", line 6, in _
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/__init__.py", line 2279, in _assert
    raise AssertionError(message)
AssertionError: Guard failed: y.size()[0] == 8

基本概念:符号与守卫#

为了启用动态性,export() 提供了一个 dynamic_shapes 参数。处理动态形状最简单的方法是使用 Dim.AUTO 并查看返回的程序。动态行为是在输入维度级别指定的;对于每个输入,我们可以指定一个值元组

from torch.export.dynamic_shapes import Dim

dynamic_shapes = {
    "w": (Dim.AUTO, Dim.AUTO),
    "x": (Dim.AUTO,),
    "y": (Dim.AUTO, Dim.AUTO),
    "z": (Dim.AUTO,),
}
ep = export(model, (w, x, y, z), dynamic_shapes=dynamic_shapes)

在查看生成的程序之前,让我们先了解指定 dynamic_shapes 意味着什么,以及它如何与导出交互。对于指定了 Dim 对象的每个输入维度,都会 分配 一个符号,其范围为 [2, inf](为什么不是 [0, inf][1, inf]?我们稍后将在 0/1 特化部分解释)。

然后 Export 运行模型追踪,观察模型执行的每个操作。每个单独的操作都可以发出所谓的 “守卫” (guards);基本上是程序有效所必须满足的布尔条件。当守卫涉及到为输入维度分配的符号时,程序将包含对哪些输入形状有效的限制;即程序的动态行为。符号形状子系统负责接收所有发出的守卫,并生成一个符合所有这些守卫的最终程序表示。在我们看到 ExportedProgram 中的这种 “最终表示” 之前,让我们先看看我们正在追踪的玩具模型发出的守卫。

在这里,每个前向输入张量都标记有追踪开始时分配的符号

class DynamicModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.l = torch.nn.Linear(5, 3)

    def forward(
        self,
        w: torch.Tensor,  # [s0, s1]
        x: torch.Tensor,  # [s2]
        y: torch.Tensor,  # [s3, s4]
        z: torch.Tensor,  # [s5]
    ):
        x0 = x + y  # guard: s2 == s4
        x1 = self.l(w)  # guard: s1 == 5
        x2 = x0.flatten()  # no guard added here
        x3 = x2 + z  # guard: s3 * s4 == s5
        return x1, x3

让我们了解每个操作和发出的守卫

  • x0 = x + y:这是一个带有广播的逐元素加法,因为 x 是一个 1 维张量,而 y 是一个 2 维张量。x 沿着 y 的最后一个维度进行广播,发出守卫 s2 == s4

  • x1 = self.l(w):调用 nn.Linear() 与模型参数执行矩阵乘法。在导出中,参数、缓冲区和常量被认为是程序状态,即静态的,因此这是动态输入 (w: [s0, s1]) 与静态形状张量之间的 matmul。这会发出守卫 s1 == 5

  • x2 = x0.flatten():这次调用实际上没有发出任何守卫!(至少没有与输入形状相关的守卫)

  • x3 = x2 + z:展平后 x2 的形状为 [s3*s4],这个逐元素加法发出 s3 * s4 == s5

写下所有这些守卫并进行总结几乎就像数学证明,这正是符号形状子系统试图做的事情!总之,我们可以得出结论,程序必须具有以下输入形状才能有效

  • w: [s0, 5]

  • x: [s2]

  • y: [s3, s2]

  • z: [s2*s3]

当我们最终打印出导出的程序以查看结果时,这些形状就是我们在相应输入上看到的标注

print(ep)
ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, p_l_weight: "f32[3, 5]", p_l_bias: "f32[3]", w: "f32[s15, 5]", x: "f32[s77]", y: "f32[s17, s77]", z: "f32[s17*s77]"):
            # File: /var/lib/workspace/intermediate_source/torch_export_tutorial.py:268 in forward, code: x0 = x + y  # [8, 4]
            add: "f32[s17, s77]" = torch.ops.aten.add.Tensor(x, y);  x = y = None

            # File: /var/lib/ci-user/.local/lib/python3.10/site-packages/torch/nn/modules/linear.py:134 in forward, code: return F.linear(input, self.weight, self.bias)
            linear: "f32[s15, 3]" = torch.ops.aten.linear.default(w, p_l_weight, p_l_bias);  w = p_l_weight = p_l_bias = None

            # File: /var/lib/workspace/intermediate_source/torch_export_tutorial.py:270 in forward, code: x2 = x0.flatten()  # [32]
            flatten: "f32[s17*s77]" = torch.ops.aten.flatten.using_ints(add);  add = None

            # File: /var/lib/workspace/intermediate_source/torch_export_tutorial.py:271 in forward, code: x3 = x2 + z  # [32]
            add_1: "f32[s17*s77]" = torch.ops.aten.add.Tensor(flatten, z);  flatten = z = None
            return (linear, add_1)

Graph signature:
    # inputs
    p_l_weight: PARAMETER target='l.weight'
    p_l_bias: PARAMETER target='l.bias'
    w: USER_INPUT
    x: USER_INPUT
    y: USER_INPUT
    z: USER_INPUT

    # outputs
    linear: USER_OUTPUT
    add_1: USER_OUTPUT

Range constraints: {s15: VR[2, int_oo], s77: VR[2, int_oo], s17: VR[2, int_oo], s17*s77: VR[4, int_oo]}

另一个需要注意的特性是上面的 range_constraints 字段,它包含每个符号的有效范围。目前这并不是很有趣,因为导出调用没有发出任何与符号边界相关的守卫,并且每个基础符号都有一个通用边界,但这将在后面出现。

到目前为止,由于我们一直在导出这个玩具模型,这种体验并不能代表调试动态形状守卫和问题的典型难度。在大多数情况下,并不清楚正在发出哪些守卫,以及哪些操作和部分用户代码对此负责。对于这个玩具模型,我们可以精确锁定到行,而且守卫也非常直观。

在更复杂的情况下,有帮助的第一步通常是启用详细日志记录。这可以通过环境变量 TORCH_LOGS="+dynamic" 来完成,或者通过交互式方式使用 torch._logging.set_logs(dynamic=10)

torch._logging.set_logs(dynamic=10)
ep = export(model, (w, x, y, z), dynamic_shapes=dynamic_shapes)
I0708 23:34:22.403000 33372 torch/fx/experimental/symbolic_shapes.py:4058] create_env
I0708 23:34:22.405000 33372 torch/fx/experimental/symbolic_shapes.py:5743] create_symbol s15 = 6 for L['w'].size()[0] [2, int_oo] (_export/non_strict_utils.py:221 in fakify), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="s15" or to suppress this message run with TORCHDYNAMO_EXTENDED_ADVICE="0"
I0708 23:34:22.405000 33372 torch/fx/experimental/symbolic_shapes.py:5743] create_symbol s21 = 5 for L['w'].size()[1] [2, int_oo] (_export/non_strict_utils.py:221 in fakify), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="s21" or to suppress this message run with TORCHDYNAMO_EXTENDED_ADVICE="0"
I0708 23:34:22.407000 33372 torch/fx/experimental/symbolic_shapes.py:5743] create_symbol s77 = 4 for L['x'].size()[0] [2, int_oo] (_export/non_strict_utils.py:221 in fakify), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="s77" or to suppress this message run with TORCHDYNAMO_EXTENDED_ADVICE="0"
I0708 23:34:22.408000 33372 torch/fx/experimental/symbolic_shapes.py:5743] create_symbol s17 = 8 for L['y'].size()[0] [2, int_oo] (_export/non_strict_utils.py:221 in fakify), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="s17" or to suppress this message run with TORCHDYNAMO_EXTENDED_ADVICE="0"
I0708 23:34:22.409000 33372 torch/fx/experimental/symbolic_shapes.py:5743] create_symbol s94 = 4 for L['y'].size()[1] [2, int_oo] (_export/non_strict_utils.py:221 in fakify), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="s94" or to suppress this message run with TORCHDYNAMO_EXTENDED_ADVICE="0"
I0708 23:34:22.411000 33372 torch/fx/experimental/symbolic_shapes.py:5743] create_symbol s68 = 32 for L['z'].size()[0] [2, int_oo] (_export/non_strict_utils.py:221 in fakify), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="s68" or to suppress this message run with TORCHDYNAMO_EXTENDED_ADVICE="0"
V0708 23:34:22.417000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.418000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.418000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.419000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.420000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.421000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.421000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.422000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.423000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.424000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
I0708 23:34:22.426000 33372 torch/fx/experimental/symbolic_shapes.py:8081] eval Eq(s77, s94) [guard added] (_subclasses/fake_impls.py:1960 in infer_size), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED="Eq(s77, s94)"
I0708 23:34:22.427000 33372 torch/fx/experimental/symbolic_shapes.py:7674] set_replacement s94 = s77 (solve) VR[2, int_oo]
V0708 23:34:22.428000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
I0708 23:34:22.435000 33372 torch/fx/experimental/symbolic_shapes.py:8081] eval Eq(s21, 5) [guard added] (_meta_registrations.py:2518 in meta_mm), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED="Eq(s21, 5)"
V0708 23:34:22.435000 33372 torch/fx/experimental/symbolic_shapes.py:7503] _update_var_to_range s21 = VR[5, 5] (update)
I0708 23:34:22.436000 33372 torch/fx/experimental/symbolic_shapes.py:7674] set_replacement s21 = 5 (range_refined_to_singleton) VR[5, 5]
V0708 23:34:22.448000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.451000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
I0708 23:34:22.453000 33372 torch/fx/experimental/symbolic_shapes.py:8081] eval Eq(s17*s77, s68) [guard added] (_subclasses/fake_impls.py:1960 in infer_size), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED="Eq(s17*s77, s68)"
V0708 23:34:22.454000 33372 torch/fx/experimental/symbolic_shapes.py:7503] _update_var_to_range s68 = VR[4, int_oo] (update)
I0708 23:34:22.455000 33372 torch/fx/experimental/symbolic_shapes.py:7674] set_replacement s68 = s17*s77 (solve) VR[4, int_oo]
I0708 23:34:22.460000 33372 torch/fx/experimental/symbolic_shapes.py:5889] produce_guards
V0708 23:34:22.460000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['w'].size()[0] s15 None
V0708 23:34:22.460000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['w'].size()[1] 5 None
V0708 23:34:22.461000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['w'].stride()[0] 5 None
V0708 23:34:22.461000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['w'].stride()[1] 1 None
V0708 23:34:22.461000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['w'].storage_offset() 0 None
V0708 23:34:22.461000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].size()[0] s77 None
V0708 23:34:22.461000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].stride()[0] 1 None
V0708 23:34:22.461000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].storage_offset() 0 None
V0708 23:34:22.461000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].size()[0] s17 None
V0708 23:34:22.462000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].size()[1] s77 None
V0708 23:34:22.462000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].stride()[0] s77 None
V0708 23:34:22.462000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].stride()[1] 1 None
V0708 23:34:22.462000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].storage_offset() 0 None
V0708 23:34:22.462000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['z'].size()[0] s17*s77 None
V0708 23:34:22.462000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['z'].stride()[0] 1 None
V0708 23:34:22.463000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['z'].storage_offset() 0 None
V0708 23:34:22.474000 33372 torch/fx/experimental/symbolic_shapes.py:8352] eval 5 [trivial]

即使对于这个简单的玩具模型,它也会吐出相当多的内容。这里的日志行在前后已被截断以忽略不必要的信息,但通过查看日志,我们可以看到与我们上面描述的内容相关的行;例如符号的分配

"""
create_symbol s0 = 6 for L['w'].size()[0] [2, int_oo] (_dynamo/variables/builder.py:2841 in <lambda>)
create_symbol s1 = 5 for L['w'].size()[1] [2, int_oo] (_dynamo/variables/builder.py:2841 in <lambda>)
runtime_assert True == True [statically known]
create_symbol s2 = 4 for L['x'].size()[0] [2, int_oo] (_dynamo/variables/builder.py:2841 in <lambda>)
create_symbol s3 = 8 for L['y'].size()[0] [2, int_oo] (_dynamo/variables/builder.py:2841 in <lambda>)
create_symbol s4 = 4 for L['y'].size()[1] [2, int_oo] (_dynamo/variables/builder.py:2841 in <lambda>)
create_symbol s5 = 32 for L['z'].size()[0] [2, int_oo] (_dynamo/variables/builder.py:2841 in <lambda>)
"""
"\ncreate_symbol s0 = 6 for L['w'].size()[0] [2, int_oo] (_dynamo/variables/builder.py:2841 in <lambda>)\ncreate_symbol s1 = 5 for L['w'].size()[1] [2, int_oo] (_dynamo/variables/builder.py:2841 in <lambda>)\nruntime_assert True == True [statically known]\ncreate_symbol s2 = 4 for L['x'].size()[0] [2, int_oo] (_dynamo/variables/builder.py:2841 in <lambda>)\ncreate_symbol s3 = 8 for L['y'].size()[0] [2, int_oo] (_dynamo/variables/builder.py:2841 in <lambda>)\ncreate_symbol s4 = 4 for L['y'].size()[1] [2, int_oo] (_dynamo/variables/builder.py:2841 in <lambda>)\ncreate_symbol s5 = 32 for L['z'].size()[0] [2, int_oo] (_dynamo/variables/builder.py:2841 in <lambda>)\n"

带有 create_symbol 的行显示了新符号分配的时间,日志还识别了它们分配到的张量变量名称和维度。在其他行中,我们还可以看到发出的守卫

"""
runtime_assert Eq(s2, s4) [guard added] x0 = x + y  # output shape: [8, 4]  # dynamic_shapes_tutorial.py:16 in forward (_subclasses/fake_impls.py:845 in infer_size), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED="Eq(s2, s4)"
runtime_assert Eq(s1, 5) [guard added] x1 = self.l(w)  # [6, 3]  # dynamic_shapes_tutorial.py:17 in forward (_meta_registrations.py:2127 in meta_mm), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED="Eq(s1, 5)"
runtime_assert Eq(s2*s3, s5) [guard added] x3 = x2 + z  # [32]  # dynamic_shapes_tutorial.py:19 in forward (_subclasses/fake_impls.py:845 in infer_size), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED="Eq(s2*s3, s5)"
"""
'\nruntime_assert Eq(s2, s4) [guard added] x0 = x + y  # output shape: [8, 4]  # dynamic_shapes_tutorial.py:16 in forward (_subclasses/fake_impls.py:845 in infer_size), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED="Eq(s2, s4)"\nruntime_assert Eq(s1, 5) [guard added] x1 = self.l(w)  # [6, 3]  # dynamic_shapes_tutorial.py:17 in forward (_meta_registrations.py:2127 in meta_mm), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED="Eq(s1, 5)"\nruntime_assert Eq(s2*s3, s5) [guard added] x3 = x2 + z  # [32]  # dynamic_shapes_tutorial.py:19 in forward (_subclasses/fake_impls.py:845 in infer_size), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED="Eq(s2*s3, s5)"\n'

[guard added] 消息旁边,我们还能看到负责的用户代码行——幸运的是这里的模型足够简单。在许多现实世界的案例中,这并没有那么直接:高级 torch 操作可能有复杂的伪内核实现或运算符分解,从而使守卫发出的位置和内容变得复杂。在这些情况下,深入研究和调查的最佳方法是遵循日志建议,并使用环境变量 TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED="..." 重新运行,以进一步归因感兴趣的守卫。

Dim.AUTO 只是与 dynamic_shapes 交互的可用选项之一;在撰写本文时,还有另外 2 个选项可用:Dim.DYNAMICDim.STATICDim.STATIC 只是简单地将某个维度标记为静态,而 Dim.DYNAMIC 在各方面都类似于 Dim.AUTO,除了一点:它在特化为常量时会引发错误;这是为了保持动态性。例如,看看当在标记为动态的维度上发出静态守卫时会发生什么

dynamic_shapes["w"] = (Dim.AUTO, Dim.DYNAMIC)
try:
    export(model, (w, x, y, z), dynamic_shapes=dynamic_shapes)
except Exception:
    tb.print_exc()
I0708 23:34:22.479000 33372 torch/fx/experimental/symbolic_shapes.py:4058] create_env
I0708 23:34:22.480000 33372 torch/fx/experimental/symbolic_shapes.py:5743] create_symbol s15 = 6 for L['w'].size()[0] [2, int_oo] (_export/non_strict_utils.py:221 in fakify), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="s15" or to suppress this message run with TORCHDYNAMO_EXTENDED_ADVICE="0"
I0708 23:34:22.481000 33372 torch/fx/experimental/symbolic_shapes.py:5743] create_symbol s21 = 5 for L['w'].size()[1] [2, int_oo] (_export/non_strict_utils.py:221 in fakify), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="s21" or to suppress this message run with TORCHDYNAMO_EXTENDED_ADVICE="0"
I0708 23:34:22.483000 33372 torch/fx/experimental/symbolic_shapes.py:5743] create_symbol s77 = 4 for L['x'].size()[0] [2, int_oo] (_export/non_strict_utils.py:221 in fakify), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="s77" or to suppress this message run with TORCHDYNAMO_EXTENDED_ADVICE="0"
I0708 23:34:22.484000 33372 torch/fx/experimental/symbolic_shapes.py:5743] create_symbol s17 = 8 for L['y'].size()[0] [2, int_oo] (_export/non_strict_utils.py:221 in fakify), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="s17" or to suppress this message run with TORCHDYNAMO_EXTENDED_ADVICE="0"
I0708 23:34:22.485000 33372 torch/fx/experimental/symbolic_shapes.py:5743] create_symbol s94 = 4 for L['y'].size()[1] [2, int_oo] (_export/non_strict_utils.py:221 in fakify), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="s94" or to suppress this message run with TORCHDYNAMO_EXTENDED_ADVICE="0"
I0708 23:34:22.487000 33372 torch/fx/experimental/symbolic_shapes.py:5743] create_symbol s68 = 32 for L['z'].size()[0] [2, int_oo] (_export/non_strict_utils.py:221 in fakify), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="s68" or to suppress this message run with TORCHDYNAMO_EXTENDED_ADVICE="0"
V0708 23:34:22.493000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.493000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.494000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.495000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.496000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.497000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.498000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.498000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.499000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.500000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
I0708 23:34:22.502000 33372 torch/fx/experimental/symbolic_shapes.py:8081] eval Eq(s77, s94) [guard added] (_subclasses/fake_impls.py:1960 in infer_size), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED="Eq(s77, s94)"
I0708 23:34:22.503000 33372 torch/fx/experimental/symbolic_shapes.py:7674] set_replacement s94 = s77 (solve) VR[2, int_oo]
V0708 23:34:22.505000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
I0708 23:34:22.511000 33372 torch/fx/experimental/symbolic_shapes.py:8081] eval Eq(s21, 5) [guard added] (_meta_registrations.py:2518 in meta_mm), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED="Eq(s21, 5)"
V0708 23:34:22.512000 33372 torch/fx/experimental/symbolic_shapes.py:7503] _update_var_to_range s21 = VR[5, 5] (update)
I0708 23:34:22.512000 33372 torch/fx/experimental/symbolic_shapes.py:7674] set_replacement s21 = 5 (range_refined_to_singleton) VR[5, 5]
V0708 23:34:22.525000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.527000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
I0708 23:34:22.530000 33372 torch/fx/experimental/symbolic_shapes.py:8081] eval Eq(s17*s77, s68) [guard added] (_subclasses/fake_impls.py:1960 in infer_size), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED="Eq(s17*s77, s68)"
V0708 23:34:22.531000 33372 torch/fx/experimental/symbolic_shapes.py:7503] _update_var_to_range s68 = VR[4, int_oo] (update)
I0708 23:34:22.532000 33372 torch/fx/experimental/symbolic_shapes.py:7674] set_replacement s68 = s17*s77 (solve) VR[4, int_oo]
I0708 23:34:22.537000 33372 torch/fx/experimental/symbolic_shapes.py:5889] produce_guards
V0708 23:34:22.537000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['w'].size()[0] s15 None
V0708 23:34:22.537000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['w'].size()[1] 5 RelaxedUnspecConstraint(warn_only=False)
V0708 23:34:22.537000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['w'].stride()[0] 5 None
V0708 23:34:22.537000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['w'].stride()[1] 1 None
V0708 23:34:22.538000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['w'].storage_offset() 0 None
V0708 23:34:22.538000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].size()[0] s77 None
V0708 23:34:22.538000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].stride()[0] 1 None
V0708 23:34:22.538000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].storage_offset() 0 None
V0708 23:34:22.538000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].size()[0] s17 None
V0708 23:34:22.538000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].size()[1] s77 None
V0708 23:34:22.539000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].stride()[0] s77 None
V0708 23:34:22.539000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].stride()[1] 1 None
V0708 23:34:22.539000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].storage_offset() 0 None
V0708 23:34:22.539000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['z'].size()[0] s17*s77 None
V0708 23:34:22.539000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['z'].stride()[0] 1 None
V0708 23:34:22.539000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['z'].storage_offset() 0 None
Traceback (most recent call last):
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2039, in _export_to_aten_ir_make_fx
    produce_guards_callback(gm)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2197, in _produce_guards_callback
    return produce_guards_and_solve_constraints(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_export/non_strict_utils.py", line 625, in produce_guards_and_solve_constraints
    raise constraint_violation_error
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_export/non_strict_utils.py", line 584, in produce_guards_and_solve_constraints
    shape_env.produce_guards(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/symbolic_shapes.py", line 5851, in produce_guards
    return self.produce_guards_verbose(*args, **kwargs, langs=("python",))[0].exprs
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/symbolic_shapes.py", line 6731, in produce_guards_verbose
    raise ConstraintViolationError(
torch.fx.experimental.symbolic_shapes.ConstraintViolationError: Constraints violated (L['w'].size()[1])! For more information, run with TORCH_LOGS="+dynamic".
  - You marked L['w'].size()[1] as dynamic but your code specialized it to be a constant (5). If you're using mark_dynamic, either remove it or use maybe_mark_dynamic. If you're using Dim.DYNAMIC, replace it with either Dim.STATIC or Dim.AUTO.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/var/lib/workspace/intermediate_source/torch_export_tutorial.py", line 418, in <module>
    export(model, (w, x, y, z), dynamic_shapes=dynamic_shapes)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/__init__.py", line 205, in export
    raise e
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/__init__.py", line 171, in export
    return _export(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1344, in wrapper
    raise e
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1310, in wrapper
    ep = fn(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/exported_program.py", line 124, in wrapper
    return fn(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_utils_internal.py", line 96, in wrapper_function
    return function(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2516, in _export
    ep = _export_for_training(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1344, in wrapper
    raise e
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1310, in wrapper
    ep = fn(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/exported_program.py", line 124, in wrapper
    return fn(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2304, in _export_for_training
    export_artifact = export_func(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2233, in _non_strict_export
    aten_export_artifact = _to_aten_func(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2041, in _export_to_aten_ir_make_fx
    raise UserError(UserErrorType.CONSTRAINT_VIOLATION, str(e))  # noqa: B904
torch._dynamo.exc.UserError: Constraints violated (L['w'].size()[1])! For more information, run with TORCH_LOGS="+dynamic".
  - You marked L['w'].size()[1] as dynamic but your code specialized it to be a constant (5). If you're using mark_dynamic, either remove it or use maybe_mark_dynamic. If you're using Dim.DYNAMIC, replace it with either Dim.STATIC or Dim.AUTO.

The error above occurred when calling torch.export.export. If you would like to view some more information about this error, and get a list of all other errors that may occur in your export call, you can replace your `export()` call with `draft_export()`.

静态守卫并不总是模型固有的;它们也可以来自用户规范。事实上,导致形状特化的一个常见陷阱是用户为等效维度指定了冲突的标记;一个是动态的,另一个是静态的。当 x.shape[0]y.shape[1] 出现这种情况时,会引发相同类型的错误

dynamic_shapes["w"] = (Dim.AUTO, Dim.AUTO)
dynamic_shapes["x"] = (Dim.STATIC,)
dynamic_shapes["y"] = (Dim.AUTO, Dim.DYNAMIC)
try:
    export(model, (w, x, y, z), dynamic_shapes=dynamic_shapes)
except Exception:
    tb.print_exc()
I0708 23:34:22.551000 33372 torch/fx/experimental/symbolic_shapes.py:4058] create_env
I0708 23:34:22.552000 33372 torch/fx/experimental/symbolic_shapes.py:5743] create_symbol s15 = 6 for L['w'].size()[0] [2, int_oo] (_export/non_strict_utils.py:221 in fakify), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="s15" or to suppress this message run with TORCHDYNAMO_EXTENDED_ADVICE="0"
I0708 23:34:22.553000 33372 torch/fx/experimental/symbolic_shapes.py:5743] create_symbol s21 = 5 for L['w'].size()[1] [2, int_oo] (_export/non_strict_utils.py:221 in fakify), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="s21" or to suppress this message run with TORCHDYNAMO_EXTENDED_ADVICE="0"
I0708 23:34:22.555000 33372 torch/fx/experimental/symbolic_shapes.py:5743] create_symbol s17 = 8 for L['y'].size()[0] [2, int_oo] (_export/non_strict_utils.py:221 in fakify), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="s17" or to suppress this message run with TORCHDYNAMO_EXTENDED_ADVICE="0"
I0708 23:34:22.556000 33372 torch/fx/experimental/symbolic_shapes.py:5743] create_symbol s94 = 4 for L['y'].size()[1] [2, int_oo] (_export/non_strict_utils.py:221 in fakify), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="s94" or to suppress this message run with TORCHDYNAMO_EXTENDED_ADVICE="0"
I0708 23:34:22.558000 33372 torch/fx/experimental/symbolic_shapes.py:5743] create_symbol s68 = 32 for L['z'].size()[0] [2, int_oo] (_export/non_strict_utils.py:221 in fakify), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="s68" or to suppress this message run with TORCHDYNAMO_EXTENDED_ADVICE="0"
V0708 23:34:22.564000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.565000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.566000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.567000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.568000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.568000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.570000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.570000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
I0708 23:34:22.576000 33372 torch/fx/experimental/symbolic_shapes.py:8081] eval Eq(s94, 4) [guard added] (_subclasses/fake_impls.py:1960 in infer_size), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED="Eq(s94, 4)"
V0708 23:34:22.577000 33372 torch/fx/experimental/symbolic_shapes.py:7503] _update_var_to_range s94 = VR[4, 4] (update)
I0708 23:34:22.578000 33372 torch/fx/experimental/symbolic_shapes.py:7674] set_replacement s94 = 4 (range_refined_to_singleton) VR[4, 4]
I0708 23:34:22.585000 33372 torch/fx/experimental/symbolic_shapes.py:8081] eval Eq(s21, 5) [guard added] (_meta_registrations.py:2518 in meta_mm), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED="Eq(s21, 5)"
V0708 23:34:22.586000 33372 torch/fx/experimental/symbolic_shapes.py:7503] _update_var_to_range s21 = VR[5, 5] (update)
I0708 23:34:22.586000 33372 torch/fx/experimental/symbolic_shapes.py:7674] set_replacement s21 = 5 (range_refined_to_singleton) VR[5, 5]
I0708 23:34:22.607000 33372 torch/fx/experimental/symbolic_shapes.py:8081] eval Eq(4*s17, s68) [guard added] (_subclasses/fake_impls.py:1960 in infer_size), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED="Eq(4*s17, s68)"
V0708 23:34:22.609000 33372 torch/fx/experimental/symbolic_shapes.py:7503] _update_var_to_range s68 = VR[8, int_oo] (update)
I0708 23:34:22.611000 33372 torch/fx/experimental/symbolic_shapes.py:7674] set_replacement s68 = 4*s17 (solve) VR[8, int_oo]
I0708 23:34:22.616000 33372 torch/fx/experimental/symbolic_shapes.py:5889] produce_guards
V0708 23:34:22.616000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['w'].size()[0] s15 None
V0708 23:34:22.617000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['w'].size()[1] 5 None
V0708 23:34:22.617000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['w'].stride()[0] 5 None
V0708 23:34:22.617000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['w'].stride()[1] 1 None
V0708 23:34:22.617000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['w'].storage_offset() 0 None
V0708 23:34:22.617000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].size()[0] 4 None
V0708 23:34:22.617000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].stride()[0] 1 None
V0708 23:34:22.617000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].storage_offset() 0 None
V0708 23:34:22.618000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].size()[0] s17 None
V0708 23:34:22.618000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].size()[1] 4 RelaxedUnspecConstraint(warn_only=False)
V0708 23:34:22.618000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].stride()[0] 4 None
V0708 23:34:22.618000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].stride()[1] 1 None
V0708 23:34:22.618000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].storage_offset() 0 None
V0708 23:34:22.618000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['z'].size()[0] 4*s17 None
V0708 23:34:22.619000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['z'].stride()[0] 1 None
V0708 23:34:22.619000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['z'].storage_offset() 0 None
Traceback (most recent call last):
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2039, in _export_to_aten_ir_make_fx
    produce_guards_callback(gm)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2197, in _produce_guards_callback
    return produce_guards_and_solve_constraints(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_export/non_strict_utils.py", line 625, in produce_guards_and_solve_constraints
    raise constraint_violation_error
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_export/non_strict_utils.py", line 584, in produce_guards_and_solve_constraints
    shape_env.produce_guards(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/symbolic_shapes.py", line 5851, in produce_guards
    return self.produce_guards_verbose(*args, **kwargs, langs=("python",))[0].exprs
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/symbolic_shapes.py", line 6731, in produce_guards_verbose
    raise ConstraintViolationError(
torch.fx.experimental.symbolic_shapes.ConstraintViolationError: Constraints violated (L['y'].size()[1])! For more information, run with TORCH_LOGS="+dynamic".
  - You marked L['y'].size()[1] as dynamic but your code specialized it to be a constant (4). If you're using mark_dynamic, either remove it or use maybe_mark_dynamic. If you're using Dim.DYNAMIC, replace it with either Dim.STATIC or Dim.AUTO.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/var/lib/workspace/intermediate_source/torch_export_tutorial.py", line 431, in <module>
    export(model, (w, x, y, z), dynamic_shapes=dynamic_shapes)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/__init__.py", line 205, in export
    raise e
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/__init__.py", line 171, in export
    return _export(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1344, in wrapper
    raise e
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1310, in wrapper
    ep = fn(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/exported_program.py", line 124, in wrapper
    return fn(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_utils_internal.py", line 96, in wrapper_function
    return function(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2516, in _export
    ep = _export_for_training(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1344, in wrapper
    raise e
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1310, in wrapper
    ep = fn(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/exported_program.py", line 124, in wrapper
    return fn(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2304, in _export_for_training
    export_artifact = export_func(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2233, in _non_strict_export
    aten_export_artifact = _to_aten_func(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2041, in _export_to_aten_ir_make_fx
    raise UserError(UserErrorType.CONSTRAINT_VIOLATION, str(e))  # noqa: B904
torch._dynamo.exc.UserError: Constraints violated (L['y'].size()[1])! For more information, run with TORCH_LOGS="+dynamic".
  - You marked L['y'].size()[1] as dynamic but your code specialized it to be a constant (4). If you're using mark_dynamic, either remove it or use maybe_mark_dynamic. If you're using Dim.DYNAMIC, replace it with either Dim.STATIC or Dim.AUTO.

The error above occurred when calling torch.export.export. If you would like to view some more information about this error, and get a list of all other errors that may occur in your export call, you can replace your `export()` call with `draft_export()`.

在这里您可能会问,为什么导出要进行 “特化”,即为什么我们通过走静态路线来解决这种静态/动态冲突。答案是因为上述的符号形状系统,即符号和守卫。当 x.shape[0] 被标记为静态时,我们不分配符号,并以具体整数 4 来编译处理此形状。为 y.shape[1] 分配了一个符号,所以我们最终发出了守卫 s3 == 4,导致了特化。

导出的一个特性是,在追踪过程中,诸如断言、torch._check()if/else 条件之类的语句也会发出守卫。看看当我们用此类语句增强现有模型时会发生什么

class DynamicModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.l = torch.nn.Linear(5, 3)

    def forward(self, w, x, y, z):
        assert w.shape[0] <= 512
        torch._check(x.shape[0] >= 4)
        if w.shape[0] == x.shape[0] + 2:
            x0 = x + y
            x1 = self.l(w)
            x2 = x0.flatten()
            x3 = x2 + z
            return x1, x3
        else:
            return w

dynamic_shapes = {
    "w": (Dim.AUTO, Dim.AUTO),
    "x": (Dim.AUTO,),
    "y": (Dim.AUTO, Dim.AUTO),
    "z": (Dim.AUTO,),
}
try:
    ep = export(DynamicModel(), (w, x, y, z), dynamic_shapes=dynamic_shapes)
except Exception:
    tb.print_exc()
I0708 23:34:22.628000 33372 torch/fx/experimental/symbolic_shapes.py:4058] create_env
I0708 23:34:22.630000 33372 torch/fx/experimental/symbolic_shapes.py:5743] create_symbol s15 = 6 for L['w'].size()[0] [2, int_oo] (_export/non_strict_utils.py:221 in fakify), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="s15" or to suppress this message run with TORCHDYNAMO_EXTENDED_ADVICE="0"
I0708 23:34:22.630000 33372 torch/fx/experimental/symbolic_shapes.py:5743] create_symbol s21 = 5 for L['w'].size()[1] [2, int_oo] (_export/non_strict_utils.py:221 in fakify), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="s21" or to suppress this message run with TORCHDYNAMO_EXTENDED_ADVICE="0"
I0708 23:34:22.632000 33372 torch/fx/experimental/symbolic_shapes.py:5743] create_symbol s77 = 4 for L['x'].size()[0] [2, int_oo] (_export/non_strict_utils.py:221 in fakify), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="s77" or to suppress this message run with TORCHDYNAMO_EXTENDED_ADVICE="0"
I0708 23:34:22.634000 33372 torch/fx/experimental/symbolic_shapes.py:5743] create_symbol s17 = 8 for L['y'].size()[0] [2, int_oo] (_export/non_strict_utils.py:221 in fakify), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="s17" or to suppress this message run with TORCHDYNAMO_EXTENDED_ADVICE="0"
I0708 23:34:22.634000 33372 torch/fx/experimental/symbolic_shapes.py:5743] create_symbol s94 = 4 for L['y'].size()[1] [2, int_oo] (_export/non_strict_utils.py:221 in fakify), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="s94" or to suppress this message run with TORCHDYNAMO_EXTENDED_ADVICE="0"
I0708 23:34:22.636000 33372 torch/fx/experimental/symbolic_shapes.py:5743] create_symbol s68 = 32 for L['z'].size()[0] [2, int_oo] (_export/non_strict_utils.py:221 in fakify), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="s68" or to suppress this message run with TORCHDYNAMO_EXTENDED_ADVICE="0"
V0708 23:34:22.642000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.643000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.643000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.644000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.645000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.646000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.646000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.647000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.648000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.649000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
I0708 23:34:22.654000 33372 torch/fx/experimental/symbolic_shapes.py:8081] eval s15 <= 512 [guard added] (workspace/intermediate_source/torch_export_tutorial.py:450 in forward), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED="s15 <= 512"
V0708 23:34:22.655000 33372 torch/fx/experimental/symbolic_shapes.py:7503] _update_var_to_range s15 = VR[2, 512] (update)
I0708 23:34:22.658000 33372 torch/fx/experimental/symbolic_shapes.py:8081] eval s77 >= 4 [guard added] (workspace/intermediate_source/torch_export_tutorial.py:451 in forward), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED="s77 >= 4"
V0708 23:34:22.658000 33372 torch/fx/experimental/symbolic_shapes.py:7503] _update_var_to_range s77 = VR[4, int_oo] (update)
I0708 23:34:22.663000 33372 torch/fx/experimental/symbolic_shapes.py:8081] eval Eq(s15, s77 + 2) [guard added] (workspace/intermediate_source/torch_export_tutorial.py:452 in forward), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED="Eq(s15, s77 + 2)"
V0708 23:34:22.666000 33372 torch/fx/experimental/symbolic_shapes.py:7503] _update_var_to_range s77 = VR[4, 510] (update)
V0708 23:34:22.666000 33372 torch/fx/experimental/symbolic_shapes.py:7503] _update_var_to_range s15 = VR[6, 512] (update)
I0708 23:34:22.667000 33372 torch/fx/experimental/symbolic_shapes.py:7674] set_replacement s15 = s77 + 2 (solve) VR[6, 512]
I0708 23:34:22.671000 33372 torch/fx/experimental/symbolic_shapes.py:8081] eval Eq(s77, s94) [guard added] (_subclasses/fake_impls.py:1960 in infer_size), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED="Eq(s77, s94)"
V0708 23:34:22.671000 33372 torch/fx/experimental/symbolic_shapes.py:7503] _update_var_to_range s94 = VR[4, 510] (update)
I0708 23:34:22.672000 33372 torch/fx/experimental/symbolic_shapes.py:7674] set_replacement s94 = s77 (solve) VR[4, 510]
V0708 23:34:22.675000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
I0708 23:34:22.682000 33372 torch/fx/experimental/symbolic_shapes.py:8081] eval Eq(s21, 5) [guard added] (_meta_registrations.py:2518 in meta_mm), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED="Eq(s21, 5)"
V0708 23:34:22.682000 33372 torch/fx/experimental/symbolic_shapes.py:7503] _update_var_to_range s21 = VR[5, 5] (update)
I0708 23:34:22.683000 33372 torch/fx/experimental/symbolic_shapes.py:7674] set_replacement s21 = 5 (range_refined_to_singleton) VR[5, 5]
V0708 23:34:22.698000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.701000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
I0708 23:34:22.709000 33372 torch/fx/experimental/symbolic_shapes.py:8081] eval Eq(s17*s77, s68) [guard added] (_subclasses/fake_impls.py:1960 in infer_size), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED="Eq(s17*s77, s68)"
V0708 23:34:22.710000 33372 torch/fx/experimental/symbolic_shapes.py:7503] _update_var_to_range s68 = VR[8, int_oo] (update)
I0708 23:34:22.711000 33372 torch/fx/experimental/symbolic_shapes.py:7674] set_replacement s68 = s17*s77 (solve) VR[8, int_oo]
I0708 23:34:22.716000 33372 torch/fx/experimental/symbolic_shapes.py:5889] produce_guards
V0708 23:34:22.717000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['w'].size()[0] s77 + 2 None
V0708 23:34:22.717000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['w'].size()[1] 5 None
V0708 23:34:22.717000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['w'].stride()[0] 5 None
V0708 23:34:22.717000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['w'].stride()[1] 1 None
V0708 23:34:22.717000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['w'].storage_offset() 0 None
V0708 23:34:22.718000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].size()[0] s77 None
V0708 23:34:22.718000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].stride()[0] 1 None
V0708 23:34:22.718000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].storage_offset() 0 None
V0708 23:34:22.718000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].size()[0] s17 None
V0708 23:34:22.718000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].size()[1] s77 None
V0708 23:34:22.718000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].stride()[0] s77 None
V0708 23:34:22.719000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].stride()[1] 1 None
V0708 23:34:22.719000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].storage_offset() 0 None
V0708 23:34:22.719000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['z'].size()[0] s17*s77 None
V0708 23:34:22.719000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['z'].stride()[0] 1 None
V0708 23:34:22.719000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['z'].storage_offset() 0 None
V0708 23:34:22.735000 33372 torch/fx/experimental/symbolic_shapes.py:8352] eval 5 [trivial]

每个语句都会发出一个额外的守卫,导出的程序显示了这些变化;为了支持 s2 + 2 而消除了 s0,并且 s2 现在包含下限和上限,这反映在 range_constraints 中。

对于 if/else 条件,您可能会问为什么选择了 True 分支,以及为什么追踪没有发出 w.shape[0] != x.shape[0] + 2 的守卫。答案是导出由追踪提供的示例输入引导,并在所取分支上进行特化。如果提供的示例输入形状未能通过 if 条件,则导出将追踪并发出与 else 分支对应的守卫。此外,您可能会问为什么我们只追踪了 if 分支,以及是否可以在程序中维护控制流并保持两个分支都处于激活状态。为此,请参阅上面的 “控制流操作” 部分重新编写您的模型代码。

0/1 特化#

既然我们在讨论守卫和特化,现在是谈论我们之前提到的 0/1 特化问题的好时机。其底线是,导出将在值为 0 或 1 的样本输入维度上进行特化,因为这些形状具有在追踪时不具有通用性的属性。例如,大小为 1 的张量可以进行广播,而其他大小则不行;大小为 0 则……这只意味着,当您希望程序将其硬编码时,应指定 0/1 的样本输入;当希望具有动态行为时,应指定非 0/1 的样本输入。看看当我们导出这个线性层时,在运行时会发生什么

ep = export(
    torch.nn.Linear(4, 3),
    (torch.randn(1, 4),),
    dynamic_shapes={
        "input": (Dim.AUTO, Dim.STATIC),
    },
)
try:
    ep.module()(torch.randn(2, 4))
except Exception:
    tb.print_exc()
I0708 23:34:22.740000 33372 torch/fx/experimental/symbolic_shapes.py:4058] create_env
I0708 23:34:22.753000 33372 torch/fx/experimental/symbolic_shapes.py:5889] produce_guards
V0708 23:34:22.753000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['input'].size()[0] 1 None
V0708 23:34:22.753000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['input'].size()[1] 4 None
V0708 23:34:22.753000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['input'].stride()[0] 4 None
V0708 23:34:22.753000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['input'].stride()[1] 1 None
V0708 23:34:22.753000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['input'].storage_offset() 0 None
W0708 23:34:22.756000 33372 torch/_export/non_strict_utils.py:656] dimension inputs['input'].shape[0] 0/1 specialized; Dim.AUTO was specified along with a sample input with hint = 1.
Traceback (most recent call last):
  File "/var/lib/workspace/intermediate_source/torch_export_tutorial.py", line 500, in <module>
    ep.module()(torch.randn(2, 4))
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/graph_module.py", line 1000, in call_wrapped
    return self._wrapped_call(self, *args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/graph_module.py", line 507, in __call__
    raise e
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/graph_module.py", line 493, in __call__
    return super(self.cls, obj).__call__(*args, **kwargs)  # type: ignore[misc]
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1884, in _call_impl
    return inner()
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1832, in inner
    result = forward_call(*args, **kwargs)
  File "<eval_with_key>.163", line 9, in forward
    _guards_fn = self._guards_fn(input_1);  _guards_fn = None
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
    return forward_call(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_dynamo/external_utils.py", line 215, in inner
    return func(*args, **kwargs)
  File "<string>", line 3, in _
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/__init__.py", line 2279, in _assert
    raise AssertionError(message)
AssertionError: Guard failed: input.size()[0] == 1

命名维度 (Named Dims)#

到目前为止,我们只讨论了 3 种指定动态形状的方法:Dim.AUTODim.DYNAMICDim.STATIC。它们的吸引力在于低摩擦的用户体验;在模型追踪过程中发出的所有守卫都会被遵守,并且像最小值/最大值范围、关系以及静态/动态维度之类的动态行为都会在导出底层自动确定。动态形状子系统本质上充当了一个 “发现” 过程,总结这些守卫并呈现导出认为的程序整体动态行为。这种设计的缺点在用户对这些模型的动态行为有更强的期望或信念时显现出来——也许是对动态性有强烈的渴望,必须不惜一切代价避免特定维度上的特化;或者也许我们只是想捕捉随着原始模型代码、底层分解或元内核的变化而发生的动态行为变化。这些变化将无法被检测到,且 export() 调用极有可能成功,除非已有测试检查产生的 ExportedProgram 表示。

对于此类情况,我们的立场是推荐 “传统” 的指定动态形状的方式,导出的长期用户可能对此比较熟悉:命名的 Dims

dx = Dim("dx", min=4, max=256)
dh = Dim("dh", max=512)
dynamic_shapes = {
    "x": (dx, None),
    "y": (2 * dx, dh),
}

这种风格的动态形状允许用户指定为输入维度分配了哪些符号、这些符号的最小/最大边界,并对生成的 ExportedProgram 的动态行为进行限制;如果模型追踪发出的守卫与给定的关系或静态/动态规范发生冲突,将引发 ConstraintViolation 错误。例如,在上述规范中,断言了以下内容

  • x.shape[0] 的范围为 [4, 256],且通过 y.shape[0] == 2 * x.shape[0]y.shape[0] 关联。

  • x.shape[1] 是静态的。

  • y.shape[1] 的范围为 [2, 512],且与其他维度均无关联。

在此设计中,我们允许使用单变量线性表达式指定维度之间的关系:可以为任何维度指定 A * dim + B。这允许用户为动态维度指定更复杂的约束,如整数可除性

dx = Dim("dx", min=4, max=512)
dynamic_shapes = {
    "x": (4 * dx, None)  # x.shape[0] has range [16, 2048], and is divisible by 4.
}

违反约束及其建议修复#

这种规范风格(在 Dim.AUTO 引入之前)的一个常见问题是,规范通常与模型追踪产生的结果不匹配。这将导致 ConstraintViolation 错误并导出建议的修复——例如,在这个模型和规范中,模型本身要求 xy 的维度 0 之间具有相等性,并且要求维度 1 是静态的。

class Foo(torch.nn.Module):
    def forward(self, x, y):
        w = x + y
        return w + torch.ones(4)

dx, dy, d1 = torch.export.dims("dx", "dy", "d1")
try:
    ep = export(
        Foo(),
        (torch.randn(6, 4), torch.randn(6, 4)),
        dynamic_shapes={
            "x": (dx, d1),
            "y": (dy, d1),
        },
    )
except Exception:
    tb.print_exc()
I0708 23:34:22.765000 33372 torch/fx/experimental/symbolic_shapes.py:4058] create_env
I0708 23:34:22.767000 33372 torch/fx/experimental/symbolic_shapes.py:5743] create_symbol s77 = 6 for L['x'].size()[0] [2, int_oo] (_export/non_strict_utils.py:221 in fakify), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="s77" or to suppress this message run with TORCHDYNAMO_EXTENDED_ADVICE="0"
I0708 23:34:22.768000 33372 torch/fx/experimental/symbolic_shapes.py:5743] create_symbol s27 = 4 for L['x'].size()[1] [2, int_oo] (_export/non_strict_utils.py:221 in fakify), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="s27" or to suppress this message run with TORCHDYNAMO_EXTENDED_ADVICE="0"
I0708 23:34:22.771000 33372 torch/fx/experimental/symbolic_shapes.py:5743] create_symbol s17 = 6 for L['y'].size()[0] [2, int_oo] (_export/non_strict_utils.py:221 in fakify), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="s17" or to suppress this message run with TORCHDYNAMO_EXTENDED_ADVICE="0"
I0708 23:34:22.772000 33372 torch/fx/experimental/symbolic_shapes.py:5743] create_symbol s94 = 4 for L['y'].size()[1] [2, int_oo] (_export/non_strict_utils.py:221 in fakify), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="s94" or to suppress this message run with TORCHDYNAMO_EXTENDED_ADVICE="0"
V0708 23:34:22.778000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.779000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.780000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.782000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.782000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.783000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
I0708 23:34:22.787000 33372 torch/fx/experimental/symbolic_shapes.py:8081] eval Eq(s27, s94) [guard added] (_subclasses/fake_impls.py:1960 in infer_size), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED="Eq(s27, s94)"
I0708 23:34:22.789000 33372 torch/fx/experimental/symbolic_shapes.py:7674] set_replacement s94 = s27 (solve) VR[2, int_oo]
I0708 23:34:22.790000 33372 torch/fx/experimental/symbolic_shapes.py:8081] eval Eq(s77, s17) [guard added] (_subclasses/fake_impls.py:1960 in infer_size), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED="Eq(s77, s17)"
I0708 23:34:22.791000 33372 torch/fx/experimental/symbolic_shapes.py:7674] set_replacement s77 = s17 (solve) VR[2, int_oo]
V0708 23:34:22.794000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
I0708 23:34:22.803000 33372 torch/fx/experimental/symbolic_shapes.py:8081] eval Eq(s27, 4) [guard added] (_subclasses/fake_impls.py:1960 in infer_size), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED="Eq(s27, 4)"
V0708 23:34:22.804000 33372 torch/fx/experimental/symbolic_shapes.py:7503] _update_var_to_range s27 = VR[4, 4] (update)
I0708 23:34:22.804000 33372 torch/fx/experimental/symbolic_shapes.py:7674] set_replacement s27 = 4 (range_refined_to_singleton) VR[4, 4]
I0708 23:34:22.810000 33372 torch/fx/experimental/symbolic_shapes.py:5889] produce_guards
V0708 23:34:22.811000 33372 torch/fx/experimental/symbolic_shapes.py:7503] _update_var_to_range s94 = VR[4, 4] (update)
I0708 23:34:22.811000 33372 torch/fx/experimental/symbolic_shapes.py:7674] set_replacement s94 = 4 (find) VR[4, 4]
V0708 23:34:22.811000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].size()[0] s17 StrictMinMaxConstraint(warn_only=False, vr=VR[0, int_oo])
V0708 23:34:22.811000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].size()[1] 4 StrictMinMaxConstraint(warn_only=False, vr=VR[0, int_oo])
V0708 23:34:22.811000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].stride()[0] 4 None
V0708 23:34:22.812000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].stride()[1] 1 None
V0708 23:34:22.812000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].storage_offset() 0 None
V0708 23:34:22.812000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].size()[0] s17 StrictMinMaxConstraint(warn_only=False, vr=VR[0, int_oo])
V0708 23:34:22.812000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].size()[1] 4 StrictMinMaxConstraint(warn_only=False, vr=VR[0, int_oo])
V0708 23:34:22.812000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].stride()[0] 4 None
V0708 23:34:22.812000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].stride()[1] 1 None
V0708 23:34:22.812000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].storage_offset() 0 None
Traceback (most recent call last):
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2039, in _export_to_aten_ir_make_fx
    produce_guards_callback(gm)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2197, in _produce_guards_callback
    return produce_guards_and_solve_constraints(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_export/non_strict_utils.py", line 625, in produce_guards_and_solve_constraints
    raise constraint_violation_error
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_export/non_strict_utils.py", line 584, in produce_guards_and_solve_constraints
    shape_env.produce_guards(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/symbolic_shapes.py", line 5851, in produce_guards
    return self.produce_guards_verbose(*args, **kwargs, langs=("python",))[0].exprs
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/symbolic_shapes.py", line 6731, in produce_guards_verbose
    raise ConstraintViolationError(
torch.fx.experimental.symbolic_shapes.ConstraintViolationError: Constraints violated (d1, dy)! For more information, run with TORCH_LOGS="+dynamic".
  - You marked d1 as dynamic but your code specialized it to be a constant (4). If you're using mark_dynamic, either remove it or use maybe_mark_dynamic. If you're using Dim.DYNAMIC, replace it with either Dim.STATIC or Dim.AUTO.
  - You marked d1 as dynamic but your code specialized it to be a constant (4). If you're using mark_dynamic, either remove it or use maybe_mark_dynamic. If you're using Dim.DYNAMIC, replace it with either Dim.STATIC or Dim.AUTO.
  - The values of dy = L['y'].size()[0] and dx = L['x'].size()[0] must always be equal.
Suggested fixes:
  d1 = 4
  dy = dx

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/var/lib/workspace/intermediate_source/torch_export_tutorial.py", line 557, in <module>
    ep = export(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/__init__.py", line 205, in export
    raise e
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/__init__.py", line 171, in export
    return _export(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1344, in wrapper
    raise e
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1310, in wrapper
    ep = fn(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/exported_program.py", line 124, in wrapper
    return fn(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_utils_internal.py", line 96, in wrapper_function
    return function(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2516, in _export
    ep = _export_for_training(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1344, in wrapper
    raise e
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1310, in wrapper
    ep = fn(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/exported_program.py", line 124, in wrapper
    return fn(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2304, in _export_for_training
    export_artifact = export_func(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2233, in _non_strict_export
    aten_export_artifact = _to_aten_func(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2041, in _export_to_aten_ir_make_fx
    raise UserError(UserErrorType.CONSTRAINT_VIOLATION, str(e))  # noqa: B904
torch._dynamo.exc.UserError: Constraints violated (d1, dy)! For more information, run with TORCH_LOGS="+dynamic".
  - You marked d1 as dynamic but your code specialized it to be a constant (4). If you're using mark_dynamic, either remove it or use maybe_mark_dynamic. If you're using Dim.DYNAMIC, replace it with either Dim.STATIC or Dim.AUTO.
  - You marked d1 as dynamic but your code specialized it to be a constant (4). If you're using mark_dynamic, either remove it or use maybe_mark_dynamic. If you're using Dim.DYNAMIC, replace it with either Dim.STATIC or Dim.AUTO.
  - The values of dy = L['y'].size()[0] and dx = L['x'].size()[0] must always be equal.
Suggested fixes:
  d1 = 4
  dy = dx

The error above occurred when calling torch.export.export. If you would like to view some more information about this error, and get a list of all other errors that may occur in your export call, you can replace your `export()` call with `draft_export()`.

对建议修复的预期是,用户可以交互式地将更改复制粘贴到其动态形状规范中,然后成功导出。

最后,关于规范选项,有一些值得了解的事项

  • None 是静态行为的一个很好的选项: - dynamic_shapes=None(默认)在整个模型都是静态的情况下导出。 - 在输入级别指定 None 会在所有张量维度都是静态的情况下导出,且对于非张量输入也是必需的。 - 在维度级别指定 None 会对该维度进行特化,尽管这已被弃用,转而支持 Dim.STATIC

  • 指定逐维度的整数值也会产生静态行为,并且还将另外检查提供的样本输入是否符合规范。

这些选项组合在下面的输入和动态形状规范中

inputs = (
    torch.randn(4, 4),
    torch.randn(3, 3),
    16,
    False,
)
dynamic_shapes = {
    "tensor_0": (Dim.AUTO, None),
    "tensor_1": None,
    "int_val": None,
    "bool_val": None,
}

数据依赖错误#

在尝试导出模型时,您可能遇到过类似 “Could not guard on data-dependent expression” 或 “Could not extract specialized integer from data-dependent expression” 之类的错误。这些错误的存在是因为 torch.export() 使用 FakeTensors 来编译程序,这些张量在符号上代表其实际张量对应物。虽然它们具有等效的符号属性(例如大小、步长、数据类型),但不同之处在于 FakeTensors 不包含任何数据值。虽然这避免了不必要的内存使用和昂贵的计算,但它确实意味着导出可能无法开箱即用地编译用户代码中依赖于数据值的部分。简而言之,如果编译器需要一个具体的、依赖于数据的值才能继续进行,它将报错,抱怨该值不可用。

数据依赖的值出现在许多地方,常见的来源是像 item()tolist()torch.unbind() 这样从张量中提取标量值的调用。这些值在导出的程序中是如何表示的?在 约束/动态形状 部分中,我们讨论了分配符号来表示动态输入维度。这里也是一样:我们为程序中出现的每个数据依赖值分配符号。重要的区别在于,这些是 “未支持的” (unbacked) 符号,与为输入维度分配的 “已支持的” (backed) 符号相反。“已支持/未支持” 这一术语指的是符号是否存在 “提示” (hint):一个支持该符号的具体数值,它可以告知编译器如何继续。

在输入形状符号(支持符号)的情况下,这些提示就是提供的样本输入形状,这解释了为什么控制流分支由样本输入属性决定。对于数据依赖的值,符号是在追踪期间从 FakeTensor “数据” 中获取的,因此编译器不知道这些符号将采用的实际值(提示)。

让我们看看这些内容是如何在导出的程序中显示的

class Foo(torch.nn.Module):
    def forward(self, x, y):
        a = x.item()
        b = y.tolist()
        return b + [a]

inps = (
    torch.tensor(1),
    torch.tensor([2, 3]),
)
ep = export(Foo(), inps)
print(ep)
I0708 23:34:22.822000 33372 torch/fx/experimental/symbolic_shapes.py:4058] create_env
I0708 23:34:22.828000 33372 torch/fx/experimental/symbolic_shapes.py:5330] create_unbacked_symint u0 [-int_oo, int_oo] (_subclasses/fake_impls.py:1233 in local_scalar_dense)
I0708 23:34:22.828000 33372 torch/fx/experimental/symbolic_shapes.py:1448] compute_unbacked_bindings [u0]
I0708 23:34:22.833000 33372 torch/fx/experimental/symbolic_shapes.py:5330] create_unbacked_symint u1 [-int_oo, int_oo] (_subclasses/fake_impls.py:1233 in local_scalar_dense)
I0708 23:34:22.833000 33372 torch/fx/experimental/symbolic_shapes.py:1448] compute_unbacked_bindings [u1]
I0708 23:34:22.834000 33372 torch/fx/experimental/symbolic_shapes.py:5330] create_unbacked_symint u2 [-int_oo, int_oo] (_subclasses/fake_impls.py:1233 in local_scalar_dense)
I0708 23:34:22.834000 33372 torch/fx/experimental/symbolic_shapes.py:1448] compute_unbacked_bindings [u2]
I0708 23:34:22.836000 33372 torch/fx/experimental/symbolic_shapes.py:5889] produce_guards
V0708 23:34:22.837000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].storage_offset() 0 None
V0708 23:34:22.837000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].size()[0] 2 None
V0708 23:34:22.837000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].stride()[0] 1 None
V0708 23:34:22.837000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].storage_offset() 0 None
ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, x: "i64[]", y: "i64[2]"):
            # File: /var/lib/workspace/intermediate_source/torch_export_tutorial.py:618 in forward, code: a = x.item()
            item: "Sym(u0)" = torch.ops.aten.item.default(x);  x = None

            # File: /var/lib/workspace/intermediate_source/torch_export_tutorial.py:619 in forward, code: b = y.tolist()
            unbind = torch.ops.aten.unbind.int(y);  y = None
            getitem: "i64[]" = unbind[0]
            getitem_1: "i64[]" = unbind[1];  unbind = None
            item_1: "Sym(u1)" = torch.ops.aten.item.default(getitem);  getitem = None
            item_2: "Sym(u2)" = torch.ops.aten.item.default(getitem_1);  getitem_1 = None
            return (item_1, item_2, item)

Graph signature:
    # inputs
    x: USER_INPUT
    y: USER_INPUT

    # outputs
    item_1: USER_OUTPUT
    item_2: USER_OUTPUT
    item: USER_OUTPUT

Range constraints: {u0: VR[-int_oo, int_oo], u1: VR[-int_oo, int_oo], u2: VR[-int_oo, int_oo]}

结果是分配并返回了 3 个未支持符号(注意它们的前缀是 “u”,而不是用于输入形状/已支持符号的常用 “s”):1 个用于 item() 调用,而 tolist() 调用的 y 中的每个元素各分配 1 个。请注意,从范围约束字段来看,这些符号的范围为 [-int_oo, int_oo],而不是分配给输入形状符号的默认 [0, int_oo] 范围,因为我们没有关于这些值是什么的信息——它们不代表大小,因此不一定具有正值。

守卫、torch._check()#

但上面的情况很容易导出,因为这些符号的具体数值并未用于编译器的任何决策;相关的只是返回值是未支持符号。本节强调的数据依赖错误是类似以下的情况,即遇到了 数据依赖的守卫

class Foo(torch.nn.Module):
    def forward(self, x, y):
        a = x.item()
        if a // 2 >= 5:
            return y + 2
        else:
            return y * 5

在这里,我们实际上需要 a 的 “提示” 或具体值,以便编译器决定是将 return y + 2 还是 return y * 5 追踪为输出。由于我们使用 FakeTensors 进行追踪,我们不知道 a // 2 >= 5 实际上评估为什么,导出报错 “Could not guard on data-dependent expression u0 // 2 >= 5 (unhinted)”。

那么我们要如何导出这个玩具模型呢?与 torch.compile() 不同,导出需要完整的图编译,我们不能只是对其进行图断裂。以下是一些基本选项:

  1. 手动特化:我们可以干预选择追踪的分支,要么通过移除控制流代码以仅包含特化分支,要么使用 torch.compiler.is_compiling() 来守卫编译时追踪的内容。

  2. torch.cond():我们可以使用 torch.cond() 重写控制流代码,这样我们就不会在某个分支上特化。

虽然这些选项是有效的,但它们有其陷阱。选项 1 有时需要对模型代码进行剧烈的、侵入性的重写以进行特化,而 torch.cond() 并非处理数据依赖错误的全面系统。正如我们将看到的,有些数据依赖错误并不涉及控制流。

通常推荐的方法是从 torch._check() 调用开始。虽然这些给人的印象纯粹是断言语句,但它们实际上是一套告知编译器符号属性的系统。虽然 torch._check() 调用在运行时确实起到了断言作用,但在编译时被追踪时,检查过的表达式会被发送到符号形状子系统进行推理,并且从表达式为真推导出的任何符号属性都将被存储为符号属性(前提是它足够聪明,能推导出这些属性)。所以即使未支持符号没有提示,如果我们能够通过 torch._check() 调用传达这些符号普遍成立的属性,我们就有可能绕过数据依赖的守卫,而无需重写违规的模型代码。

例如在上面的模型中,插入 torch._check(a >= 10) 会告知编译器始终可以返回 y + 2,而 torch._check(a == 4) 告知其返回 y * 5。看看我们重新导出此模型时会发生什么。

class Foo(torch.nn.Module):
    def forward(self, x, y):
        a = x.item()
        torch._check(a >= 10)
        torch._check(a <= 60)
        if a // 2 >= 5:
            return y + 2
        else:
            return y * 5

inps = (
    torch.tensor(32),
    torch.randn(4),
)
ep = export(Foo(), inps)
print(ep)
I0708 23:34:22.844000 33372 torch/fx/experimental/symbolic_shapes.py:4058] create_env
I0708 23:34:22.850000 33372 torch/fx/experimental/symbolic_shapes.py:5330] create_unbacked_symint u0 [-int_oo, int_oo] (_subclasses/fake_impls.py:1233 in local_scalar_dense)
I0708 23:34:22.850000 33372 torch/fx/experimental/symbolic_shapes.py:1448] compute_unbacked_bindings [u0]
I0708 23:34:22.852000 33372 torch/fx/experimental/symbolic_shapes.py:8081] runtime_assert u0 >= 10 [guard added] (workspace/intermediate_source/torch_export_tutorial.py:673 in forward), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED="u0 >= 10"
V0708 23:34:22.853000 33372 torch/fx/experimental/symbolic_shapes.py:7503] _update_var_to_range u0 = VR[10, int_oo] (update)
I0708 23:34:22.857000 33372 torch/fx/experimental/symbolic_shapes.py:8081] runtime_assert u0 <= 60 [guard added] (workspace/intermediate_source/torch_export_tutorial.py:674 in forward), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED="u0 <= 60"
V0708 23:34:22.858000 33372 torch/fx/experimental/symbolic_shapes.py:7503] _update_var_to_range u0 = VR[10, 60] (update)
V0708 23:34:22.863000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == True [statically known]
I0708 23:34:22.867000 33372 torch/fx/experimental/symbolic_shapes.py:5889] produce_guards
V0708 23:34:22.867000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].storage_offset() 0 None
V0708 23:34:22.867000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].size()[0] 4 None
V0708 23:34:22.867000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].stride()[0] 1 None
V0708 23:34:22.868000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].storage_offset() 0 None
V0708 23:34:22.869000 33372 torch/fx/experimental/symbolic_shapes.py:8616] runtime_assert u0 >= 10 == True [statically known]
V0708 23:34:22.870000 33372 torch/fx/experimental/symbolic_shapes.py:8616] runtime_assert u0 <= 60 == True [statically known]
ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, x: "i64[]", y: "f32[4]"):
            # File: /var/lib/workspace/intermediate_source/torch_export_tutorial.py:672 in forward, code: a = x.item()
            item: "Sym(u0)" = torch.ops.aten.item.default(x);  x = None
            ge_2: "Sym(u0 >= 10)" = item >= 10
            _assert_scalar_default = torch.ops.aten._assert_scalar.default(ge_2, "Runtime assertion failed for expression u0 >= 10 on node 'ge_2'");  ge_2 = _assert_scalar_default = None
            le_1: "Sym(u0 <= 60)" = item <= 60;  item = None
            _assert_scalar_default_1 = torch.ops.aten._assert_scalar.default(le_1, "Runtime assertion failed for expression u0 <= 60 on node 'le_1'");  le_1 = _assert_scalar_default_1 = None

            # File: /var/lib/workspace/intermediate_source/torch_export_tutorial.py:676 in forward, code: return y + 2
            add: "f32[4]" = torch.ops.aten.add.Tensor(y, 2);  y = None
            return (add,)

Graph signature:
    # inputs
    x: USER_INPUT
    y: USER_INPUT

    # outputs
    add: USER_OUTPUT

Range constraints: {u0: VR[10, 60]}

导出成功,并从范围约束字段中注意到 u0 的范围为 [10, 60]

那么 torch._check() 调用实际上沟通了哪些信息?随着符号形状子系统变得越来越聪明,这些信息会有所不同,但在基本层面上,这些通常都是正确的

  1. 与非数据依赖表达式的相等性:传达相等性的 torch._check() 调用,例如 u0 == s0 + 4u0 == 5

  2. 范围细化:如上所述,为符号提供下限或上限的调用。

  3. 围绕更复杂表达式的一些基本推理:插入 torch._check(a < 4) 通常会告知编译器 a >= 4 为假。对复杂表达式(如 torch._check(a ** 2 - 3 * a <= 10))的检查通常可以使您通过相同的守卫。

如前所述,torch._check() 调用在数据依赖控制流之外也有适用性。例如,在下面这个模型中,插入 torch._check() 盛行,而手动特化和 torch.cond() 则不行

class Foo(torch.nn.Module):
    def forward(self, x, y):
        a = x.item()
        return y[a]

inps = (
    torch.tensor(32),
    torch.randn(60),
)
try:
    export(Foo(), inps)
except Exception:
    tb.print_exc()
I0708 23:34:22.877000 33372 torch/fx/experimental/symbolic_shapes.py:4058] create_env
I0708 23:34:22.883000 33372 torch/fx/experimental/symbolic_shapes.py:5330] create_unbacked_symint u0 [-int_oo, int_oo] (_subclasses/fake_impls.py:1233 in local_scalar_dense)
I0708 23:34:22.883000 33372 torch/fx/experimental/symbolic_shapes.py:1448] compute_unbacked_bindings [u0]
I0708 23:34:22.885000 33372 torch/fx/experimental/symbolic_shapes.py:8249] could not evaluate u0 >= 60 due to data dependency, it was assumed to be False with no runtime assertions (_subclasses/fake_impls.py:694 in meta_select)
I0708 23:34:22.885000 33372 torch/fx/experimental/symbolic_shapes.py:8249] For C++ stack trace, run with TORCHDYNAMO_EXTENDED_DEBUG_CPP=1
I0708 23:34:22.886000 33372 torch/fx/experimental/symbolic_shapes.py:8249] could not evaluate u0 < -60 due to data dependency, it was assumed to be False with no runtime assertions (_subclasses/fake_impls.py:694 in meta_select)
I0708 23:34:22.886000 33372 torch/fx/experimental/symbolic_shapes.py:8249] For C++ stack trace, run with TORCHDYNAMO_EXTENDED_DEBUG_CPP=1
I0708 23:34:22.887000 33372 torch/fx/experimental/symbolic_shapes.py:8249] could not evaluate u0 >= 0 due to data dependency, it was assumed to be False with no runtime assertions (_subclasses/fake_impls.py:705 in meta_select)
I0708 23:34:22.887000 33372 torch/fx/experimental/symbolic_shapes.py:8249] For C++ stack trace, run with TORCHDYNAMO_EXTENDED_DEBUG_CPP=1
I0708 23:34:22.888000 33372 torch/fx/experimental/symbolic_shapes.py:8249] could not evaluate u0 < 0 due to data dependency, it was assumed to be False with no runtime assertions (_subclasses/fake_impls.py:707 in meta_select)
I0708 23:34:22.888000 33372 torch/fx/experimental/symbolic_shapes.py:8249] For C++ stack trace, run with TORCHDYNAMO_EXTENDED_DEBUG_CPP=1
I0708 23:34:22.888000 33372 torch/fx/experimental/symbolic_shapes.py:5330] create_unbacked_symint u1 [-int_oo, int_oo] (_subclasses/fake_impls.py:719 in meta_select)
I0708 23:34:22.890000 33372 torch/fx/experimental/symbolic_shapes.py:8249] could not evaluate u1 >= 0 due to data dependency, it was assumed to be True with no runtime assertions (utils/_stats.py:29 in wrapper)
I0708 23:34:22.890000 33372 torch/fx/experimental/symbolic_shapes.py:8249] For C++ stack trace, run with TORCHDYNAMO_EXTENDED_DEBUG_CPP=1
I0708 23:34:22.891000 33372 torch/fx/experimental/symbolic_shapes.py:1448] compute_unbacked_bindings [u1]
I0708 23:34:22.893000 33372 torch/fx/experimental/symbolic_shapes.py:5889] produce_guards
V0708 23:34:22.893000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].storage_offset() 0 None
V0708 23:34:22.893000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].size()[0] 60 None
V0708 23:34:22.893000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].stride()[0] 1 None
V0708 23:34:22.893000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].storage_offset() 0 None

在这种情况下,插入 torch._check() 仅仅是为了防止操作失败。导出调用将失败,提示 “Could not guard on data-dependent expression -u0 > 60”,这意味着编译器不知道这是否是一个有效的索引操作——x 的值对于 y 来说是否越界。在这里,手动特化限制太强,而 torch.cond() 也没有立足之地。相反,告知编译器 u0 的范围就足够了

class Foo(torch.nn.Module):
    def forward(self, x, y):
        a = x.item()
        torch._check(a >= 0)
        torch._check(a < y.shape[0])
        return y[a]

inps = (
    torch.tensor(32),
    torch.randn(60),
)
ep = export(Foo(), inps)
print(ep)
I0708 23:34:22.897000 33372 torch/fx/experimental/symbolic_shapes.py:4058] create_env
I0708 23:34:22.903000 33372 torch/fx/experimental/symbolic_shapes.py:5330] create_unbacked_symint u0 [-int_oo, int_oo] (_subclasses/fake_impls.py:1233 in local_scalar_dense)
I0708 23:34:22.903000 33372 torch/fx/experimental/symbolic_shapes.py:1448] compute_unbacked_bindings [u0]
I0708 23:34:22.904000 33372 torch/fx/experimental/symbolic_shapes.py:8081] runtime_assert u0 >= 0 [guard added] (workspace/intermediate_source/torch_export_tutorial.py:722 in forward), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED="u0 >= 0"
V0708 23:34:22.905000 33372 torch/fx/experimental/symbolic_shapes.py:7503] _update_var_to_range u0 = VR[0, int_oo] (update)
I0708 23:34:22.907000 33372 torch/fx/experimental/symbolic_shapes.py:8081] runtime_assert u0 < 60 [guard added] (workspace/intermediate_source/torch_export_tutorial.py:723 in forward), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED="u0 < 60"
V0708 23:34:22.908000 33372 torch/fx/experimental/symbolic_shapes.py:7503] _update_var_to_range u0 = VR[0, 59] (update)
V0708 23:34:22.910000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
V0708 23:34:22.911000 33372 torch/fx/experimental/symbolic_shapes.py:8396] eval False == False [statically known]
I0708 23:34:22.915000 33372 torch/fx/experimental/symbolic_shapes.py:5889] produce_guards
V0708 23:34:22.915000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].storage_offset() 0 None
V0708 23:34:22.915000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].size()[0] 60 None
V0708 23:34:22.915000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].stride()[0] 1 None
V0708 23:34:22.915000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].storage_offset() 0 None
V0708 23:34:22.918000 33372 torch/fx/experimental/symbolic_shapes.py:8616] runtime_assert u0 <= 59 == True [statically known]
V0708 23:34:22.919000 33372 torch/fx/experimental/symbolic_shapes.py:8616] runtime_assert u0 < 60 == True [statically known]
ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, x: "i64[]", y: "f32[60]"):
            # File: /var/lib/workspace/intermediate_source/torch_export_tutorial.py:721 in forward, code: a = x.item()
            item: "Sym(u0)" = torch.ops.aten.item.default(x);  x = None
            ge_1: "Sym(u0 >= 0)" = item >= 0
            _assert_scalar_default = torch.ops.aten._assert_scalar.default(ge_1, "Runtime assertion failed for expression u0 >= 0 on node 'ge_1'");  ge_1 = _assert_scalar_default = None
            le: "Sym(u0 <= 59)" = item <= 59
            _assert_scalar_default_1 = torch.ops.aten._assert_scalar.default(le, "Runtime assertion failed for expression u0 <= 59 on node 'le'");  le = _assert_scalar_default_1 = None
            lt_1: "Sym(u0 < 60)" = item < 60
            _assert_scalar_default_2 = torch.ops.aten._assert_scalar.default(lt_1, "Runtime assertion failed for expression u0 < 60 on node 'lt_1'");  lt_1 = _assert_scalar_default_2 = None

            # File: /var/lib/workspace/intermediate_source/torch_export_tutorial.py:724 in forward, code: return y[a]
            select: "f32[]" = torch.ops.aten.select.int(y, 0, item);  y = item = None
            return (select,)

Graph signature:
    # inputs
    x: USER_INPUT
    y: USER_INPUT

    # outputs
    select: USER_OUTPUT

Range constraints: {u0: VR[0, 59]}

特化值#

另一类数据依赖错误发生在程序尝试在追踪时提取具体的数据依赖整数/浮点数值。这看起来像是 “Could not extract specialized integer from data-dependent expression”,类似于前一类错误——如果在尝试评估具体的整数/浮点数值时发生这些错误,则在评估具体的布尔值时会产生数据依赖的守卫错误。

这种错误通常发生在对数据依赖表达式进行显式或隐式的 int() 转换时。例如,这个列表推导式有一个 range() 调用,它隐式地对列表大小执行了 int() 转换

class Foo(torch.nn.Module):
    def forward(self, x, y):
        a = x.item()
        b = torch.cat([y for y in range(a)], dim=0)
        return b + int(a)

inps = (
    torch.tensor(32),
    torch.randn(60),
)
try:
    export(Foo(), inps, strict=False)
except Exception:
    tb.print_exc()
I0708 23:34:22.926000 33372 torch/fx/experimental/symbolic_shapes.py:4058] create_env
I0708 23:34:22.933000 33372 torch/fx/experimental/symbolic_shapes.py:5330] create_unbacked_symint u0 [-int_oo, int_oo] (_subclasses/fake_impls.py:1233 in local_scalar_dense)
I0708 23:34:22.934000 33372 torch/fx/experimental/symbolic_shapes.py:1448] compute_unbacked_bindings [u0]
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416] Data dependent variable 'u0' allocated at:
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/usr/local/bin/sphinx-build", line 6, in <module>
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     sys.exit(main())
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/usr/local/lib/python3.10/dist-packages/sphinx/cmd/build.py", line 339, in main
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return make_main(argv)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/usr/local/lib/python3.10/dist-packages/sphinx/cmd/build.py", line 213, in make_main
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return make_mode.run_make_mode(argv[1:])
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/usr/local/lib/python3.10/dist-packages/sphinx/cmd/make_mode.py", line 181, in run_make_mode
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return make.run_generic_build(args[0])
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/usr/local/lib/python3.10/dist-packages/sphinx/cmd/make_mode.py", line 169, in run_generic_build
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return build_main(args + opts)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/usr/local/lib/python3.10/dist-packages/sphinx/cmd/build.py", line 293, in build_main
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     app = Sphinx(args.sourcedir, args.confdir, args.outputdir,
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/usr/local/lib/python3.10/dist-packages/sphinx/application.py", line 272, in __init__
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     self._init_builder()
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/usr/local/lib/python3.10/dist-packages/sphinx/application.py", line 343, in _init_builder
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     self.events.emit('builder-inited')
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/usr/local/lib/python3.10/dist-packages/sphinx/events.py", line 97, in emit
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     results.append(listener.handler(self.app, *args))
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/usr/local/lib/python3.10/dist-packages/sphinx_gallery/gen_gallery.py", line 757, in generate_gallery_rst
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     ) = generate_dir_rst(
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/usr/local/lib/python3.10/dist-packages/sphinx_gallery/gen_rst.py", line 606, in generate_dir_rst
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     results = parallel(
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/usr/local/lib/python3.10/dist-packages/sphinx_gallery/gen_rst.py", line 607, in <genexpr>
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     p_fun(fname, target_dir, src_dir, gallery_conf) for fname in iterator
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/workspace/conf.py", line 85, in wrapper
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     p.start()
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/usr/lib/python3.10/multiprocessing/process.py", line 121, in start
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     self._popen = self._Popen(self)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/usr/lib/python3.10/multiprocessing/context.py", line 224, in _Popen
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return _default_context.get_context().Process._Popen(process_obj)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/usr/lib/python3.10/multiprocessing/context.py", line 281, in _Popen
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return Popen(process_obj)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/usr/lib/python3.10/multiprocessing/popen_fork.py", line 19, in __init__
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     self._launch(process_obj)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/usr/lib/python3.10/multiprocessing/popen_fork.py", line 71, in _launch
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     code = process_obj._bootstrap(parent_sentinel=child_r)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/usr/lib/python3.10/multiprocessing/process.py", line 314, in _bootstrap
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     self.run()
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/usr/lib/python3.10/multiprocessing/process.py", line 108, in run
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     self._target(*self._args, **self._kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/workspace/conf.py", line 73, in call_fn
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     result = func(*args, **kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/usr/local/lib/python3.10/dist-packages/sphinx_gallery/gen_rst.py", line 1374, in generate_file_rst
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     output_blocks, time_elapsed = execute_script(
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/usr/local/lib/python3.10/dist-packages/sphinx_gallery/gen_rst.py", line 1192, in execute_script
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     execute_code_block(
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/usr/local/lib/python3.10/dist-packages/sphinx_gallery/gen_rst.py", line 1048, in execute_code_block
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     is_last_expr, mem_max = _exec_and_get_memory(
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/usr/local/lib/python3.10/dist-packages/sphinx_gallery/gen_rst.py", line 876, in _exec_and_get_memory
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     mem_max, _ = call_memory(
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/usr/local/lib/python3.10/dist-packages/sphinx_gallery/gen_rst.py", line 1725, in _sg_call_memory_noop
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return 0.0, func()
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/usr/local/lib/python3.10/dist-packages/sphinx_gallery/gen_rst.py", line 794, in __call__
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     exec(self.code, self.fake_main.__dict__)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/workspace/intermediate_source/torch_export_tutorial.py", line 756, in <module>
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     export(Foo(), inps, strict=False)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/__init__.py", line 171, in export
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return _export(
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1310, in wrapper
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     ep = fn(*args, **kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/exported_program.py", line 124, in wrapper
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return fn(*args, **kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_utils_internal.py", line 96, in wrapper_function
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return function(*args, **kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2516, in _export
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     ep = _export_for_training(
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1310, in wrapper
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     ep = fn(*args, **kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/exported_program.py", line 124, in wrapper
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return fn(*args, **kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2304, in _export_for_training
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     export_artifact = export_func(
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2233, in _non_strict_export
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     aten_export_artifact = _to_aten_func(
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2010, in _export_to_aten_ir_make_fx
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     gm, graph_signature = transform(_make_fx_helper)(
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2140, in _aot_export_non_strict
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     gm, sig = aot_export(stack, wrapped_mod, args, kwargs=kwargs, **flags)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1916, in _make_fx_helper
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     gm = make_fx(
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 3060, in wrapped
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return make_fx_tracer.trace(f, *args)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 2962, in trace
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return self._trace_inner(f, *args)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 2923, in _trace_inner
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     t = dispatch_trace(
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_compile.py", line 54, in inner
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return disable_fn(*args, **kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 1445, in _fn
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return fn(*args, **kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 1691, in dispatch_trace
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     graph = tracer.trace(root, concrete_args)  # type: ignore[arg-type]
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 2497, in trace
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     res = super().trace(root, concrete_args)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 1445, in _fn
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return fn(*args, **kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 914, in trace
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     (self.create_arg(fn(*args)),),
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 1761, in wrapped
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     out = f(*tensors)  # type:ignore[call-arg]
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "<string>", line 1, in <lambda>
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1798, in wrapped_fn
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return tuple(flat_fn(*args))
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/utils.py", line 192, in flat_fn
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     tree_out = fn(*args, **kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/graph_capture_wrappers.py", line 1536, in functional_call
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     out = mod(*args[params_len:], **kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 888, in module_call_wrapper
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return self.call_module(mod, forward, args, kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 2586, in call_module
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return Tracer.call_module(self, m, forward, args, kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 577, in call_module
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     ret_val = forward(*args, **kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 881, in forward
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return _orig_module_call(mod, *args, **kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return self._call_impl(*args, **kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return forward_call(*args, **kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2124, in forward
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     tree_out = mod(*args, **kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 888, in module_call_wrapper
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return self.call_module(mod, forward, args, kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 2586, in call_module
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return Tracer.call_module(self, m, forward, args, kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 577, in call_module
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     ret_val = forward(*args, **kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 881, in forward
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return _orig_module_call(mod, *args, **kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return self._call_impl(*args, **kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return forward_call(*args, **kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/workspace/intermediate_source/torch_export_tutorial.py", line 747, in forward
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     a = x.item()
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 1823, in __torch_function__
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return func(*args, **kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 1909, in __torch_function__
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return func(*args, **kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_export/non_strict_utils.py", line 1205, in __torch_function__
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return func(*args, **kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_ops.py", line 998, in handler
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return torch._library.utils.handle_dispatch_mode(
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_library/utils.py", line 362, in handle_dispatch_mode
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return curr_mode.__torch_dispatch__(op_overload, overload_types, args, kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_compile.py", line 54, in inner
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return disable_fn(*args, **kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 1446, in _fn
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return fn(*args, **kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/utils/_stats.py", line 29, in wrapper
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return fn(*args, **kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 2007, in __torch_dispatch__
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return proxy_call(self, func, self.pre_dispatch, args, kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 1286, in proxy_call
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     out = func(*args, **kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_ops.py", line 875, in __call__
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return self._op(*args, **kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_compile.py", line 54, in inner
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return disable_fn(*args, **kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 1446, in _fn
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return fn(*args, **kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/utils/_stats.py", line 29, in wrapper
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return fn(*args, **kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py", line 1606, in __torch_dispatch__
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return self.dispatch(func, types, args, kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py", line 2386, in dispatch
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return self._cached_dispatch_impl(func, types, args, kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py", line 1739, in _cached_dispatch_impl
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return self._dispatch_impl(func, types, args, kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py", line 3052, in _dispatch_impl
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     op_impl_out = op_impl(self, func, *args, **kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_subclasses/fake_impls.py", line 205, in dispatch_to_op_implementations_dict
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return op_implementations_dict[func](fake_mode, func, *args, **kwargs)
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_subclasses/fake_impls.py", line 1233, in local_scalar_dense
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     r = fake_mode.shape_env.create_unbacked_symint()
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]   File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/recording.py", line 297, in wrapper
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]     return retlog(fn(*args, **kwargs))
V0708 23:34:22.935000 33372 torch/fx/experimental/symbolic_shapes.py:7416]



def forward(self, arg0_1: "i64[]", arg1_1: "f32[60]"):
    # File: /var/lib/workspace/intermediate_source/torch_export_tutorial.py:747 in forward, code: a = x.item()
    item: "Sym(u0)" = torch.ops.aten.item.default(arg0_1);  arg0_1 = item = None




def forward(self, arg0_1: "i64[]", arg1_1: "f32[60]"):
    # File: /var/lib/workspace/intermediate_source/torch_export_tutorial.py:747 in forward, code: a = x.item()
    item: "Sym(u0)" = torch.ops.aten.item.default(arg0_1);  arg0_1 = item = None

Traceback (most recent call last):
  File "/var/lib/workspace/intermediate_source/torch_export_tutorial.py", line 756, in <module>
    export(Foo(), inps, strict=False)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/__init__.py", line 205, in export
    raise e
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/__init__.py", line 171, in export
    return _export(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1344, in wrapper
    raise e
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1310, in wrapper
    ep = fn(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/exported_program.py", line 124, in wrapper
    return fn(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_utils_internal.py", line 96, in wrapper_function
    return function(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2516, in _export
    ep = _export_for_training(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1344, in wrapper
    raise e
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1310, in wrapper
    ep = fn(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/exported_program.py", line 124, in wrapper
    return fn(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2304, in _export_for_training
    export_artifact = export_func(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2233, in _non_strict_export
    aten_export_artifact = _to_aten_func(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2010, in _export_to_aten_ir_make_fx
    gm, graph_signature = transform(_make_fx_helper)(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2140, in _aot_export_non_strict
    gm, sig = aot_export(stack, wrapped_mod, args, kwargs=kwargs, **flags)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1916, in _make_fx_helper
    gm = make_fx(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 3060, in wrapped
    return make_fx_tracer.trace(f, *args)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 2962, in trace
    return self._trace_inner(f, *args)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 2923, in _trace_inner
    t = dispatch_trace(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_compile.py", line 54, in inner
    return disable_fn(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 1445, in _fn
    return fn(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 1691, in dispatch_trace
    graph = tracer.trace(root, concrete_args)  # type: ignore[arg-type]
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 2497, in trace
    res = super().trace(root, concrete_args)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 1445, in _fn
    return fn(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 914, in trace
    (self.create_arg(fn(*args)),),
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 1761, in wrapped
    out = f(*tensors)  # type:ignore[call-arg]
  File "<string>", line 1, in <lambda>
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1798, in wrapped_fn
    return tuple(flat_fn(*args))
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/utils.py", line 192, in flat_fn
    tree_out = fn(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/graph_capture_wrappers.py", line 1536, in functional_call
    out = mod(*args[params_len:], **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 888, in module_call_wrapper
    return self.call_module(mod, forward, args, kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 2586, in call_module
    return Tracer.call_module(self, m, forward, args, kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 577, in call_module
    ret_val = forward(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 881, in forward
    return _orig_module_call(mod, *args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
    return forward_call(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 2124, in forward
    tree_out = mod(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 888, in module_call_wrapper
    return self.call_module(mod, forward, args, kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 2586, in call_module
    return Tracer.call_module(self, m, forward, args, kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 577, in call_module
    ret_val = forward(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 881, in forward
    return _orig_module_call(mod, *args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
    return forward_call(*args, **kwargs)
  File "/var/lib/workspace/intermediate_source/torch_export_tutorial.py", line 748, in forward
    b = torch.cat([y for y in range(a)], dim=0)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/__init__.py", line 467, in __index__
    return self.node.int_()
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/sym_node.py", line 502, in int_
    return self.guard_int("", 0)  # NB: uses Python backtrace
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/sym_node.py", line 556, in guard_int
    r = self.evaluate()
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/sym_node.py", line 550, in evaluate
    return self.shape_env.evaluate_sym_node(self, size_oblivious)
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/symbolic_shapes.py", line 8107, in evaluate_sym_node
    return self.evaluate_expr(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/symbolic_shapes.py", line 8203, in evaluate_expr
    return self._inner_evaluate_expr(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/recording.py", line 297, in wrapper
    return retlog(fn(*args, **kwargs))
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/symbolic_shapes.py", line 8226, in _inner_evaluate_expr
    return self._evaluate_expr(
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/experimental/symbolic_shapes.py", line 8459, in _evaluate_expr
    raise self._make_data_dependent_error(
torch.fx.experimental.symbolic_shapes.GuardOnDataDependentSymNode: Could not extract specialized integer from data-dependent expression u0 (unhinted: u0).  (Size-like symbols: none)


Caused by: (workspace/intermediate_source/torch_export_tutorial.py:748 in forward)
For more information, run with TORCH_LOGS="dynamic"
For extended logs when we create symbols, also add TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="u0"
If you suspect the guard was triggered from C++, add TORCHDYNAMO_EXTENDED_DEBUG_CPP=1
For more debugging help, see https://docs.google.com/document/d/1HSuTTVvYH1pTew89Rtpeu84Ht3nQEFTYhAX3Ypa_xJs/edit?usp=sharing

For C++ stack trace, run with TORCHDYNAMO_EXTENDED_DEBUG_CPP=1

The error above occurred when calling torch.export.export. If you would like to view some more information about this error, and get a list of all other errors that may occur in your export call, you can replace your `export()` call with `draft_export()`.

对于这些错误,您拥有的一些基本选项是:

  1. 避免不必要的 int() 转换调用,在本例中即 return 语句中的 int(a)

  2. 使用 torch._check() 调用;不幸的是,在这种情况下您所能做的可能只是进行特化(使用 torch._check(a == 60))。

  3. 在更高级别上重写有问题的代码。例如,列表推导式在语义上是一个 repeat() 操作,它不涉及 int() 转换。以下重写避免了数据依赖错误

class Foo(torch.nn.Module):
    def forward(self, x, y):
        a = x.item()
        b = y.unsqueeze(0).repeat(a, 1)
        return b + a

inps = (
    torch.tensor(32),
    torch.randn(60),
)
ep = export(Foo(), inps, strict=False)
print(ep)
I0708 23:34:22.956000 33372 torch/fx/experimental/symbolic_shapes.py:4058] create_env
I0708 23:34:22.963000 33372 torch/fx/experimental/symbolic_shapes.py:5330] create_unbacked_symint u0 [-int_oo, int_oo] (_subclasses/fake_impls.py:1233 in local_scalar_dense)
I0708 23:34:22.963000 33372 torch/fx/experimental/symbolic_shapes.py:1448] compute_unbacked_bindings [u0]
I0708 23:34:22.967000 33372 torch/fx/experimental/symbolic_shapes.py:8081] runtime_assert u0 >= 0 [guard added] (_meta_registrations.py:4566 in meta_repeat), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED="u0 >= 0"
V0708 23:34:22.968000 33372 torch/fx/experimental/symbolic_shapes.py:7503] _update_var_to_range u0 = VR[0, int_oo] (update)
I0708 23:34:22.973000 33372 torch/fx/experimental/symbolic_shapes.py:8249] could not evaluate Eq(u0, 0) due to data dependency, it was assumed to be False with no runtime assertions (utils/_stats.py:29 in wrapper)
I0708 23:34:22.973000 33372 torch/fx/experimental/symbolic_shapes.py:8249] For C++ stack trace, run with TORCHDYNAMO_EXTENDED_DEBUG_CPP=1
I0708 23:34:22.981000 33372 torch/fx/experimental/symbolic_shapes.py:8249] could not evaluate 60*u0 < 2 due to data dependency, it was assumed to be False with no runtime assertions (_prims_common/__init__.py:322 in is_contiguous)
I0708 23:34:22.981000 33372 torch/fx/experimental/symbolic_shapes.py:8249] For C++ stack trace, run with TORCHDYNAMO_EXTENDED_DEBUG_CPP=1
I0708 23:34:22.982000 33372 torch/fx/experimental/symbolic_shapes.py:8249] could not evaluate Eq(u0, 1) due to data dependency, it was assumed to be False with no runtime assertions (_prims_common/__init__.py:285 in check_contiguous_sizes_strides)
I0708 23:34:22.982000 33372 torch/fx/experimental/symbolic_shapes.py:8249] For C++ stack trace, run with TORCHDYNAMO_EXTENDED_DEBUG_CPP=1
V0708 23:34:22.985000 33372 torch/fx/experimental/symbolic_shapes.py:8616] runtime_assert True == True [statically known]
I0708 23:34:22.990000 33372 torch/fx/experimental/symbolic_shapes.py:5889] produce_guards
V0708 23:34:22.990000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].storage_offset() 0 None
V0708 23:34:22.990000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].size()[0] 60 None
V0708 23:34:22.990000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].stride()[0] 1 None
V0708 23:34:22.991000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['y'].storage_offset() 0 None
ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, x: "i64[]", y: "f32[60]"):
            # File: /var/lib/workspace/intermediate_source/torch_export_tutorial.py:769 in forward, code: a = x.item()
            item: "Sym(u0)" = torch.ops.aten.item.default(x);  x = None
            ge: "Sym(u0 >= 0)" = item >= 0
            _assert_scalar_default = torch.ops.aten._assert_scalar.default(ge, "Runtime assertion failed for expression u0 >= 0 on node 'ge'");  ge = _assert_scalar_default = None

            # File: /var/lib/workspace/intermediate_source/torch_export_tutorial.py:770 in forward, code: b = y.unsqueeze(0).repeat(a, 1)
            unsqueeze: "f32[1, 60]" = torch.ops.aten.unsqueeze.default(y, 0);  y = None
            repeat: "f32[u0, 60]" = torch.ops.aten.repeat.default(unsqueeze, [item, 1]);  unsqueeze = None

            # File: /var/lib/workspace/intermediate_source/torch_export_tutorial.py:771 in forward, code: return b + a
            add: "f32[u0, 60]" = torch.ops.aten.add.Tensor(repeat, item);  repeat = item = None
            return (add,)

Graph signature:
    # inputs
    x: USER_INPUT
    y: USER_INPUT

    # outputs
    add: USER_OUTPUT

Range constraints: {u0: VR[0, int_oo]}

数据依赖的错误可能复杂得多,您的工具箱中还有更多应对方案:入门级包括 torch._check_is_size()guard_size_oblivious() 或真实张量追踪。如需更深入的指南,请参考 导出编程模型,或 处理 GuardOnDataDependentSymNode 错误

自定义操作 (Custom Ops)#

torch.export 可以导出带有自定义操作符的 PyTorch 程序。关于如何用 C++ 或 Python 编写自定义操作符,请参考 此页面

以下是在 Python 中注册一个自定义操作符以便用于 torch.export 的示例。需要注意的重要一点是,自定义操作符必须具有 FakeTensor 内核

@torch.library.custom_op("my_custom_library::custom_op", mutates_args={})
def custom_op(x: torch.Tensor) -> torch.Tensor:
    print("custom_op called!")
    return torch.relu(x)

@custom_op.register_fake
def custom_op_meta(x):
    # Returns an empty tensor with the same shape as the expected output
    return torch.empty_like(x)

以下是一个导出带有自定义操作符程序的示例。

class CustomOpExample(torch.nn.Module):
    def forward(self, x):
        x = torch.sin(x)
        x = torch.ops.my_custom_library.custom_op(x)
        x = torch.cos(x)
        return x

exported_custom_op_example = export(CustomOpExample(), (torch.randn(3, 3),))
print(exported_custom_op_example)
print(exported_custom_op_example.module()(torch.randn(3, 3)))
I0708 23:34:23.007000 33372 torch/fx/experimental/symbolic_shapes.py:4058] create_env
I0708 23:34:23.017000 33372 torch/fx/experimental/symbolic_shapes.py:5889] produce_guards
V0708 23:34:23.017000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].size()[0] 3 None
V0708 23:34:23.017000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].size()[1] 3 None
V0708 23:34:23.017000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].stride()[0] 3 None
V0708 23:34:23.017000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].stride()[1] 1 None
V0708 23:34:23.017000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].storage_offset() 0 None
ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, x: "f32[3, 3]"):
            # File: /var/lib/workspace/intermediate_source/torch_export_tutorial.py:812 in forward, code: x = torch.sin(x)
            sin: "f32[3, 3]" = torch.ops.aten.sin.default(x);  x = None

            # File: /var/lib/workspace/intermediate_source/torch_export_tutorial.py:813 in forward, code: x = torch.ops.my_custom_library.custom_op(x)
            custom_op: "f32[3, 3]" = torch.ops.my_custom_library.custom_op.default(sin);  sin = None

            # File: /var/lib/workspace/intermediate_source/torch_export_tutorial.py:814 in forward, code: x = torch.cos(x)
            cos: "f32[3, 3]" = torch.ops.aten.cos.default(custom_op);  custom_op = None
            return (cos,)

Graph signature:
    # inputs
    x: USER_INPUT

    # outputs
    cos: USER_OUTPUT

Range constraints: {}

custom_op called!
tensor([[1.0000, 0.6068, 1.0000],
        [1.0000, 0.6749, 1.0000],
        [0.7340, 1.0000, 1.0000]])

注意,在 ExportedProgram 中,自定义操作符被包含在图中。

IR/分解 (Decompositions)#

torch.export 生成的图返回一个仅包含 ATen 操作符 的图,ATen 操作符是 PyTorch 中的基本计算单位。由于有超过 3000 个 ATen 操作符,导出提供了一种根据特定特征缩小图中使用的操作符集的方法,从而创建不同的 IR。

默认情况下,导出产生最通用的 IR,其中包含所有 ATen 操作符,包括函数式和非函数式操作符。函数式操作符是指不包含任何输入变动或别名的操作符。您可以在这里找到所有 ATen 操作符的列表,并且可以通过检查 op._schema.is_mutable 来检查操作符是否为函数式的,例如

print(torch.ops.aten.add.Tensor._schema.is_mutable)
print(torch.ops.aten.add_.Tensor._schema.is_mutable)
False
True

这种通用 IR 可用于在 eager PyTorch Autograd 中进行训练。

class DecompExample(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,)

ep_for_training = torch.export.export(DecompExample(), (torch.randn(1, 1, 3, 3),))
print(ep_for_training.graph)
I0708 23:34:23.028000 33372 torch/fx/experimental/symbolic_shapes.py:4058] create_env
I0708 23:34:23.063000 33372 torch/fx/experimental/symbolic_shapes.py:5889] produce_guards
V0708 23:34:23.063000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].size()[0] 1 None
V0708 23:34:23.063000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].size()[1] 1 None
V0708 23:34:23.063000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].size()[2] 3 None
V0708 23:34:23.063000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].size()[3] 3 None
V0708 23:34:23.063000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].stride()[0] 9 None
V0708 23:34:23.064000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].stride()[1] 9 None
V0708 23:34:23.064000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].stride()[2] 3 None
V0708 23:34:23.064000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].stride()[3] 1 None
V0708 23:34:23.064000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].storage_offset() 0 None
graph():
    %p_conv_weight : [num_users=1] = placeholder[target=p_conv_weight]
    %p_conv_bias : [num_users=1] = placeholder[target=p_conv_bias]
    %p_bn_weight : [num_users=1] = placeholder[target=p_bn_weight]
    %p_bn_bias : [num_users=1] = placeholder[target=p_bn_bias]
    %b_bn_running_mean : [num_users=1] = placeholder[target=b_bn_running_mean]
    %b_bn_running_var : [num_users=1] = placeholder[target=b_bn_running_var]
    %b_bn_num_batches_tracked : [num_users=1] = placeholder[target=b_bn_num_batches_tracked]
    %x : [num_users=1] = placeholder[target=x]
    %conv2d : [num_users=1] = call_function[target=torch.ops.aten.conv2d.default](args = (%x, %p_conv_weight, %p_conv_bias), kwargs = {})
    %add_ : [num_users=0] = call_function[target=torch.ops.aten.add_.Tensor](args = (%b_bn_num_batches_tracked, 1), kwargs = {})
    %batch_norm : [num_users=1] = call_function[target=torch.ops.aten.batch_norm.default](args = (%conv2d, %p_bn_weight, %p_bn_bias, %b_bn_running_mean, %b_bn_running_var, True, 0.1, 1e-05, False), kwargs = {})
    return (batch_norm,)

然后我们可以通过 API run_decompositions 将此导出的程序降级为仅包含函数式 ATen 操作符的操作符集,该 API 会将 ATen 操作符分解为分解表中指定的那些,并使图函数化。通过指定空集,我们仅执行函数化,不进行任何额外的分解。这会产生一个包含约 2000 个操作符(而不是上面的 3000 个)的 IR,是推理用例的理想选择。

/usr/lib/python3.10/copyreg.py:101: FutureWarning: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
  return cls.__new__(cls, *args)
graph():
    %p_conv_weight : [num_users=1] = placeholder[target=p_conv_weight]
    %p_conv_bias : [num_users=1] = placeholder[target=p_conv_bias]
    %p_bn_weight : [num_users=1] = placeholder[target=p_bn_weight]
    %p_bn_bias : [num_users=1] = placeholder[target=p_bn_bias]
    %b_bn_running_mean : [num_users=1] = placeholder[target=b_bn_running_mean]
    %b_bn_running_var : [num_users=1] = placeholder[target=b_bn_running_var]
    %b_bn_num_batches_tracked : [num_users=1] = placeholder[target=b_bn_num_batches_tracked]
    %x : [num_users=1] = placeholder[target=x]
    %conv2d : [num_users=1] = call_function[target=torch.ops.aten.conv2d.default](args = (%x, %p_conv_weight, %p_conv_bias), kwargs = {})
    %add : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%b_bn_num_batches_tracked, 1), kwargs = {})
    %_native_batch_norm_legit_functional : [num_users=3] = call_function[target=torch.ops.aten._native_batch_norm_legit_functional.default](args = (%conv2d, %p_bn_weight, %p_bn_bias, %b_bn_running_mean, %b_bn_running_var, True, 0.1, 1e-05), kwargs = {})
    %getitem : [num_users=1] = call_function[target=operator.getitem](args = (%_native_batch_norm_legit_functional, 0), kwargs = {})
    %getitem_3 : [num_users=1] = call_function[target=operator.getitem](args = (%_native_batch_norm_legit_functional, 3), kwargs = {})
    %getitem_4 : [num_users=1] = call_function[target=operator.getitem](args = (%_native_batch_norm_legit_functional, 4), kwargs = {})
    return (getitem_3, getitem_4, add, getitem)

正如我们所见,之前可变的操作符 torch.ops.aten.add_.default 现在已被 torch.ops.aten.add.default(一个函数式操作符)替换。

我们还可以进一步将此导出的程序降级为仅包含 核心 ATen 操作符集 的操作符集,该集合仅包含约 180 个操作符。对于不想重新实现所有 ATen 操作符的后端,此 IR 是最佳选择。

from torch.export import default_decompositions

core_aten_decomp_table = default_decompositions()
core_aten_ep = ep_for_training.run_decompositions(decomp_table=core_aten_decomp_table)
print(core_aten_ep.graph)
/usr/lib/python3.10/copyreg.py:101: FutureWarning: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
  return cls.__new__(cls, *args)
graph():
    %p_conv_weight : [num_users=1] = placeholder[target=p_conv_weight]
    %p_conv_bias : [num_users=1] = placeholder[target=p_conv_bias]
    %p_bn_weight : [num_users=1] = placeholder[target=p_bn_weight]
    %p_bn_bias : [num_users=1] = placeholder[target=p_bn_bias]
    %b_bn_running_mean : [num_users=1] = placeholder[target=b_bn_running_mean]
    %b_bn_running_var : [num_users=1] = placeholder[target=b_bn_running_var]
    %b_bn_num_batches_tracked : [num_users=1] = placeholder[target=b_bn_num_batches_tracked]
    %x : [num_users=1] = placeholder[target=x]
    %convolution : [num_users=1] = call_function[target=torch.ops.aten.convolution.default](args = (%x, %p_conv_weight, %p_conv_bias, [1, 1], [0, 0], [1, 1], False, [0, 0], 1), kwargs = {})
    %add : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%b_bn_num_batches_tracked, 1), kwargs = {})
    %_native_batch_norm_legit_functional : [num_users=3] = call_function[target=torch.ops.aten._native_batch_norm_legit_functional.default](args = (%convolution, %p_bn_weight, %p_bn_bias, %b_bn_running_mean, %b_bn_running_var, True, 0.1, 1e-05), kwargs = {})
    %getitem : [num_users=1] = call_function[target=operator.getitem](args = (%_native_batch_norm_legit_functional, 0), kwargs = {})
    %getitem_3 : [num_users=1] = call_function[target=operator.getitem](args = (%_native_batch_norm_legit_functional, 3), kwargs = {})
    %getitem_4 : [num_users=1] = call_function[target=operator.getitem](args = (%_native_batch_norm_legit_functional, 4), kwargs = {})
    return (getitem_3, getitem_4, add, getitem)

我们现在看到 torch.ops.aten.conv2d.default 已被分解为 torch.ops.aten.convolution.default。这是因为 convolution 是一个更 “核心” 的操作符,因为像 conv1dconv2d 这样的操作可以使用相同的 op 来实现。

我们也可以指定自己的分解行为

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

my_decomp_table[torch.ops.aten.conv2d.default] = my_awesome_custom_conv2d_function
my_ep = ep_for_training.run_decompositions(my_decomp_table)
print(my_ep.graph)
/usr/lib/python3.10/copyreg.py:101: FutureWarning: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
  return cls.__new__(cls, *args)
graph():
    %p_conv_weight : [num_users=1] = placeholder[target=p_conv_weight]
    %p_conv_bias : [num_users=1] = placeholder[target=p_conv_bias]
    %p_bn_weight : [num_users=1] = placeholder[target=p_bn_weight]
    %p_bn_bias : [num_users=1] = placeholder[target=p_bn_bias]
    %b_bn_running_mean : [num_users=1] = placeholder[target=b_bn_running_mean]
    %b_bn_running_var : [num_users=1] = placeholder[target=b_bn_running_var]
    %b_bn_num_batches_tracked : [num_users=1] = placeholder[target=b_bn_num_batches_tracked]
    %x : [num_users=1] = placeholder[target=x]
    %convolution : [num_users=1] = call_function[target=torch.ops.aten.convolution.default](args = (%x, %p_conv_weight, %p_conv_bias, [1, 1], [0, 0], [1, 1], False, [0, 0], 1), kwargs = {})
    %mul : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%convolution, 2), kwargs = {})
    %add : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%b_bn_num_batches_tracked, 1), kwargs = {})
    %_native_batch_norm_legit_functional : [num_users=3] = call_function[target=torch.ops.aten._native_batch_norm_legit_functional.default](args = (%mul, %p_bn_weight, %p_bn_bias, %b_bn_running_mean, %b_bn_running_var, True, 0.1, 1e-05), kwargs = {})
    %getitem : [num_users=1] = call_function[target=operator.getitem](args = (%_native_batch_norm_legit_functional, 0), kwargs = {})
    %getitem_3 : [num_users=1] = call_function[target=operator.getitem](args = (%_native_batch_norm_legit_functional, 3), kwargs = {})
    %getitem_4 : [num_users=1] = call_function[target=operator.getitem](args = (%_native_batch_norm_legit_functional, 4), kwargs = {})
    return (getitem_3, getitem_4, add, getitem)

注意,torch.ops.aten.conv2d.default 现在不再仅分解为 torch.ops.aten.convolution.default,而是分解为 torch.ops.aten.convolution.defaulttorch.ops.aten.mul.Tensor,这符合我们的自定义分解规则。

ExportDB#

torch.export 将始终从 PyTorch 程序中导出一个单一的计算图。由于这一要求,将会出现一些与 torch.export 不兼容的 Python 或 PyTorch 特性,这将需要用户重写其部分模型代码。我们在教程的前面已经看到了这方面的示例——例如,使用 cond 重写 if 语句。

ExportDB 是记载 torch.export 支持和不支持的 Python/PyTorch 特性的标准参考文档。它本质上是一个程序样本列表,每个样本代表一种特定 Python/PyTorch 特性的用法及其与 torch.export 的交互。示例还按类别加了标签,以便更容易地搜索。

例如,让我们使用 ExportDB 来更好地了解 cond 操作符中的谓词是如何工作的。我们可以查看一个名为 cond_predicate 的示例,它带有一个 torch.cond 标签。示例代码如下:

def cond_predicate(x):
    """
    The conditional statement (aka predicate) passed to ``cond()`` must be one of the following:
    - ``torch.Tensor`` with a single element
    - boolean expression
    NOTE: If the `pred` is test on a dim with batch size < 2, it will be specialized.
    """
    pred = x.dim() > 2 and x.shape[2] > 10
    return cond(pred, lambda x: x.cos(), lambda y: y.sin(), [x])

更广泛地说,当发生以下情况之一时,ExportDB 可以用作参考:

  1. 在尝试使用 torch.export 之前,您提前知道您的模型使用了一些棘手的 Python/PyTorch 特性,并且您想知道 torch.export 是否涵盖该特性。

  2. 在尝试使用 torch.export 时,发生了失败,且不清楚如何解决。

ExportDB 并不是详尽无遗的,但旨在涵盖典型 PyTorch 代码中发现的所有用例。如果有某个重要的 Python/PyTorch 特性应当添加到 ExportDB 或由 torch.export 支持,欢迎与我们联系。

运行已导出的程序#

由于 torch.export 仅是一种图捕捉机制,直接调用 torch.export 生成的产物相当于运行 eager 模块。为了优化导出程序的执行,我们可以通过 torch.compile 将此导出的产物传递给诸如 Inductor、AOTInductorTensorRT 之类的后端。

class M(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(3, 3)

    def forward(self, x):
        x = self.linear(x)
        return x

inp = torch.randn(2, 3, device="cuda")
m = M().to(device="cuda")
ep = torch.export.export(m, (inp,))

# Run it eagerly
res = ep.module()(inp)
print(res)

# Run it with torch.compile
res = torch.compile(ep.module(), backend="inductor")(inp)
print(res)
I0708 23:34:23.886000 33372 torch/fx/experimental/symbolic_shapes.py:4058] create_env
I0708 23:34:23.900000 33372 torch/fx/experimental/symbolic_shapes.py:5889] produce_guards
V0708 23:34:23.900000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].size()[0] 2 None
V0708 23:34:23.900000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].size()[1] 3 None
V0708 23:34:23.901000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].stride()[0] 3 None
V0708 23:34:23.901000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].stride()[1] 1 None
V0708 23:34:23.901000 33372 torch/fx/experimental/symbolic_shapes.py:6122] track_symint L['x'].storage_offset() 0 None
tensor([[ 0.3467, -0.3446,  0.6130],
        [-1.7386, -0.2031, -1.9849]], device='cuda:0',
       grad_fn=<AddmmBackward0>)
I0708 23:34:24.584000 33372 torch/fx/experimental/symbolic_shapes.py:4058] [2/0] create_env
/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_inductor/compile_fx.py:321: UserWarning: TensorFloat32 tensor cores for float32 matrix multiplication available but not enabled. Consider setting `torch.set_float32_matmul_precision('high')` for better performance.
  warnings.warn(
I0708 23:34:25.889000 33372 torch/fx/experimental/symbolic_shapes.py:5889] [2/0] produce_guards
I0708 23:34:25.901000 33372 torch/fx/experimental/symbolic_shapes.py:5889] [2/0] produce_guards
V0708 23:34:25.901000 33372 torch/fx/experimental/symbolic_shapes.py:6122] [2/0] track_symint L['x'].size()[0] 2 None
V0708 23:34:25.901000 33372 torch/fx/experimental/symbolic_shapes.py:6122] [2/0] track_symint L['x'].size()[1] 3 None
V0708 23:34:25.901000 33372 torch/fx/experimental/symbolic_shapes.py:6122] [2/0] track_symint L['x'].stride()[0] 3 None
V0708 23:34:25.901000 33372 torch/fx/experimental/symbolic_shapes.py:6122] [2/0] track_symint L['x'].stride()[1] 1 None
V0708 23:34:25.902000 33372 torch/fx/experimental/symbolic_shapes.py:6122] [2/0] track_symint L['x'].storage_offset() 0 None
V0708 23:34:25.902000 33372 torch/fx/experimental/symbolic_shapes.py:6122] [2/0] track_symint L['self']._modules['linear']._parameters['weight'].size()[0] 3 None
V0708 23:34:25.902000 33372 torch/fx/experimental/symbolic_shapes.py:6122] [2/0] track_symint L['self']._modules['linear']._parameters['weight'].size()[1] 3 None
V0708 23:34:25.902000 33372 torch/fx/experimental/symbolic_shapes.py:6122] [2/0] track_symint L['self']._modules['linear']._parameters['weight'].stride()[0] 3 None
V0708 23:34:25.902000 33372 torch/fx/experimental/symbolic_shapes.py:6122] [2/0] track_symint L['self']._modules['linear']._parameters['weight'].stride()[1] 1 None
V0708 23:34:25.902000 33372 torch/fx/experimental/symbolic_shapes.py:6122] [2/0] track_symint L['self']._modules['linear']._parameters['weight'].storage_offset() 0 None
V0708 23:34:25.902000 33372 torch/fx/experimental/symbolic_shapes.py:6122] [2/0] track_symint L['self']._modules['linear']._parameters['bias'].size()[0] 3 None
V0708 23:34:25.902000 33372 torch/fx/experimental/symbolic_shapes.py:6122] [2/0] track_symint L['self']._modules['linear']._parameters['bias'].stride()[0] 1 None
V0708 23:34:25.902000 33372 torch/fx/experimental/symbolic_shapes.py:6122] [2/0] track_symint L['self']._modules['linear']._parameters['bias'].storage_offset() 0 None
V0708 23:34:25.903000 33372 torch/fx/experimental/symbolic_shapes.py:6386] [2/0] Skipping guard L['x'].size()[0] == 2
V0708 23:34:25.903000 33372 torch/fx/experimental/symbolic_shapes.py:6386] [2/0] Skipping guard L['x'].size()[1] == 3
V0708 23:34:25.903000 33372 torch/fx/experimental/symbolic_shapes.py:6386] [2/0] Skipping guard L['x'].stride()[0] == 3
V0708 23:34:25.903000 33372 torch/fx/experimental/symbolic_shapes.py:6386] [2/0] Skipping guard L['x'].stride()[1] == 1
V0708 23:34:25.903000 33372 torch/fx/experimental/symbolic_shapes.py:6386] [2/0] Skipping guard L['x'].storage_offset() == 0
V0708 23:34:25.903000 33372 torch/fx/experimental/symbolic_shapes.py:6386] [2/0] Skipping guard L['self']._modules['linear']._parameters['weight'].size()[0] == 3
V0708 23:34:25.903000 33372 torch/fx/experimental/symbolic_shapes.py:6386] [2/0] Skipping guard L['self']._modules['linear']._parameters['weight'].size()[1] == 3
V0708 23:34:25.904000 33372 torch/fx/experimental/symbolic_shapes.py:6386] [2/0] Skipping guard L['self']._modules['linear']._parameters['weight'].stride()[0] == 3
V0708 23:34:25.904000 33372 torch/fx/experimental/symbolic_shapes.py:6386] [2/0] Skipping guard L['self']._modules['linear']._parameters['weight'].stride()[1] == 1
V0708 23:34:25.904000 33372 torch/fx/experimental/symbolic_shapes.py:6386] [2/0] Skipping guard L['self']._modules['linear']._parameters['weight'].storage_offset() == 0
V0708 23:34:25.904000 33372 torch/fx/experimental/symbolic_shapes.py:6386] [2/0] Skipping guard L['self']._modules['linear']._parameters['bias'].size()[0] == 3
V0708 23:34:25.904000 33372 torch/fx/experimental/symbolic_shapes.py:6386] [2/0] Skipping guard L['self']._modules['linear']._parameters['bias'].stride()[0] == 1
V0708 23:34:25.904000 33372 torch/fx/experimental/symbolic_shapes.py:6386] [2/0] Skipping guard L['self']._modules['linear']._parameters['bias'].storage_offset() == 0
tensor([[ 0.3467, -0.3446,  0.6130],
        [-1.7386, -0.2031, -1.9849]], device='cuda:0',
       grad_fn=<CompiledFunctionBackward>)
import torch._inductor

# Note: these APIs are subject to change
# Compile the exported program to a PT2 archive using ``AOTInductor``
with torch.no_grad():
    pt2_path = torch._inductor.aoti_compile_and_package(ep)

# Load and run the .so file in Python.
# To load and run it in a C++ environment, see:
# https://docs.pytorch.ac.cn/docs/stable/torch.compiler_aot_inductor.html
aoti_compiled = torch._inductor.aoti_load_package(pt2_path)
res = aoti_compiled(inp)

结论#

我们介绍了 torch.export,这是 PyTorch 2.X 从 PyTorch 程序中导出单一计算图的新方式。特别是,我们演示了几种为了导出图而需要进行的编码修改和注意事项(控制流操作、约束等)。

脚本总运行时间:(0 分钟 4.651 秒)