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?

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

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