PyTorch torch.export

torch.export.export captures a model as a strict graph for deployment tools. A graph here is a list of operations plus the values flowing between them.

Exporting a small module

A two-layer Sequential exports to a short graph. The node count comes straight from the exported object.

import torch

torch.manual_seed(0)
model = torch.nn.Sequential(torch.nn.Linear(4, 3), torch.nn.ReLU())
exported = torch.export.export(model, (torch.randn(2, 4),))
print(len(list(exported.graph.nodes)))

Reading the node kinds

Each node names its kind in its op field. Inputs come first, and the output node closes the graph.

import torch

torch.manual_seed(0)
model = torch.nn.Sequential(torch.nn.Linear(4, 3), torch.nn.ReLU())
exported = torch.export.export(model, (torch.randn(2, 4),))
print([n.op for n in exported.graph.nodes])

Running the exported module

The module built from the graph answers like the original. The equality check compares every value.

import torch

torch.manual_seed(0)
model = torch.nn.Sequential(torch.nn.Linear(4, 3), torch.nn.ReLU())
x = torch.randn(2, 4)
exported = torch.export.export(model, (x,))
with torch.no_grad():
    same = torch.equal(model(x), exported.module()(x))
print(same)

Example

import torch

class Small(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = torch.nn.Linear(4, 3)

    def forward(self, x):
        return torch.relu(self.fc(x))

torch.manual_seed(0)
model = Small()
example = torch.randn(2, 4)
exported = torch.export.export(model, (example,))
nodes = list(exported.graph.nodes)
print(f"Graph nodes: {len(nodes)}")
print(f"Node kinds: {[n.op for n in nodes]}")
with torch.no_grad():
    expected = model(example)
    got = exported.module()(example)
print(f"Outputs match: {torch.equal(expected, got)}")
print(f"First output row: {[round(v, 4) for v in got[0].tolist()]}")

Key points

Node: one entry in the graph, like one input or one operation
ExportedProgram: the object that export returns, holding graph and weights
Use torch.export.export(model, (example,)) to capture a model
Use exported.module() to run the exported graph as a model
The exported graph is the model, written as operations other tools can read

In this example, a small module wraps one Linear layer and a ReLU. Export captures it with one example batch of two rows. The graph holds six nodes, and their kinds print in order. Three placeholders carry the input and the two parameters. Two call_function nodes do the maths, and one output node ends the graph. The exported module matches the original output exactly.

Note

torch.export changed across 2.x releases, and this lesson targets torch 2.14.1+cpu

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

import torch

class Small(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = torch.nn.Linear(4, 3)

    def forward(self, x):
        return torch.relu(self.fc(x))

torch.manual_seed(0)
model = Small()
example = torch.randn(2, 4)
exported = torch.export.export(model, (example,))
nodes = list(exported.graph.nodes)
print(f"Graph nodes: {len(nodes)}")
print(f"Node kinds: {[n.op for n in nodes]}")
with torch.no_grad():
    expected = model(example)
    got = exported.module()(example)
print(f"Outputs match: {torch.equal(expected, got)}")
print(f"First output row: {[round(v, 4) for v in got[0].tolist()]}")

Exercise

What does exported.module() give back?

A runnable model built from the graph The training loop A text file of weights The optimizer