Exporting Models to ONNX
Overview
Exporting a model to ONNX makes it "portable". You can take a PyTorch model and run it in the browser (ONNX.js), on mobile, or in a highly optimized inference server.
PyTorch Example
``python
import torch
import torchvision
# 1. Load Model
model = torchvision.models.resnet18(pretrained=True)
model.eval()
# 2. Define Dummy Input (Shape is critical)
# (Batch Size, Channels, Height, Width)
dummy_input = torch.randn(1, 3, 224, 224)
# 3. Export
torch.onnx.export(
model,
dummy_input,
"resnet18.onnx",
input_names=['input_image'],
output_names=['class_probs'],
dynamic_axes={'input_image': {0: 'batch_size'}} # Allow variable batch size
)
`
Validation
Always verify the export worked.
`python`
import onnx
model = onnx.load("resnet18.onnx")
onnx.checker.check_model(model)
Common Pitfalls
- Dynamic Logic: Loops (for i in range(x)) or if` statements inside the model can fail if they depend on the data values. Scripting/Tracing methods handle these differently.
- Custom Layers: If your model uses a weird custom layer that isn't in the ONNX standard opset, export will fail.