EMZETT.
Login

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