Want to prepare a PyTorch model for deployment without leaving the PyTorch ecosystem? `torch.export` takes a `torch.nn.Module` and an example input and produces a traced graph: a map of the operations and data flow that turn inputs into outputs. The result is an `ExportedProgram` that can be inspected or used for ahead-of-time compilation. section 17.2.2
from torch.export import export dummy_input = torch.randn(2, 10) exported_program = export(model, args=(dummy_input,))
Pass the model and an example input as an `args` tuple to create an `ExportedProgram`. section 17.2.2
Printing the exported program makes the graph less abstract. In the example, its `forward` graph applies a linear operation, ReLU, and another linear operation, with tensor shapes shown along the way. To run the exported model on an input, call `exported_program.module()(dummy_input)`. section 17.2.2
Listing 17.9 Examining exported program graph representation
# In[]
print(exported_program)
# Out[]
ExportedProgram:
class GraphModule(torch.nn.Module):
def forward(self, p_fc1_weight: "f32[5, 10]", p_fc1_bias: "f32[5]",
p_fc2_weight: "f32[2, 5]", p_fc2_bias: "f32[2]", x: "f32[2, 10]"):
# File: .py:10 in forward, code: x = self.fc1(x)
linear: "f32[2, 5]" = torch.ops.aten.linear.default(x, p_fc1_weight, p_fc1_bias);
x = p_fc1_weight = p_fc1_bias = None
# File: .py:11 in forward, code: x = self.relu(x)
relu: "f32[2, 5]" = torch.ops.aten.relu.default(linear); linear = None
# File: .py:12 in forward, code: x = self.fc2(x)
linear_1: "f32[2, 2]" = torch.ops.aten.linear.default(relu, p_fc2_weight, p_fc2_bias);
relu = p_fc2_weight = p_fc2_bias = None
return (linear_1,)
The printed graph shows how a small model’s layers become explicit operations and data flow. section 17.2.2
For a packaged ahead-of-time result, the example passes the exported program to `torch._inductor.aoti_compile_and_package`, writes `model.pt2`, then loads it with `torch._inductor.aoti_load_package` for execution. section 17.2.2
Full graph capture has a limit: an `if` statement whose condition depends on an input tensor can make `torch.export.export` fail rather than record only the path taken by the example input. `torch.compile` instead permits graph breaks when it encounters behavior it cannot compile, trading some potential performance for flexibility. section 17.3.1
# In[]
example_input = torch.randn(5, 10)
torch.export.export(model, (example_input,))
# Out[]
from user code:
File "/var/...py", line 13, in forward
if x.sum() > 0: #1
UserError: Dynamic control flow is not supported at the moment.
An input-dependent conditional produces a dynamic-control-flow error during export. section 17.3.1
Choose ONNX when cross-framework compatibility is the priority. For PyTorch-centered deployment, `torch.export` is an alternative; PyTorch’s ONNX export can also use this export technology by setting `dynamo=True`. section 17.2.2
