PyTorch FX Symbolic Tracing

Symbolic tracing records a module's forward pass as a graph whose nodes are single computation steps. torch.fx.symbolic_trace does the recording and returns a GraphModule that can run and print the graph.

Tracing returns a GraphModule

The trace wraps the same layers in a runnable graph module. The isinstance check proves the wrap happened.

import torch


class TinyNet(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = torch.nn.Linear(4, 8)
        self.fc2 = torch.nn.Linear(8, 2)

    def forward(self, x):
        x = self.fc1(x)
        x = torch.relu(x)
        return self.fc2(x)


torch.manual_seed(0)
net = TinyNet()

traced = torch.fx.symbolic_trace(net)
print("Traced module is a GraphModule:", isinstance(traced, torch.fx.GraphModule))

Reading the nodes in order

Each line names a node's kind and its name. The order matches the forward pass from input to output.

for node in traced.graph.nodes:
    print(node.op, node.name)

The traced module still runs

A GraphModule is a real module and answers calls. Three rows in give three rows out.

x = torch.randn(3, 4)
with torch.no_grad():
    out = traced(x)
print(out.shape)

Example

import torch


class TinyNet(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = torch.nn.Linear(4, 8)
        self.fc2 = torch.nn.Linear(8, 2)

    def forward(self, x):
        x = self.fc1(x)
        x = torch.relu(x)
        return self.fc2(x)


torch.manual_seed(0)
net = TinyNet()
traced = torch.fx.symbolic_trace(net)

print(traced.graph)
print("Node count:", len(list(traced.graph.nodes)))

x = torch.randn(3, 4)
with torch.no_grad():
    plain = net(x)
    graphed = traced(x)
print("Outputs equal:", torch.equal(plain, graphed))

Key points

Use torch.fx.symbolic_trace(module) to trace a module
GraphModule: the runnable module that holds a traced graph
placeholder: the node that stands for the input
call_module: a node that runs one of the module's own layers
call_function: a node that runs a plain function, like torch.relu
The output node returns the final node's value
In short, tracing turns running code into a list of steps that a program can inspect
Tracing records the path the example input takes

In this example, a two-layer module is traced and its graph is printed. The placeholder node is the input x. The fc1 node runs the first Linear layer on it. The relu node applies torch.relu as a plain function call. The fc2 node runs the second Linear layer, and the return line hands its value back. Five nodes carry the whole forward pass. The traced module then runs on fresh input and matches the original exactly.

Note

A branch the input never triggers leaves no node in the graph.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

import torch


class TinyNet(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = torch.nn.Linear(4, 8)
        self.fc2 = torch.nn.Linear(8, 2)

    def forward(self, x):
        x = self.fc1(x)
        x = torch.relu(x)
        return self.fc2(x)


torch.manual_seed(0)
net = TinyNet()
traced = torch.fx.symbolic_trace(net)

print(traced.graph)
print("Node count:", len(list(traced.graph.nodes)))

x = torch.randn(3, 4)
with torch.no_grad():
    plain = net(x)
    graphed = traced(x)
print("Outputs equal:", torch.equal(plain, graphed))

Exercise

What does a call_module node do when the graph runs?

It prints the graph It runs one of the module's own layers It creates a new module It stops the trace