ONNX export in PyTorch - Model Metrics & Evaluation

Metrics & Evaluation - ONNX export

Which metric matters for ONNX export and WHY

When exporting a PyTorch model to ONNX format, the key metric is model output consistency. This means the ONNX model should produce the same predictions as the original PyTorch model for the same inputs. This ensures the exported model works correctly in other environments.

Metrics like accuracy, precision, or recall are not directly relevant to the export process itself but are important to verify before export. The main focus during export is that the model's outputs match closely between PyTorch and ONNX.

Confusion matrix or equivalent visualization

For ONNX export, a confusion matrix is not applicable. Instead, we compare outputs from PyTorch and ONNX models using numerical checks.

PyTorch output: [0.8, 0.1, 0.1]
ONNX output:    [0.799, 0.101, 0.100]
Difference:    [0.001, 0.001, 0.000]
Max difference < tolerance (e.g., 1e-3) means export is good.
    

Tradeoff: Export fidelity vs model complexity

Exporting complex models to ONNX can sometimes cause small differences in outputs due to unsupported operations or precision changes.

There is a tradeoff between:

• High fidelity: Outputs match very closely, but export may require simplifying or changing some model parts.
• Full complexity: Export the exact model, but outputs may differ slightly or export may fail.

Choosing the right balance depends on your use case. For critical applications, prioritize fidelity and test outputs carefully.

What "good" vs "bad" export looks like

Good export:

• ONNX model runs without errors.
• Outputs match PyTorch outputs within a small tolerance (e.g., max difference < 1e-3).
• Model size and structure are reasonable.

Bad export:

• Export fails or crashes.
• ONNX outputs differ greatly from PyTorch outputs (large numerical differences).
• Unsupported operations cause missing or incorrect behavior.

Common pitfalls when exporting to ONNX

• Dynamic axes not set: Causes fixed input sizes, limiting model flexibility.
• Unsupported ops: Some PyTorch operations may not be supported in ONNX, causing export errors or wrong outputs.
• Model in training mode: Exporting while model is in training mode can cause different outputs (e.g., dropout active).
• Not testing outputs: Failing to compare PyTorch and ONNX outputs can hide export issues.
• Ignoring input preprocessing: Differences in input data format or normalization can cause output mismatches.

Self-check question

Your PyTorch model and ONNX model outputs differ by a max of 0.05 on some inputs. Is this export good?

Answer: No, a max difference of 0.05 is quite large and may cause wrong predictions or behavior. You should investigate the cause, check for unsupported ops, and try to reduce the difference below a small tolerance like 1e-3.

Key Result

ONNX export quality is measured by how closely the ONNX model outputs match the original PyTorch model outputs.