PythiloLesson 30 of 30 · progress kept in this browser

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.

This lesson uses torch, which Python in this browser does not include. Read it here, and run and check it in the Pythilo app on iPhone or iPad.

All lessons