ONNX and model export: exchanging models between frameworks
In short: ONNX is an open exchange format for trained models; it lets networks trained in PyTorch or TensorFlow run quickly and across platforms independent of the framework.
In more detail: A model is often trained in Python/PyTorch but should run in C++, Java, the browser, on a phone or edge device. ONNX (Open Neural Network Exchange) describes the model as a graph of operators with weights in one file.
In Depth
Export and running
import torch
model.eval()
example = torch.randn(1, 3, 224, 224)
torch.onnx.export(model, example, 'model.onnx'
input_names=['image'], output_names=['classes']
dynamic_axes={'image': {0: 'batch'}}, opset_version=17)import onnxruntime as ort, numpy as np
sess = ort.InferenceSession('model.onnx', providers=['CPUExecutionProvider'])
out = sess.run(None, {'image': np.random.randn(1,3,224,224).astype('float32')})ONNX Runtime runs on CPU, GPU (CUDA, DirectML), NPU, mobile and in the browser (onnxruntime-web, WebGPU/WASM). Check results against the original (np.allclose).
Optimising
- Graph optimisation: fuse operators, fold constants.
- Quantisation (INT8/FP16) for smaller, faster models (see quantisation).
- Pruning/distillation before export.
- Batching, fixed vs. dynamic input shapes (
dynamic_axes). - Hardware accelerators: TensorRT (NVIDIA), OpenVINO (Intel), Core ML (Apple), TFLite (Android/edge), ncnn.
- LLMs: GGUF/llama.cpp, safetensors; ONNX GenAI.
- Profile (latency, memory) on target hardware.
Formats overview
- ONNX: neutral, broad runtime support.
- TorchScript/torch.export: PyTorch-internal.
- SavedModel/TFLite: TensorFlow/edge.
- Core ML (.mlmodel): Apple.
- safetensors: safe weight files (no code execution like Pickle; pickle files from unknown sources are unsafe!).
- GGUF: quantised LLMs for llama.cpp/Ollama.
- Hugging Face Hub: share models, read model cards and licences.
Pitfalls
- Not every operator is supported; dynamic control flow (if/loops) is hard to export.
- Pre-/post-processing (tokenizer, normalisation, scaling) must be rebuilt identically: frequent source of errors!
- Version dependence (opset); check accuracy differences FP16/INT8.
- Licence of the model and training data (commercial use?).
- Security: treat models from unknown sources like software (signatures, hash, safetensors).
- Maintenance: model versioning, monitoring (MLOps), rollback.
See also: MLOps, Quantisation, AI frameworks