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

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.

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.