PyTorch FX symbolic_trace 学习笔记

PyTorch

中文

理解 torch.fx symbolic tracing、GraphModule 和 FX graph nodes。

Author

Wangkun Xu

Published

August 3, 2026

Modified

August 3, 2026

1. 核心概念

torch.fx.symbolic_trace() 可以理解为一次“不使用真实数值的模拟执行”：

FX 为 forward() 的输入创建 Proxy；
使用这些 Proxy 执行原始 forward()；
Proxy 拦截 module、function 和 tensor method 调用；
tracer 将这些操作记录为 Node，形成 Graph；
FX 根据 Graph 生成新的 forward()；
最终返回可执行的 fx.GraphModule。

Tracing 阶段和正常执行阶段的区别如下：

阶段	输入	操作结果
Symbolic tracing	`Proxy`	创建 graph nodes，不进行真实 tensor 计算
`GraphModule` runtime	`Tensor`	执行真实数值计算

Tracing 时，Proxy 是 Node 的 Python 包装器；tracing 完成后，Graph 保存的是 Node。

2. Residual network 示例

import torch
from torch import fx, nn

class ResidualMLP(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.linear = nn.Linear(4, 4)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.linear(x) + x

model = ResidualMLP().eval()
x = torch.tensor([[0.2, -0.4, 0.7, 0.1]])
traced = fx.symbolic_trace(model)

with torch.no_grad():
    torch.testing.assert_close(traced(x), model(x))

traced 是一个 fx.GraphModule。它继承自 nn.Module，包含 FX graph、graph 使用的子模块和参数，以及根据 graph 生成的 forward()，所以可以直接执行：

output = traced(x)

3. `GraphModule` 的常用属性

属性或方法	用途
`traced.graph`	FX 中间表示，包含拓扑有序的 nodes
`traced.code`	FX 根据 graph 生成的 Python `forward()`
`traced.forward`	可执行的生成函数
`traced.<submodule>`	graph 引用的子模块，例如 `traced.linear`
`traced.named_modules()`	根模块和所有子模块
`traced.named_parameters()`	所有 `nn.Parameter`
`traced.state_dict()`	parameters 和 persistent buffers
`traced.recompile()`	修改 graph 后重新生成 `forward()`

print(traced.graph)
print(traced.code)
print("Modules:", [name for name, _ in traced.named_modules()])
print("Parameters:", [name for name, _ in traced.named_parameters()])
print("State dict:", list(traced.state_dict()))

输出大致为：

graph():
    %x : torch.Tensor [num_users=2] = placeholder[target=x]
    %linear : [num_users=1] = call_module[target=linear](args=(%x,), kwargs={})
    %add : [num_users=1] = call_function[target=operator.add](args=(%linear, %x), kwargs={})
    return add

def forward(self, x: torch.Tensor) -> torch.Tensor:
    linear = self.linear(x)
    add = linear + x
    linear = x = None
    return add

Modules: ['', 'linear']
Parameters: ['linear.weight', 'linear.bias']
State dict: ['linear.weight', 'linear.bias']

其中：

Modules 中的 '' 表示根模块 traced，'linear' 表示 traced.linear；
Parameters 包含可训练的 linear.weight 和 linear.bias；
state_dict 包含 parameters 和 persistent buffers。本例没有 buffers，因此键与 Parameters 相同；
生成代码中的 linear = x = None 用于尽早释放不再使用的局部引用，不改变计算结果。

4. 如何阅读一行 FX graph

%x : torch.Tensor [num_users=2] = placeholder[target=x]

%x：对名为 x 的 node 输出的引用；
torch.Tensor：来自 forward() 的 Python 类型标注，不是 dtype；
num_users=2：有两个不同的 consumer nodes 使用该输出；
placeholder：该 node 表示模型输入；
target=x：对应 forward(self, x) 中的参数 x。

本例中 x 的两个 consumers 是 linear 和 add：

         linear
        /      \
x -----        add
        \______/

num_users 统计不同 consumer nodes，而不是在参数中出现的总次数。例如 x + x 中，x 在 add.args 出现两次，但 consumer node 只有一个 add。

5. `Node` 的关键属性

属性	含义
`name`	graph 内该次操作的唯一名称
`op`	FX 的操作类别
`target`	输入名、attribute path、module path、callable 或 method 名
`args`	positional arguments，其中的 `Node` 引用形成 graph edges
`kwargs`	keyword arguments
`users`	使用当前 node 输出的 consumer nodes
`all_input_nodes`	从嵌套 `args/kwargs` 中提取的所有输入 nodes
`meta`	shape、dtype、module stack 等附加信息
`type`	Python type annotation，不是 tensor dtype

可以这样记忆：

target          = 调用什么
args/kwargs     = 输入和调用参数是什么
all_input_nodes = 输入由哪些 nodes 产生
users           = 输出被哪些 nodes 使用
name            = 这一次 graph 操作叫什么

查看全部 nodes：

for node in traced.graph.nodes:
    print(
        "name=", node.name,
        "op=", node.op,
        "target=", node.target,
        "args=", node.args,
        "kwargs=", node.kwargs,
        "inputs=", [value.name for value in node.all_input_nodes],
        "users=", [user.name for user in node.users],
        "meta=", node.meta,
    )

6. 六种 `node.op` 及其 `target`

node.op 表示“如何执行”，node.target 进一步说明“具体执行什么”。

`op`	`target` 的含义	典型来源
`placeholder`	`forward()` 参数名，例如 `"x"`	`forward(self, x)`
`get_attr`	module attribute path，例如 `"weight"`	parameter、buffer、constant
`call_module`	submodule path，例如 `"block.conv"`	`self.block.conv(x)`
`call_function`	Python/PyTorch callable，例如 `operator.add`	`a + b`、`torch.relu(x)`
`call_method`	tensor method 名，例如 `"view"`	`x.view(...)`
`output`	通常是字符串 `"output"`	`return value`

`call_module`

对于：

%conv1 = call_module[target=block.conv1](args=(%x,))

target="block.conv1" 是相对于 GraphModule 根对象的子模块地址：

module = traced.get_submodule("block.conv1")
# 等价于 traced.block.conv1

同一个共享 module 可以被调用多次：

name="shared",   target="shared", args=(x,)
name="shared_1", target="shared", args=(y,)

两个 node 的 name 不同，因为它们是两次计算；target 相同，因为它们使用同一个 module 和同一组参数。

同一数学操作可能有不同表示

PyTorch 写法	`op`	常见 `target`
`self.relu(x)`	`call_module`	`"relu"`
`torch.relu(x)`	`call_function`	`torch.relu`
`F.relu(x)`	`call_function`	`torch.nn.functional.relu`
`x.relu()`	`call_method`	`"relu"`
`x + y`	`call_function`	`operator.add`
`torch.add(x, y)`	`call_function`	`torch.add`
`x.add(y)`	`call_method`	`"add"`

对 call_module，通过 traced.get_submodule(node.target) 取得真实 module；对 call_function，target 本身就是 callable。

7. `args`、`users` 与 residual connection

对于：

y = self.linear(x) + x

FX 记录：

add_node.args == (linear_node, x_node)

意思是 add 的两个输入分别由 linear_node 和 x_node 产生。这里保存的是 Node 引用，不是字符串、Proxy 或真实 Tensor。

反向查看：

set(x_node.users) == {linear_node, add_node}

因此：

node.args  = 我的输入来自谁
node.users = 我的输出被谁使用

all_input_nodes 是对 args/kwargs 的方便展开。例如：

torch.cat([a, b], dim=1)

即使 a_node 和 b_node 嵌套在 list 中，all_input_nodes 仍能返回：

[a_node, b_node]

8. 找到 graph inputs 和 outputs

FX 没有直接提供 traced.inputs 和 traced.outputs。输入是所有 placeholder nodes；返回结构保存在 output node 的 args 中。

input_nodes = [
    node for node in traced.graph.nodes
    if node.op == "placeholder"
]

output_node = next(
    node for node in traced.graph.nodes
    if node.op == "output"
)

print("inputs:", [(node.name, node.target) for node in input_nodes])
print("output structure:", output_node.args[0])

output_node.args[0] 可能是单个 Node，也可能是由 tuple、list 或 dict 组成的多输出结构。

9. `meta` 与 shape propagation

symbolic_trace() 主要捕获操作和连接，通常不会提供完整的 shape 和 dtype。ShapeProp 使用 representative input 执行 graph，并把 tensor_meta 写入每个 node 的 meta：

from torch.fx.passes.shape_prop import ShapeProp

ShapeProp(traced).propagate(x)

for node in traced.graph.nodes:
    tensor_meta = node.meta.get("tensor_meta")
    if tensor_meta is not None:
        print(
            node.name,
            "shape=", tuple(tensor_meta.shape),
            "dtype=", tensor_meta.dtype,
        )

需要区分：

node.type                  = Python type annotation，例如 torch.Tensor
node.meta["tensor_meta"]  = 实际传播得到的 tensor metadata
tensor_meta.dtype          = 实际 dtype，例如 torch.float32

10. Symbolic tracing 的限制

FX 无法直接处理依赖 tensor 数值的 Python control flow：

def forward(self, x):
    if x.sum() > 0:
        return self.a(x)
    return self.b(x)

Tracing 时 x.sum() > 0 仍然是 Proxy，没有真实值，Python if 无法决定分支，因此通常抛出 TraceError。

普通 Python 常量操作会在 tracing 时直接执行。例如：

constant = 2 + 3
y = x * constant

graph 通常只记录 x * 5，不会为 2 + 3 创建 node。

11. 快速检查函数

def summarize_fx(graph_module: fx.GraphModule) -> None:
    for node in graph_module.graph.nodes:
        tensor_meta = node.meta.get("tensor_meta")
        raw_shape = getattr(tensor_meta, "shape", None)
        shape = tuple(raw_shape) if raw_shape is not None else None

        print(
            f"{node.name:12} "
            f"op={node.op:13} "
            f"target={str(node.target):24} "
            f"inputs={[value.name for value in node.all_input_nodes]} "
            f"users={[user.name for user in node.users]} "
            f"shape={shape}"
        )


summarize_fx(traced)

最值得优先查看的属性是：

op      → 这是什么类别的操作？
target  → 具体调用什么？
args    → 输入和参数是什么？
users   → 输出流向哪里？
meta    → tensor 的 shape、dtype 等信息是什么？

1. 核心概念

2. Residual network 示例

3. GraphModule 的常用属性

4. 如何阅读一行 FX graph

5. Node 的关键属性

6. 六种 node.op 及其 target

call_module