Model Export

ONNX for Model Interoperability

What is ONNX?
ONNX (Open Neural Network Exchange) is an open format for representing machine learning models, enabling interoperability between frameworks.

Why ONNX?
| Benefit | Description |
|---------|-------------|
| Portability | Train in PyTorch, deploy anywhere |
| Optimization | Use ONNX Runtime for fast inference |
| Hardware support | Deploy to various accelerators |
| Tool ecosystem | Quantization, profiling, editing |

Exporting PyTorch to ONNX

Basic Export

import torch

model = YourModel()
model.eval()

# Create dummy input matching expected shape
dummy_input = torch.randn(1, 512)  # (batch, seq_len)

torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={
        "input": {0: "batch", 1: "seq_len"},
        "output": {0: "batch"},
    },
    opset_version=14,
)

For Transformers

from transformers import AutoModelForCausalLM
from optimum.exporters.onnx import main_export

# Export with optimum
main_export(
    model_name_or_path="meta-llama/Llama-2-7b-hf",
    output="./llama-onnx",
    task="text-generation",
)

ONNX Runtime Inference

Basic Usage

import onnxruntime as ort
import numpy as np

# Create session
session = ort.InferenceSession("model.onnx")

# Run inference
inputs = {"input": np.array([[1, 2, 3, 4, 5]], dtype=np.int64)}
outputs = session.run(None, inputs)

Optimizations

# Optimize for target hardware
sess_options = ort.SessionOptions()
sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL

# Use specific providers
providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
session = ort.InferenceSession("model.onnx", sess_options, providers=providers)

ONNX Ecosystem
| Tool | Purpose |
|------|---------|
| ONNX Runtime | Fast inference engine |
| onnx-simplifier | Simplify ONNX graphs |
| onnxoptimizer | Graph optimizations |
| Netron | Visualize ONNX models |

Limitations for LLMs
- Dynamic KV cache handling is complex
- Large models may have export issues
- Some custom ops need converter extensions

When to Use ONNX
| Scenario | Recommendation |
|----------|----------------|
| Cross-framework deployment | Yes |
| Edge/mobile deployment | Yes |
| NVIDIA GPU serving | Consider TensorRT directly |
| CPU inference | ONNX Runtime is excellent |