Lab 01 — ONNX Export & Optimization
Phase: 05 — ML Compilers & IR | Difficulty: ⭐⭐⭐⭐☆ | Time: 3–4 hours
ONNX is the universal interchange format between training frameworks and inference runtimes (TensorRT, ONNX Runtime, CoreML, QNN).
What you build
- Export PyTorch model to ONNX opset 17
- Apply graph optimization passes (constant folding, op fusion)
- Validate numerical equivalence PyTorch ↔ ONNX Runtime
- Measure latency: PyTorch eager vs ORT vs optimized ORT
Key concepts
| Concept | What to understand |
|---|---|
| ONNX opset | Versioned set of ops — opset 17 is required for newer transformer ops |
torch.onnx.export | Traces model; use dynamo_export for torch.compile models |
| ORT optimizations | OrtOptimizationLevel.ORT_ENABLE_ALL applies graph-level rewrites |
| Dynamic shapes | dynamic_axes allows variable batch/sequence length |
Files
| File | Purpose |
|---|---|
lab.py | Stubs with TODO markers |
solution.py | Complete working implementation |
test_lab.py | pytest suite |
requirements.txt | Dependencies |
Run
pip install -r requirements.txt
python lab.py
pytest test_lab.py -v