Description

This pull request introduces ONNX support for the UME model, enabling model export, inference, and benchmarking with ONNX Runtime. Key changes include the addition of ONNX utility functions, updates to the UME model for export compatibility, and a new example script demonstrating ONNX workflows.

ONNX Support Enhancements

• ONNX Export and Forward Compatibility:

  • Added forward and export_onnx methods to the UME model in src/lobster/model/_ume.py for ONNX export and inference compatibility. These methods streamline the export process and ensure compatibility with PyTorch workflows.

• ONNX Utility Functions:

  • Created src/lobster/model/onnx_utils.py with utility functions for ONNX inference (run_onnx_inference), output comparison (compare_onnx_pytorch), and performance benchmarking (benchmark_onnx_pytorch). These functions simplify ONNX workflows and provide tools for evaluating ONNX models against PyTorch implementations.

Example and Documentation

• ONNX Example Script:
  • Added examples/onnx_example.py, a comprehensive script demonstrating ONNX export, inference, comparison, and benchmarking for UME models. This script serves as a practical guide for users to understand and utilize ONNX features.

Dependency Updates

• ONNX Runtime Dependency:
  • Updated pyproject.toml to include the onnxruntime package, which is essential for running ONNX models.

Testing Configuration

• Pytest Markers:
  • Added markers for slow, integration, and unit tests in pytest.ini, along with default options to run only fast tests. This improves test organization and execution control.

Type of Change

• Bug fix
• New feature
• Documentation update
• Performance improvement
• Code refactoring

Testing

• Tests pass locally
• Added new tests for new functionality
• Updated existing tests if needed

Checklist

• Code follows style guidelines
• Self-review completed
• Documentation updated if needed
• No breaking changes (or clearly documented)