Day 42: Exporting the model: ONNX vs TorchScript
From a training checkpoint to a servable artifact
A state_dict needs your Python model class to load. For serving — especially outside Python, or optimized for speed — you export the model into a self-contained, framework-portable format. The two options: TorchScript (PyTorch's own serialized format, runs without the original class) and ONNX (an open standard many runtimes accept, and the on-ramp to Stage 6A's TensorRT optimization).
import torch
model.eval()
example = torch.randn(1, 3, 224, 224) # a sample input defines the shape
# TorchScript via tracing
traced = torch.jit.trace(model, example)
traced.save("garment_classifier.torchscript")
# ONNX — the format Stage 6A's TensorRT/ONNX Runtime will consume
torch.onnx.export(
model, example, "garment_classifier.onnx",
input_names=["image"], output_names=["logits"],
dynamic_axes={"image": {0: "batch"}}, # allow variable batch size
opset_version=17,
)Why ONNX is the strategic choice here
The roadmap points every model toward Stage 6A's optimization ladder: eager → torch.compile → ONNX → TensorRT → INT8. Exporting to ONNX now means the classifier is already positioned for that ladder later. It also decouples serving from PyTorch — ONNX Runtime can serve it on CPU efficiently, which is exactly your droplet demo path.
Trace vs script
TorchScript offers two capture modes: tracing runs an example input and records the operations (simple, but misses data-dependent control flow like if on tensor values), while scripting analyzes the code itself (handles control flow, stricter). For a straightforward CNN, tracing is fine; know that models with dynamic control flow may need scripting.
Key terms
- TorchScript
- PyTorch's serialized model format that runs without the original Python class, via tracing or scripting.
- ONNX
- An open, framework-agnostic model format consumed by many runtimes (ONNX Runtime, TensorRT); the on-ramp to inference optimization.
- Tracing
- Capturing a model by recording the operations executed on a sample input (misses data-dependent control flow).
Why does the roadmap prefer exporting the classifier to ONNX rather than only saving a PyTorch state_dict?
