Exporting and serving models

Options for getting a model into production

Inference needs less than training: call model.eval(), wrap calls in torch.inference_mode(), and batch incoming requests when you can, because one batch of 16 is much cheaper than 16 single calls on a GPU.

The simplest deployment is to serve the Python model directly behind a web API such as FastAPI: load the weights once at startup, not per request. When you need the model outside Python or in an optimised runtime, export it. torch.export.export(model, example_args) traces the model into an ExportedProgram, a full graph without Python control flow, which is the foundation for PyTorch's ahead-of-time compilation and mobile/edge tooling. ONNX export (torch.onnx.export) produces a file that ONNX Runtime and other engines can run on servers, browsers and phones. TorchScript is an older export path and is not the recommended choice for new projects.

Quantization to int8 can shrink a model and speed up CPU inference at some accuracy cost; always measure that cost on your validation set. Whatever you choose, check the exported model's outputs against the original eager model.

Export and verify

assert_close compares outputs within floating-point tolerances. Export after eval() so dropout and batch norm are in inference behaviour. SmallCNN comes from the CNN topic.

import torch

model = SmallCNN()
model.load_state_dict(torch.load("model.pt", map_location="cpu", weights_only=True))
model.eval()

example = torch.randn(1, 3, 32, 32)

exported = torch.export.export(model, (example,))
with torch.inference_mode():
    torch.testing.assert_close(exported.module()(example), model(example))

torch.onnx.export(model, (example,), "small_cnn.onnx",
                  input_names=["image"], output_names=["logits"])

Load once, validate inputs

In a web service, load the model at startup and reuse it; loading per request is slow. Validate input size and type before converting to a tensor, cap batch and payload sizes, and never load model files uploaded by users, since torch.load of untrusted files is a code-execution risk.

त्वरित जाँच: Before trusting an exported model in production, what should you do?

• Retrain it from scratch in the export format
• Compare its outputs with the original eager model on sample inputs
• Convert all weights to int64
• Delete the original PyTorch weights