Exporting a model
torch.export.export(model, example_inputs) runs forward once symbolically and records it as a graph of core PyTorch operators: an ExportedProgram, with no Python code left in it. That graph is what deployment tools take, such as ExecuTorch for phones or the converter to Core ML.
import torch
from torch import nn
model = nn.Sequential(nn.Linear(4, 8),
nn.ReLU())
example = (torch.randn(2, 4),)
ep = torch.export.export(model, example)
for node in ep.graph.nodes:
if node.op == "call_function":
print(node.target)
print(ep.module()(torch.randn(2, 4)).shape)Every size is fixed to the example's, and the exported module refuses any other batch size. A simple way to serve any number of inputs with it is to run them in chunks of that fixed size, padding the last chunk with zeros and dropping the padded rows from the result.
torch.export.save(ep, "model.pt2") writes the program to a file, and torch.export.load reads it back, with its weights.
Your task
Write export_model(model, example) returning torch.export.export(model, (example,)). Write ops(ep) returning the targets of the call_function nodes of ep.graph, in order, as strings. Write run_batches(ep, x, batch) that runs ep.module() on x in chunks of exactly batch rows, padding the last chunk with zero rows, and returns the outputs for the real rows only, concatenated. Write round_trip(ep, path) that saves ep to path and returns the loaded program.
The explanation opens here once every check passes.